Merge pull request AI-Hypercomputer#2912 from AI-Hypercomputer:chengnuojin-zero1-test

PiperOrigin-RevId: 852931065

Author: Google-ML-Automation

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