397b moe rollout fixes

examples/deepswe/train_maxtext_nb.py

@@ -231,6 +231,20 @@ def str2bool(v):
     default=None,
     help="Optional override for rollout mesh TP dimension.",
 )
+parser.add_argument(
+    "--rollout_mesh_ep",
+    type=int,
+    default=1,
+    help=(
+        "Rollout expert parallelism. Required for large MoE models: the vLLM"
+        " adapter pads the MoE MLP dim up to a multiple of 2 * num_lanes, and"
+        " with ep=1 that padding grows at the same rate as the per-chip shard"
+        " shrinks, so per-chip weight memory is invariant under tp/dp. Expert"
+        " parallelism is the only axis that reduces experts-per-chip."
+        " Constraint: tp * ep * dcp must divide gdn_num_key_heads for models"
+        " with GatedDeltaNet layers."
+    ),
+)
 parser.add_argument(
     "--train_mesh_fsdp",
     type=int,
@@ -787,6 +801,7 @@ def mixed_type_batch_fn(elements):
 # 1. Resolve Rollout Mesh Dimensions
 rollout_fsdp = args.rollout_mesh_fsdp
 rollout_tp = args.rollout_mesh_tp
+ROLLOUT_EP = args.rollout_mesh_ep
 if rollout_fsdp is None and rollout_tp is None:
   num_rollout_devices = int(total_devices * args.rollout_split_fraction)
   rollout_tp = 2
@@ -890,13 +905,41 @@ def mixed_type_batch_fn(elements):
 # ==========================================
 # 7. Model Initialization via MaxText
 # ==========================================
+from etils import epath
 from orbax.checkpoint._src.serialization import jax_array_handlers
 from orbax.checkpoint._src.serialization import type_handler_registry
 
-# Ensure standard ArrayHandler is used for OCDBT base model restore
-type_handler_registry.register_type_handler(
-    jax.Array, jax_array_handlers.ArrayHandler(), override=True
-)
+# pathwaysutils registers CloudPathwaysArrayHandler on init, which reads
+# checkpoint shards on the Pathways workers. It does not support OCDBT yet
+# (b/365549911), so an OCDBT checkpoint has to fall back to the standard
+# ArrayHandler -- but that one reads on the client and materializes whole arrays
+# in the head container's host RAM.
+#
+# So only pay that cost when the checkpoint really is OCDBT. The client-side
+# restore scales with the largest single array rather than with model size,
+# which is why it goes unnoticed at 35B ([256, 10, 2048, 512] = 5.4 GiB, ~18 GB
+# peak) and is fatal at 397B: the scan axis makes every MoE tensor
+# [512, 15, 4096, 1024] = 64 GiB with ~12 in flight, so the head needs ~690 GB
+# and is OOMKilled. Keeping the reads on the workers peaks at 22 GB instead.
+# Note that checkpoint_storage_concurrent_gb does not bound this path.
+#
+# Outside Pathways this is a no-op either way: ArrayHandler is already the
+# default handler for jax.Array.
+if (epath.Path(MODEL_PATH) / "manifest.ocdbt").exists():
+  print(
+      "Base checkpoint is OCDBT, using the standard ArrayHandler:"
+      f" {MODEL_PATH}",
+      flush=True,
+  )
+  type_handler_registry.register_type_handler(
+      jax.Array, jax_array_handlers.ArrayHandler(), override=True
+  )
+else:
+  print(
+      "Base checkpoint is not OCDBT, keeping the registered handler so reads"
+      f" stay on the Pathways workers: {MODEL_PATH}",
+      flush=True,
+  )
 
 (
     qwen_reference,
@@ -1007,6 +1050,10 @@ def get_lora_model(base_model, model_mesh):
     "rollout_vllm_init_with_random_weights": True,
     "tensor_parallel_size": rollout_mesh.shape.get("model", 1),
     "data_parallel_size": rollout_mesh.shape.get("data", 1),
+    # RolloutConfig.expert_parallel_size defaults to 1 and was never set here,
+    # so expert parallelism was unreachable from configuration. See
+    # --rollout_mesh_ep for why large MoE models require it.
+    "expert_parallel_size": ROLLOUT_EP,
     "rollout_vllm_max_num_seqs": VLLM_MAX_NUM_SEQS,
     "rollout_vllm_max_num_batched_tokens": VLLM_MAX_BATCHED_TOKENS,
     "rollout_vllm_kwargs": {