[feat] Add SP and PP support for qwen_moe true on policy

miles/backends/megatron_utils/actor.py

@@ -344,6 +344,16 @@ def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
         else:
             return self.train_actor(rollout_id, rollout_data)
 
+    def _sync_before_rank_subset_logging(self) -> None:
+        if not getattr(self.args, "true_on_policy_mode", False):
+            return
+        if (
+            self.args.tensor_model_parallel_size <= 1
+            and self.args.pipeline_model_parallel_size <= 1
+        ):
+            return
+        dist.barrier(group=get_gloo_group())
+
     def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
         # Create data iterator for log_probs and train.
         data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
@@ -438,7 +448,9 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
             if self.rollout_data_postprocess is not None:
                 self.rollout_data_postprocess(self.args)
 
+            self._sync_before_rank_subset_logging()
             log_rollout_data(rollout_id, self.args, rollout_data)
+            self._sync_before_rank_subset_logging()
 
             # Train
             self._set_replay_stage("replay_backward")

miles/true_on_policy/config.py

@@ -33,6 +33,7 @@ class TrueOnPolicyParallelLayout:
     """Training and rollout topology relevant to true-on-policy parity."""
 
     train_tensor_parallel_size: int
+    train_sequence_parallel: bool
     train_context_parallel_size: int
     train_pipeline_parallel_size: int
     train_expert_model_parallel_size: int
@@ -44,6 +45,10 @@ class TrueOnPolicyParallelLayout:
     def uses_train_tp(self) -> bool:
         return self.train_tensor_parallel_size > 1
 
+    @property
+    def uses_train_sp(self) -> bool:
+        return self.train_sequence_parallel and self.uses_train_tp
+
     @property
     def uses_ulysses_cp(self) -> bool:
         return self.train_context_parallel_size > 1
@@ -192,6 +197,7 @@ class TrueOnPolicyConfig:
     model_profile: TrueOnPolicyModelProfile
     train_backend: TrainBackend
     tensor_model_parallel_size: int
+    sequence_parallel: bool
     context_parallel_size: int
     pipeline_model_parallel_size: int
     rollout_num_gpus_per_engine: int
@@ -205,6 +211,7 @@ class TrueOnPolicyConfig:
     def parallel_layout(self) -> TrueOnPolicyParallelLayout:
         return TrueOnPolicyParallelLayout(
             train_tensor_parallel_size=self.tensor_model_parallel_size,
+            train_sequence_parallel=self.sequence_parallel,
             train_context_parallel_size=self.context_parallel_size,
             train_pipeline_parallel_size=self.pipeline_model_parallel_size,
             train_expert_model_parallel_size=self.expert_model_parallel_size,
@@ -245,6 +252,8 @@ def validate(self) -> None:
             raise ValueError(f"{self.model_profile.family} does not support Ulysses CP true-on-policy")
         if layout.uses_train_pp and "pp" not in self.model_profile.supported_train_layouts:
             raise ValueError(f"{self.model_profile.family} does not support PP true-on-policy")
+        if layout.uses_train_sp and "sp" not in self.model_profile.supported_train_layouts:
+            raise ValueError(f"{self.model_profile.family} does not support SP true-on-policy")
         if layout.uses_train_ep and "ep" not in self.model_profile.supported_train_layouts:
             raise ValueError(f"{self.model_profile.family} does not support EP true-on-policy")
         if layout.uses_train_expert_tp and "expert_tp" not in self.model_profile.supported_train_layouts:
@@ -253,16 +262,49 @@ def validate(self) -> None:
             raise ValueError(f"{self.model_profile.family} does not support rollout EP true-on-policy")
         if self.sglang_target == "fsdp_tp" and not self.model_profile.supports_tp_invariant:
             raise ValueError(f"{self.model_profile.family} does not support TP-invariant true-on-policy")
-        if self.train_backend == "megatron" and layout.uses_train_tp and layout.uses_train_ep:
-            # TODO: Enable this once true-on-policy supports Megatron sequence parallel
-            # for MoE + tensor-parallel training.
+        self._validate_megatron_moe_rollout_topology()
+        if (
+            self.train_backend == "megatron"
+            and layout.uses_train_tp
+            and layout.uses_train_ep
+            and not layout.uses_train_sp
+        ):
             raise ValueError(
-                "Megatron MoE true-on-policy does not support train TP with EP yet. "
-                "Megatron requires sequence parallel for MoE + tensor-parallel training, "
-                "and the current true-on-policy path intentionally disables sequence parallel."
+                "Megatron MoE true-on-policy requires sequence parallel when train TP and EP "
+                "are both enabled."
             )
         self._validate_megatron_train_topology()
 
+    def _validate_megatron_moe_rollout_topology(self) -> None:
+        if self.train_backend != "megatron" or self.model_profile.family != "qwen3_moe":
+            return
+
+        if self.rollout_expert_parallel_size < 1:
+            raise ValueError("SGLang rollout EP must be at least 1 for MoE true-on-policy.")
+        if self.rollout_num_gpus_per_engine % self.rollout_expert_parallel_size != 0:
+            raise ValueError(
+                "Qwen3 MoE true-on-policy requires rollout_num_gpus_per_engine to be "
+                "divisible by sglang_expert_parallel_size "
+                f"({self.rollout_num_gpus_per_engine} % {self.rollout_expert_parallel_size} != 0)."
+            )
+
+        # TODO(true-on-policy): factor in sglang_moe_data_parallel_size when that
+        # flag is wired through the qwen3_moe true-on-policy launch path. The
+        # SGLang MoE TP is rollout_num_gpus_per_engine / (sglang_ep * sglang_moe_dp);
+        # today no script sets sglang_moe_data_parallel_size > 1, so dividing by
+        # rollout_expert_parallel_size alone is correct, but enabling MoE DP later
+        # will need this validation to multiply rollout_expert_parallel_size by
+        # sglang_moe_data_parallel_size before computing rollout_moe_tp_size.
+        rollout_moe_tp_size = self.rollout_num_gpus_per_engine // self.rollout_expert_parallel_size
+        if rollout_moe_tp_size != self.expert_tensor_parallel_size:
+            raise ValueError(
+                "Qwen3 MoE true-on-policy requires SGLang MoE TP to match Megatron "
+                "expert tensor parallelism: "
+                "rollout_num_gpus_per_engine / sglang_expert_parallel_size "
+                f"= {rollout_moe_tp_size}, but expert_tensor_parallel_size "
+                f"= {self.expert_tensor_parallel_size}."
+            )
+
     def _validate_megatron_train_topology(self) -> None:
         if self.train_backend != "megatron" or self.train_world_size is None:
             return
@@ -370,6 +412,7 @@ def build_true_on_policy_config(args: Any) -> TrueOnPolicyConfig | None:
         model_profile=profile,
         train_backend=args.train_backend,
         tensor_model_parallel_size=_get_required_int(args, "tensor_model_parallel_size"),
+        sequence_parallel=bool(getattr(args, "use_sequence_parallel", False)),
         context_parallel_size=_get_required_int(args, "context_parallel_size"),
         pipeline_model_parallel_size=_get_required_int(args, "pipeline_model_parallel_size"),
         expert_model_parallel_size=_get_optional_int(args, "expert_model_parallel_size", 1),

miles/true_on_policy/model_profiles.py

@@ -83,7 +83,7 @@ def megatron_model_type_for(self, model_name: str) -> str:
         "Qwen3-4B-Base": "qwen3-4B",
         "Qwen3-4B-Instruct-2507": "qwen3-4B-Instruct-2507",
     },
-    supported_train_layouts=("dp", "tp", "pp", "ulysses_cp"),
+    supported_train_layouts=("dp", "tp", "sp", "pp", "ulysses_cp"),
     supported_rollout_layouts=("dp", "tp"),
     contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
 )
@@ -92,7 +92,7 @@ def megatron_model_type_for(self, model_name: str) -> str:
     family="qwen3_moe",
     model_names=("Qwen3-30B-A3B",),
     megatron_model_types={"Qwen3-30B-A3B": "qwen3-30B-A3B"},
-    supported_train_layouts=("dp", "tp", "expert_tp", "ep", "pp", "ulysses_cp"),
+    supported_train_layouts=("dp", "tp", "sp", "expert_tp", "ep", "pp", "ulysses_cp"),
     supported_rollout_layouts=("dp", "tp", "ep"),
     contract=QWEN3_MOE_TRUE_ON_POLICY_V1,
 )

miles/true_on_policy/schema.py

@@ -32,7 +32,7 @@ class TrueOnPolicyContractSchema:
     logprob_contract="sglang_prefill",
     sglang_attention_backend="fa3",
     fsdp_attention_implementation="flash_attention_3",
-    disable_megatron_sequence_parallel=True,
+    disable_megatron_sequence_parallel=False,
 )
 
 QWEN3_MOE_TRUE_ON_POLICY_V1_SCHEMA = TrueOnPolicyContractSchema(
@@ -42,5 +42,5 @@ class TrueOnPolicyContractSchema:
     logprob_contract="sglang_prefill",
     sglang_attention_backend="fa3",
     fsdp_attention_implementation="flash_attention_3",
-    disable_megatron_sequence_parallel=True,
+    disable_megatron_sequence_parallel=False,
 )

miles/utils/external_utils/command_utils.py

@@ -54,13 +54,14 @@ def convert_checkpoint(
     num_nodes: int | None = None,
     extra_args: str = "",
     dir_dst: str = "/root",
+    path_dst: str | None = None,
     hf_checkpoint: str | None = None,
     megatron_path: str = "/root/Megatron-LM",
 ):
     hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}"
 
     # TODO shall we make it in host-mapped folder and thus can cache it to speedup CI
-    path_dst = f"{dir_dst}/{model_name}_torch_dist"
+    path_dst = path_dst or f"{dir_dst}/{model_name}_torch_dist"
     tracker = Path(path_dst) / "latest_checkpointed_iteration.txt"
     if tracker.exists() and tracker.read_text().strip() == "release":
         print(f"convert_checkpoint skip {path_dst} since tracker is 'release'")
@@ -78,6 +79,7 @@ def convert_checkpoint(
         fn = exec_command
     fn(
         f"source {repo_base_dir}/scripts/models/{megatron_model_type}.sh && "
+        f"export CUDA_DEVICE_MAX_CONNECTIONS=1 && "
         f"PYTHONPATH={megatron_path} "
         f"torchrun "
         f"--nproc-per-node {num_gpus_per_node} "

miles/utils/reloadable_process_group.py

@@ -70,6 +70,7 @@ def new_function(*args, **kwargs):
     dist.all_to_all = get_new_function(dist.all_to_all)
     dist.all_to_all_single = get_new_function(dist.all_to_all_single)
     dist.broadcast = get_new_function(dist.broadcast)
+    dist.broadcast_object_list = get_new_function(dist.broadcast_object_list)
     dist.reduce = get_new_function(dist.reduce)
     dist.reduce_scatter = get_new_function(dist.reduce_scatter)
     dist.reduce_scatter_tensor = get_new_function(dist.reduce_scatter_tensor)

scripts/run_qwen3_30b_a3b.py

@@ -81,7 +81,13 @@ def __post_init__(self):
             assert self.rollout_mxfp8, "train_mxfp8 requires rollout_mxfp8 to be enabled"
 
 
-def prepare(args: ScriptArgs):
+def prepare(args: ScriptArgs, *, convert_checkpoint_kwargs: dict | None = None):
+    """Download HF assets and (optionally) convert to torch-dist.
+
+    `convert_checkpoint_kwargs` lets callers (e.g. the deterministic variant)
+    override or extend the convert_checkpoint call without baking
+    true-on-policy concerns into this base script.
+    """
     U.exec_command(f"mkdir -p {args.model_dir} {args.data_dir}")
     U.exec_command(f"hf download Qwen/{args.model_name} --local-dir {args.model_dir}/{args.model_name}")
     U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir)
@@ -111,17 +117,19 @@ def prepare(args: ScriptArgs):
             dir_dst=args.model_dir,
             hf_checkpoint=f"{args.model_dir}/{args.model_name}",
             megatron_path=args.megatron_path,
+            **(convert_checkpoint_kwargs or {}),
         )
 
 
 # TODO improve layering: split algorithm vs infra
 def execute(args: ScriptArgs):
     is_debug_mode = args.mode == "debug_minimal"
-    ref_load_path = (
+    megatron_load_path = (
         f"{args.model_dir}/{args.model_name}/"
         if args.enable_megatron_bridge
         else f"{args.model_dir}/{args.model_name}_torch_dist"
     )
+    ref_load_path = megatron_load_path
     load_save_path = f"{args.output_dir}/{args.run_id}/checkpoints"
 
     if args.rollout_fp8:
@@ -135,7 +143,7 @@ def execute(args: ScriptArgs):
     ckpt_args = (
         f"--hf-checkpoint {hf_checkpoint}/ "
         f"--ref-load {ref_load_path} "
-        f"--load {load_save_path} "
+        f"--load {megatron_load_path} "
         f"--save {load_save_path} "
         f"--save-interval {2 if is_debug_mode else 20} "
         f"--save-retain-interval {2 if is_debug_mode else 20} "

scripts/run_qwen3_30b_a3b_deterministic.py

@@ -3,6 +3,13 @@
 Extends the base script with true-on-policy contract, EP parity defaults,
 and deterministic launch plan injection. Use this script when you need
 exact-zero logprob alignment between SGLang rollout and Megatron training.
+
+This module owns every true-on-policy concern (topology-suffixed torch_dist
+checkpoints, vocab padding overrides, fused-grad workaround, debug-one-sample
+plumbing). The base ``scripts/run_qwen3_30b_a3b`` script stays
+true-on-policy-agnostic; we cooperate with it through ``prepare``'s
+``convert_checkpoint_kwargs`` hook and by appending overrides at the end of
+``args.extra_args`` so argparse last-wins picks them up.
 """
 
 import os
@@ -11,7 +18,7 @@
 
 import typer
 from scripts.run_qwen3_30b_a3b import ScriptArgs as BaseScriptArgs
-from scripts.run_qwen3_30b_a3b import prepare
+from scripts.run_qwen3_30b_a3b import prepare as base_prepare
 
 import miles.utils.external_utils.command_utils as U
 from miles.true_on_policy import apply_true_on_policy_script_defaults, build_true_on_policy_launch_plan
@@ -26,6 +33,7 @@ class ScriptArgs(BaseScriptArgs):
     true_on_policy_default_rollout_ep: bool = True
 
     def __post_init__(self):
+        rollout_engine_size_was_default = self.rollout_num_gpus_per_engine is None
         super().__post_init__()
         if (
             self.sglang_expert_parallel_size == 1
@@ -36,12 +44,91 @@ def __post_init__(self):
         if (
             self.true_on_policy
             and self.sglang_expert_parallel_size > 1
-            and self.rollout_num_gpus_per_engine < self.sglang_expert_parallel_size
+            and rollout_engine_size_was_default
         ):
-            self.rollout_num_gpus_per_engine = self.sglang_expert_parallel_size
+            self.rollout_num_gpus_per_engine = (
+                self.sglang_expert_parallel_size * self.expert_tensor_parallel_size
+            )
         apply_true_on_policy_script_defaults(self)
 
 
+def _uses_topology_aware_torch_dist(args: ScriptArgs) -> bool:
+    return bool(args.true_on_policy and args.expert_model_parallel_size > 1)
+
+
+def _megatron_torch_dist_path(args: ScriptArgs) -> str:
+    if _uses_topology_aware_torch_dist(args):
+        topology = (
+            f"tp{args.tensor_model_parallel_size}"
+            f"_pp{args.pipeline_model_parallel_size}"
+            f"_ep{args.expert_model_parallel_size}"
+            f"_etp{args.expert_tensor_parallel_size}"
+        )
+        return f"{args.model_dir}/{args.model_name}_torch_dist_{topology}"
+    return f"{args.model_dir}/{args.model_name}_torch_dist"
+
+
+def _megatron_torch_dist_conversion_args(args: ScriptArgs) -> str:
+    if not _uses_topology_aware_torch_dist(args):
+        return ""
+    return (
+        f"--tensor-model-parallel-size {args.tensor_model_parallel_size} "
+        f"--pipeline-model-parallel-size {args.pipeline_model_parallel_size} "
+        f"--expert-model-parallel-size {args.expert_model_parallel_size} "
+        f"--expert-tensor-parallel-size {args.expert_tensor_parallel_size} "
+        f"{_true_on_policy_vocab_padding_args(args)}"
+    )
+
+
+def _true_on_policy_vocab_padding_args(args: ScriptArgs) -> str:
+    if not (
+        args.true_on_policy
+        and args.expert_model_parallel_size > 1
+        and args.tensor_model_parallel_size > 1
+    ):
+        return ""
+    # Qwen3-30B-A3B's real vocab is already divisible by supported TP sizes.
+    # Avoid Megatron's default extra padding so HF -> torch-dist conversion and
+    # true-on-policy scoring use the same shard width.
+    return "--make-vocab-size-divisible-by 1 "
+
+
+def _true_on_policy_sequence_parallel_backward_args(args: ScriptArgs) -> str:
+    if not (
+        args.true_on_policy
+        and args.expert_model_parallel_size > 1
+        and args.tensor_model_parallel_size > 1
+        and args.use_sequence_parallel
+    ):
+        return ""
+    # TP+EP sequence-parallel true-on-policy currently produces nonfinite local
+    # wgrad buckets with Megatron's fused gradient accumulation path. Keep the
+    # unfused path for correctness until the fused kernel path is audited.
+    return "--no-gradient-accumulation-fusion "
+
+
+def _topology_aware_extra_args(args: ScriptArgs) -> str:
+    """Args injected at the end of train_args to override the base script.
+
+    Relies on argparse last-wins semantics for ``--load`` / ``--ref-load`` so
+    the topology-suffixed torch_dist directory wins over the base's default.
+    """
+    if not _uses_topology_aware_torch_dist(args):
+        return ""
+    load_path = _megatron_torch_dist_path(args)
+    parts = [
+        f"--load {load_path}",
+        f"--ref-load {load_path}",
+    ]
+    vocab = _true_on_policy_vocab_padding_args(args).strip()
+    if vocab:
+        parts.append(vocab)
+    grad = _true_on_policy_sequence_parallel_backward_args(args).strip()
+    if grad:
+        parts.append(grad)
+    return " ".join(parts) + " "
+
+
 def _debug_one_sample_args(args: ScriptArgs) -> str:
     train_data_parallel_size = (
         args.num_nodes
@@ -60,6 +147,16 @@ def _debug_one_sample_args(args: ScriptArgs) -> str:
     )
 
 
+def prepare(args: ScriptArgs):
+    convert_kwargs: dict = {}
+    if _uses_topology_aware_torch_dist(args):
+        convert_kwargs = {
+            "path_dst": _megatron_torch_dist_path(args),
+            "extra_args": _megatron_torch_dist_conversion_args(args),
+        }
+    base_prepare(args, convert_checkpoint_kwargs=convert_kwargs)
+
+
 def execute(args: ScriptArgs):
     from scripts.run_qwen3_30b_a3b import execute as base_execute
 
@@ -70,9 +167,10 @@ def execute(args: ScriptArgs):
     debug_args = _debug_one_sample_args(args) if args.mode == "debug_one_sample" else ""
     if args.mode == "debug_one_sample":
         args.enable_eval = False
+    topology_extra = _topology_aware_extra_args(args)
     base_mode = args.mode
     args.mode = "debug_minimal" if args.mode == "debug_one_sample" else args.mode
-    args.extra_args = f"{plan.train_args} {debug_args} {args.extra_args}"
+    args.extra_args = f"{plan.train_args} {debug_args} {topology_extra} {args.extra_args}"
     try:
         base_execute(args)
     finally: