From 6fc7721ba19f992a7d02eb6a7effd76b5f032295 Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 3 Sep 2026 12:23:35 -0700 Subject: [PATCH 1/2] Only override the Pathways array handler for OCDBT base checkpoints pathwaysutils registers CloudPathwaysArrayHandler, which reads checkpoint shards on the Pathways workers. It does not support OCDBT (b/365549911), so train_maxtext_nb.py unconditionally re-registered the standard Orbax ArrayHandler before the base model restore. That handler reads on the client and materializes whole arrays in the head container's host RAM. The cost scales with the largest single array, not with model size, so it is invisible at 35B and fatal at 397B: 35B scanned [256, 10, 2048, 512] 5.4 GiB/array ~18 GB peak, fine 397B scanned [512, 15, 4096, 1024] 64 GiB/array ~690 GB, OOMKilled The scan axis (60 layers / inhomogeneous_layer_cycle_interval 4 = 15) inflates every MoE tensor 15x, and ~12 are in flight at once (4 layer-slots x {wi_0, wi_1, wo}). The head container is OOMKilled on a 733 GiB node about 12 seconds into the restore, before training starts. checkpoint_storage_concurrent_gb does not bound this path: dropping it from 96 to 8 produced an identical allocation curve. Make the override conditional on the checkpoint actually being OCDBT, detected via manifest.ocdbt. OCDBT checkpoints keep the previous behaviour; non-OCDBT ones keep CloudPathwaysArrayHandler, so the reads stay on the workers. Measured on v5p-512 with Qwen3.5-397B-A17B, everything else identical: before: ~690 GB host RAM, ~1 GB/s, OOMKilled after: 22 GB host RAM, 475 GB/s, completes The 35B path is unaffected -- its checkpoint contains manifest.ocdbt, so it still takes the override. --- examples/deepswe/train_maxtext_nb.py | 36 ++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/examples/deepswe/train_maxtext_nb.py b/examples/deepswe/train_maxtext_nb.py index 1dd5e7b79..0b2e5864a 100644 --- a/examples/deepswe/train_maxtext_nb.py +++ b/examples/deepswe/train_maxtext_nb.py @@ -890,13 +890,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, From eb039e5d98e8f0ca8185c53cbd4965d2370bbb5f Mon Sep 17 00:00:00 2001 From: Param Bole Date: Fri, 4 Sep 2026 11:11:47 -0700 Subject: [PATCH 2/2] Add --rollout_mesh_ep to enable expert parallelism on the rollout base_rollout_dict sets tensor_parallel_size and data_parallel_size from the rollout mesh but never sets expert_parallel_size, so it falls through to the RolloutConfig default of 1 (tunix/common/configs.py). vllm_sampler.py forwards the field into additional_config["sharding"]["sharding_strategy"], and generate/utils.py only infers a value when the field is -1, so expert parallelism is unreachable from configuration and no warning is emitted. This matters for large MoE models because the vLLM adapter pads the MoE MLP dimension up to a multiple of 2 * num_lanes (adapter.py:87). With ep=1 that padding grows at exactly the rate the per-chip shard shrinks, so per-chip weight memory is invariant under tp/dp: Qwen3.5-397B-A17B, bf16, per chip: tp=32, padded 1024 -> 8192, ep=1 512 experts/chip 180.0 GB OOM tp=4, padded, ep=1 512 experts/chip 180.0 GB OOM tp=32, padded, ep=8 64 experts/chip 22.5 GB fits Adding chips does not help. Expert parallelism is the only axis that reduces experts-per-chip, so it is a requirement rather than an optimisation. Before this change the model raised, during create_sharded_state and before any generation: RESOURCE_EXHAUSTED: HLO temporaries (149.96G) exceeds available HBM (95.73G) With --rollout_mesh_ep 8 on v5p, measured 77.4 / 95.0 GB per chip. Constraint worth noting for models with GatedDeltaNet layers: ATTN_HEAD is ('model', 'expert', 'dcp') when NEW_MODEL_DESIGN is set, and gdn_attention.py computes n_kq = gdn_num_key_heads // prod(ATTN_HEAD). So tp * ep * dcp must divide gdn_num_key_heads or the GDN kernel receives zero heads and fails with "Need at least one array to stack". For Qwen3.5-397B-A17B that means tp * ep * dcp <= 16. Default is 1, so existing runs are unaffected. --- examples/deepswe/train_maxtext_nb.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/examples/deepswe/train_maxtext_nb.py b/examples/deepswe/train_maxtext_nb.py index 0b2e5864a..7178f4102 100644 --- a/examples/deepswe/train_maxtext_nb.py +++ b/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 @@ -1035,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": {