in dags/pytorch_xla/r2_7.py [0:0]
def huggingface():
accelerate_v2_8 = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-accelerate-smoke-v2-8-1vm", reserved=True
),
US_CENTRAL1_C,
)
accelerate_v4_8 = task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-accelerate-smoke-v4-8-1vm"
),
US_CENTRAL2_B,
)
accelerate_v4_8 >> accelerate_v2_8
task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-hf-bert-pjrt-func-v4-8-1vm"
),
US_CENTRAL2_B,
)
# Stable Diffusion 2
task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-stable-diffusion-2-train-func-v6e-4-1vm",
network=BM_NETWORKS,
subnetwork=V5P_BM_SUBNETWORKS,
),
US_EAST5_B_CLOUD_ML_BENCHMARKING,
)
task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-stable-diffusion-2-train-func-v5p-8-1vm",
reserved=True,
network=V5_NETWORKS,
subnetwork=V5P_SUBNETWORKS,
),
US_EAST5_A_TPU_PROD_ENV_AUTOMATED,
)
task.run_queued_resource_test(
test_config.JSonnetTpuVmTest.from_pytorch(
"pt-2-7-stable-diffusion-2-train-func-v4-8-1vm"
),
US_CENTRAL2_B,
)