[OPD] Stabilize two-node GB200 HybridEP training

examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh

@@ -63,6 +63,7 @@ MODEL_ARGS=(
    --untie-embeddings-and-output-weights --vocab-size 248320 --rotary-base 10000000
    --moe-ffn-hidden-size 512 --moe-shared-expert-intermediate-size 512
    --moe-router-score-function softmax --moe-token-dispatcher-type flex
+   --moe-flex-dispatcher-backend hybridep
    --moe-router-topk 8
    --moe-layer-freq "[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]"
    --num-experts 256 --moe-grouped-gemm --moe-token-drop-policy probs --moe-router-dtype fp32
@@ -99,7 +100,7 @@ PERF_ARGS=(
    --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1
    --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1
    --recompute-granularity full --recompute-method uniform --recompute-num-layers 1
-   --use-dynamic-batch-size --max-tokens-per-gpu 16384 --log-probs-chunk-size 4096
+   --use-dynamic-batch-size --max-tokens-per-gpu 8192 --log-probs-chunk-size 4096
 )
 # Blackwell (GB200 sm100 / B300 sm103) backends, per scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py:
 # - moe-runner-backend flashinfer_cutlass: the default triton fused-MoE mis-shards
@@ -110,6 +111,9 @@ SGLANG_ARGS=(
    --sglang-watchdog-timeout 1800 --sglang-enable-metrics
    --sglang-moe-runner-backend flashinfer_cutlass --sglang-attention-backend trtllm_mha
    --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay
+   # Hybrid linear-attention CUDA graphs keep reading shared recurrent state
+   # during replay. Disable overlap so the next batch cannot overwrite it early.
+   --sglang-disable-overlap-schedule
    --sglang-mamba-scheduler-strategy extra_buffer
 )
 MISC_ARGS=(

miles/backends/megatron_utils/model.py

@@ -269,6 +269,10 @@ def forward_step(
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,
+            pad_to_ep_max=(
+                getattr(args, "moe_token_dispatcher_type", None) == "flex"
+                and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep"
+            ),
         )
         unconcat_tokens = batch["unconcat_tokens"]
         tokens = batch["tokens"]
@@ -421,6 +425,10 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
             args.data_pad_size_multiplier,
             args.qkv_format,
             allgather_cp=args.allgather_cp,
+            pad_to_ep_max=(
+                getattr(args, "moe_token_dispatcher_type", None) == "flex"
+                and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep"
+            ),
         )
 
         from miles.utils.replay_base import all_replay_managers

miles/backends/training_utils/data.py

@@ -107,6 +107,7 @@ def get_batch(
     qkv_format: str = "thd",
     get_position_ids: bool = False,
     allgather_cp: bool = False,
+    pad_to_ep_max: bool = False,
 ) -> dict[str, torch.Tensor | list[torch.Tensor] | None]:
     """
     Generate a CP-ready micro-batch with packed sequence parameters.
@@ -121,6 +122,8 @@ def get_batch(
         data_iterator: Iterator providing micro-batch data.
         keys: List of keys to fetch from the iterator.
         pad_multiplier: Multiplier for padding size calculation (default: 128).
+        pad_to_ep_max: Pad THD token rows to the largest aligned row count in
+            the expert-parallel group. Required by HybridEP for variable batches.
 
     Returns a dict including:
     - "tokens": torch.LongTensor of shape [1, T_padded] on the current CUDA device
@@ -195,8 +198,23 @@ def get_batch(
 
             tokens = torch.cat(tokens)
 
-            # Always pad to reduce memory fragmentation and maybe make the computation faster
-            pad = (pad_size - tokens.size(0) % pad_size) % pad_size
+            # HybridEP assumes every rank contributes the same number of token rows
+            # to metadata all-gather. Dynamic batches violate that assumption unless
+            # the local packed stream is padded to an EP-wide maximum.
+            local_target = (tokens.size(0) + pad_size - 1) // pad_size * pad_size
+            if pad_to_ep_max:
+                assert parallel_state.ep.group is not None
+                target = torch.tensor(
+                    [local_target],
+                    dtype=torch.int64,
+                    device=tokens.device,
+                )
+                dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group)
+                target_num_tokens = int(target.item())
+            else:
+                target_num_tokens = local_target
+
+            pad = target_num_tokens - tokens.size(0)
             if pad != 0:
                 tokens = F.pad(tokens, (0, pad), value=pad_token_id)
                 cu_seqlens.append(cu_seqlens[-1] + pad)

miles/backends/training_utils/replay_data.py

@@ -1,6 +1,7 @@
 from typing import Protocol
 
 import torch
+import torch.distributed as dist
 
 from .cp_utils import slice_with_cp
 from .parallel import get_parallel_state
@@ -115,7 +116,23 @@ def pad_func(data, pad):
                         replay_data[i] = torch.where(r != -1, r + offset, r)
                         offset += r.shape[0]
                 replay_data = torch.cat(replay_data, dim=0)
-                pad = (pad_size - replay_data.size(0) % pad_size) % pad_size
+                local_target = (replay_data.size(0) + pad_size - 1) // pad_size * pad_size
+                if (
+                    getattr(args, "moe_token_dispatcher_type", None) == "flex"
+                    and getattr(args, "moe_flex_dispatcher_backend", None) == "hybridep"
+                ):
+                    assert parallel_state.ep.group is not None
+                    target = torch.tensor(
+                        [local_target],
+                        dtype=torch.int64,
+                        device=torch.cuda.current_device(),
+                    )
+                    dist.all_reduce(target, op=dist.ReduceOp.MAX, group=parallel_state.ep.group)
+                    target_num_tokens = int(target.item())
+                else:
+                    target_num_tokens = local_target
+
+                pad = target_num_tokens - replay_data.size(0)
                 if pad != 0:
                     replay_data = pad_func(replay_data, pad)