From 1125fa60db38b424abb2a8bb3e394725f4e0ba6d Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Tue, 6 Jan 2026 18:03:50 +0000 Subject: [PATCH] add tpu zero1 ga test --- tests/integration_tests/train_tests.py | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/integration_tests/train_tests.py b/tests/integration_tests/train_tests.py index 44f1234089..41f8d8a7c9 100644 --- a/tests/integration_tests/train_tests.py +++ b/tests/integration_tests/train_tests.py @@ -405,6 +405,34 @@ def test_gpu_cudnn_flash_jax(self): def test_base_model_shardy_false(self): train_main(TrainTests.CONFIGS["base"] + ["shardy=False"]) + @pytest.mark.integration_test + @pytest.mark.tpu_only + def test_tpu_zero1_gradient_accumulation(self): + zero1_ga = [ # tests Zero-1 optimizer sharding with gradient accumulation + None, + os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"), + "base_output_directory=gs://runner-maxtext-logs", + "run_name=runner_test", + "dataset_path=gs://maxtext-dataset", + "steps=10", + "enable_checkpointing=False", + "enable_goodput_recording=False", + "dataset_type=synthetic", + "remat_policy=minimal", + "max_target_length=8192", + "per_device_batch_size=2", + "ici_data_parallelism=-1", + "dcn_data_parallelism=1", + "ici_fsdp_parallelism=1", + "dcn_fsdp_parallelism=1", + "gradient_accumulation_steps=8", + "shard_optimizer_over_data=True", + "shard_mode=explicit", + "decoder_block=llama2", + rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizer.llama2')}", + ] + train_main(zero1_ga) + @pytest.mark.integration_test @pytest.mark.gpu_only @pytest.mark.scheduled_only