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

[maocheng23]

Stacked on top of #1059.

Summary

Adds the Miles/orchestration side of the Qwen3-30B-A3B MoE SP+PP true-on-policy stack.

This PR is one of three coupled PRs that should be reviewed and landed together because they extend the qwen3_moe_true_on_policy_v1 contract to allow Megatron sequence parallel and pipeline parallel training while preserving exact SGLang rollout parity.

Companion PRs:

• SGLang: [feat] Add Qwen3 MoE true-on-policy SP/PP support sgl-project/sglang#26317
• Megatron-LM: [feat] Add SP and PP kernels for qwen_moe true on policy Megatron-LM#31

Split-out cleanup:

• Remove prefill-logprob recomputation path: [cleanup] Remove prefill logprobs recomputation path #1192

Main Changes

• Add _sync_before_rank_subset_logging in the actor train loop for true-on-policy TP/PP rank-subset logging.
• Add topology-specific raw torch-dist checkpoint conversion/load paths for Qwen3 MoE TP/PP/EP/ETP runs.
• Wire SP/PP options through true-on-policy config, model profile, schema, scripts, checkpoint conversion, and launch fixtures.
• Deduplicate Qwen3 MoE rollout TP/EP topology validation into the true-on-policy config validator.
• Keep the validated SP+PP path on raw torch-dist loading; bridge-mode monkey patches and defensive/debug-only cleanup code were removed from this PR.

Validation

Cleaned 8-GPU ion7 real-workload run:

• Run id: moe_sppp_pr_tp2_pp2_ep2_onpolicy_real3_cleaned_noci_nosave_260524_ion7
• Ray job: raysubmit_a9NCtiu6K5dRVCsc
• Topology: Megatron TP=2, PP=2, EP=2, ETP=1, sequence parallel enabled; SGLang 4 rollout engines, each TP=2, EP=2.
• Result: Ray job succeeded, all 8 GPUs returned to 0 MiB.
• Step 1: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0377418305, weight version 1.0, mixed version 0.0.
• Step 2: train/train_rollout_logprob_abs_diff=0.0, train/train_rollout_kl=0.0, grad_norm=0.0389085777, weight version 2.0, mixed version 0.0.
• Timing: cleaned step times 485.2571s and 462.8948s, about 2.6-2.9% slower than the previous on-policy no-debug timing run and in the same overhead band versus off-policy.

Focused checks after cleanup:

• git diff --check
• python -m py_compile over changed Python files after splitting out prefill recomputation removal
• Miles fast checks: 53 passed before the PR split
• Megatron targeted extension/MoE checks: 3 passed
• Megatron 8-rank tensor-parallel mapping check passed on all ranks
• Full SP+PP 8-GPU E2E exact-logprob rerun on ion7

Local record:
recovery/qwen3_moe_sppp_clean/journal/2026-05-23-sppp-e2e.md

[gemini-code-assist]

Code Review

This pull request enables Megatron sequence parallel support for true-on-policy training and rollouts, specifically targeting Qwen3 models. Key updates include patching the Megatron Bridge to correctly handle global layer and expert indexing across Pipeline and Expert Parallelism, enhancing weight auditing with detailed layer summaries, and updating checkpoint conversion tools to support topology-specific paths. Feedback identifies a critical runtime error in the bridge plugin due to incorrect process group method calls, a logic error that skips expert globalization when pipeline parallelism is disabled, and an opportunity to improve synchronization logic by accounting for expert parallelism.

[gemini-code-assist]

The torch.distributed.ProcessGroup object does not have .size() or .rank() methods. Accessing them will raise an AttributeError at runtime. Use bridge_model_bridge.get_pg_size(ep_group) and bridge_model_bridge.parallel_state.get_expert_model_parallel_rank() instead.

Suggested change

	num_experts_per_rank = num_experts // ep_group.size()
	def _update_expert_number(param_name: str, param_type: str) -> str:
	local_expert_number = int(param_name.split(f".{param_type}")[-1])
	global_expert_number = num_experts_per_rank * ep_group.rank() + local_expert_number
	num_experts_per_rank = num_experts // bridge_model_bridge.get_pg_size(ep_group)
	def _update_expert_number(param_name: str, param_type: str) -> str:
	local_expert_number = int(param_name.split(f".{param_type}")[-1])
	global_expert_number = num_experts_per_rank * bridge_model_bridge.parallel_state.get_expert_model_parallel_rank() + local_expert_number

[gemini-code-assist]

The condition to skip the synchronization barrier should also account for Expert Parallelism (EP). If expert_model_parallel_size > 1, ranks are split across experts and synchronization is likely required for consistent logging in true-on-policy mode, similar to TP and PP.

Suggested change

	if (
	self.args.tensor_model_parallel_size <= 1
	and self.args.pipeline_model_parallel_size <= 1
	):
	if (
	self.args.tensor_model_parallel_size <= 1
	and self.args.pipeline_model_parallel_size <= 1
	and getattr(self.args, "expert_model_parallel_size", 1) <= 1
	):

References

1. Model parameters should be retrieved from the model configuration rather than being hardcoded.

[gemini-code-assist]

The current logic returns early and calls the original globalization function if PP size is 1. This effectively skips the _apply_ep_global_expert_number call, meaning expert indices will not be globalized in topologies where PP=1 and EP>1. The patch should ensure both PP and EP globalization are applied independently of each other.

Suggested change

	if "decoder.layers." not in param_name or bridge_model_bridge.get_pg_size(pp_group) <= 1:
	return original(models, config, param_name, vp_stage)
	def patched_megatron_local_name_to_global(models, config, param_name: str, vp_stage=None) -> str:
	param_name = original(models, config, param_name, vp_stage)
	pp_group = bridge_model_bridge.parallel_state.get_pipeline_model_parallel_group()
	if "decoder.layers." in param_name and bridge_model_bridge.get_pg_size(pp_group) > 1:
	pp_rank = bridge_model_bridge.parallel_state.get_pipeline_model_parallel_rank()
	layer_offset = get_transformer_layer_offset(config, vp_stage=vp_stage, pp_rank=pp_rank)
	param_name = _globalize_decoder_layer_name(param_name, layer_offset)
	return _apply_ep_global_expert_number(param_name, config)