Merge pull request #3521 from AI-Hypercomputer:mattdavidow-shorten-ga-test

PiperOrigin-RevId: 893171714

Author: Google-ML-Automation

tests/integration/gradient_accumulation_test.py

@@ -66,10 +66,11 @@ def test_grad_accumulate_same_loss(self):
         "gradient_clipping_threshold=0",  # Ensures we are testing raw scales of gradients (clipping off)
         "enable_checkpointing=False",
         "enable_goodput_recording=False",
+        "decoder_block=simple",
         "base_emb_dim=256",
         "base_num_decoder_layers=4",
         rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}",
-        "steps=20",
+        "steps=2",
     ]
     # Run with gradient accumulation with accumulate_steps=10, per_device_batch=1 --> simulating per_device_batch=10
     train_main(