From 76c7435b2cec973eb127a4e57e6abdf087f9b34d Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Fri, 10 May 2024 01:19:47 +0000 Subject: [PATCH] Add fsdpv2 test to tpu ci --- test/tpu/run_tests.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh index 2653d0ed8c2..1658bca8977 100755 --- a/test/tpu/run_tests.sh +++ b/test/tpu/run_tests.sh @@ -11,6 +11,7 @@ python3 test/spmd/test_xla_distributed_checkpoint.py python3 test/spmd/test_train_spmd_linear_model.py python3 test/spmd/test_xla_spmd_python_api_interaction.py python3 test/spmd/test_xla_auto_sharding.py +python3 test/spmd/test_fsdp_v2.py XLA_EXPERIMENTAL=nonzero:masked_select:nms python3 test/ds/test_dynamic_shape_models.py -v XLA_EXPERIMENTAL=nonzero:masked_select:nms python3 test/ds/test_dynamic_shapes.py -v python3 test/test_autocast.py