in dags/pytorch_xla/nightly.py [0:0]
def llama():
llama_inference_v4_8 = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-nightly-llama2-infer-func-v4-8-1vm"
),
US_CENTRAL2_B,
)
llama_train_v4_8 = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-nightly-llama2-train-spmd-func-v4-8-1vm"
),
US_CENTRAL2_B,
)
llama_3_train_trillium = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-nightly-llama3-train-func-v6e-4-1vm",
network=V5_NETWORKS,
subnetwork=V6E_SUBNETWORKS,
),
US_CENTRAL2_B_TPU_PROD_ENV,
)
llama_3_train_v5p_2_slices = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-nightly-llama3-train-2-slice-func-v5p-8-1vm",
reserved=True,
network=V5_NETWORKS,
subnetwork=V5P_SUBNETWORKS,
num_slices=2,
),
US_EAST5_A_TPU_PROD_ENV_AUTOMATED,
)
llama_3_train_v5p_8 = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-nightly-llama3-train-func-v5p-8-1vm",
reserved=True,
network=V5_NETWORKS,
subnetwork=V5P_SUBNETWORKS,
),
US_EAST5_A_TPU_PROD_ENV_AUTOMATED,
)