Merge pull request AI-Hypercomputer#4285 from huytransformer:htn/deepseek-ragged-a2a-scheduled

Google-ML-Automation · Google-ML-Automation · commit e6ce71b06c2b · 2026-06-26T16:52:17.000-07:00
PiperOrigin-RevId: 938834190
diff --git a/tests/integration/pipeline_parallelism_test.py b/tests/integration/pipeline_parallelism_test.py
@@ -345,6 +345,7 @@ def test_circular_deepseek_megablox_same_output_and_grad(self):
         config, single_pipeline_stage_class=deepseek.DeepSeekMoELayerToLinen
     )
 
+  @pytest.mark.scheduled_only
   @pytest.mark.tpu_only
   def test_deepseek_ragged_a2a_ep_same_output_and_grad(self):
     config = pyconfig.initialize(

Original file line number	Diff line number	Diff line change
`@@ -345,6 +345,7 @@ def test_circular_deepseek_megablox_same_output_and_grad(self):`
`345`	`345`	`config, single_pipeline_stage_class=deepseek.DeepSeekMoELayerToLinen`
`346`	`346`	`)`
`347`	`347`
	`348`	`+ @pytest.mark.scheduled_only`
`348`	`349`	`@pytest.mark.tpu_only`
`349`	`350`	`def test_deepseek_ragged_a2a_ep_same_output_and_grad(self):`
`350`	`351`	`config = pyconfig.initialize(`