Merge pull request #2574 from AI-Hypercomputer:chengnuojin-add-tiled

Google-ML-Automation · Google-ML-Automation · commit a022ce3024d8 · 2025-10-31T17:04:06.000-07:00
PiperOrigin-RevId: 826678381
diff --git a/tests/forward_pass_logit_checker.py b/tests/forward_pass_logit_checker.py
@@ -279,7 +279,7 @@ def main(config, test_args):  # pylint: disable=W0621
           rngs={"aqt": init_rng},
       )
 
-      full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits)
+      full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits, tiled=True)
       # if full_train_logits shape is [num_hosts, batch_size, seq_len, vocab_size]
       if full_train_logits.ndim == 4:
         full_train_logits = jnp.reshape(full_train_logits, (-1, config.max_target_length, config.vocab_size))

Original file line number	Diff line number	Diff line change
`@@ -279,7 +279,7 @@ def main(config, test_args): # pylint: disable=W0621`
`279`	`279`	`rngs={"aqt": init_rng},`
`280`	`280`	`)`
`281`	`281`
`282`		`- full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits)`
	`282`	`+ full_train_logits = jax.experimental.multihost_utils.process_allgather(full_train_logits, tiled=True)`
`283`	`283`	`# if full_train_logits shape is [num_hosts, batch_size, seq_len, vocab_size]`
`284`	`284`	`if full_train_logits.ndim == 4:`
`285`	`285`	`full_train_logits = jnp.reshape(full_train_logits, (-1, config.max_target_length, config.vocab_size))`