[Sync][B] models/backends: MiniMax-M2.5 + rollout GPU-placement validation (slime #1929/#1992/#1934/#1944)

[aoshen02]

Scope — PR-B (models/backends mega-PR), part of RFC #107

This PR lands the non-image half of mega-PR B. The two image-gated pieces (#1947 FlashQLA, #1952 disable-param-backup) are deliberately held for follow-up commits because both modify the build (Dockerfile / docker/patch/latest/megatron.patch) and require a gb200+h200 image rebuild before their code can run — see "Pending" below.

Included (B-core, runs on the current image)

• slime #1934 + #1944 — rollout GPU-placement validation (07877b8): vime/ray/rollout_validation.py + validate_server_group_gpu_indices, wired into rollout start + registered in CI (test_rollout_validation.py, 0-GPU). Reused as-is from the already-translated branch.
• slime #1929 (minus #1992-deleted scripts) — MiniMax-M2.5 support (a262bd4): megatron layer spec (vime_plugins/models/minimax_m2.py), mbridge online-sync bridge, megatron→hf converter, and minimax-m2 / run-minimax-m2 launchers. Rollout path translated sglang→vLLM (--sglang-* → --vllm-*, --sglang-ep-size N → --vllm-enable-expert-parallel, pkill sglang → pkill -f "vllm serve"). The two convert-minimax-m2-*.sh scripts that #1929 added are omitted (deleted by follow-up #1992).

Pending (B-docker — separate commits, image-rebuild-gated)

• slime #1947 — FlashQLA GDN backend: edits docker/Dockerfile (pip install FlashQLA) → needs arm+x86 image rebuild. This is the Megatron train-side GDN kernel for Qwen3.5/Qwen3-Next (vLLM rollout already provides qwen_gdn_attention_core).
• slime #1952 — disable param-buffers CPU backup: edits docker/patch/latest/megatron.patch (DDP _ParamAndGradBuffer surgery) + actor/arguments python that depends on that patch (disable_param_buffers_cpu_backup) → also rebuild-gated; train-memory, higher-risk, will get its own careful review.

CI

• cpu/arm gate on the F+C+B stack: in progress (placement test test_rollout_validation.py registered).
• Full gb200 matrix + h200 precision: runs on the cumulative stack (stack-3) per the solo CI workflow.

🤖 Generated with Claude Code

---

Fidelity ledger (audited 2026-06-04)

op	detail
Ported	#1929 MiniMax-M2 plugins (megatron_to_hf + mbridge + models) ; #1934/#1944 ray/rollout_validation.py GPU-placement pre-flight
Translated (sglang→vllm)	GPU-placement validation adapted to vLLM engine layout
Deferred	FlashQLA/GDN (qwen_gdn_backend + reloadable memory-skip) → #1947, edits Dockerfile, rebuild-gated
Verified	minimax present + test_rollout_validation PASS on rebuilt stack (8914)

[gemini-code-assist]

Code Review

This pull request adds support for the MiniMax-M2.5 (229B MoE) model, including training run scripts, a Megatron-to-HF converter, a model bridge, and a custom SelfAttention implementation to handle full-dimension QK Norm. Feedback focuses on optimizing the custom attention mechanism to avoid expensive tensor-parallel gather/scatter operations by using an all-reduce on the sum of squares instead. Additionally, the reviewer identified a regex matching bug in the parameter converter, a typo in the Python buffering environment variable, and opportunities to simplify the code.

[gemini-code-assist]

The regular expression r"mlp.experts\.(.+)\.weight(\d+)" will fail to match standard Megatron MoE expert parameter names. In PyTorch/Megatron, expert weights are typically named either mlp.experts.{expert_idx}.linear_fc1.weight (where the index is in the middle) or mlp.experts.linear_fc1.weight.{expert_idx} (where the index is at the end with a dot separator). The current regex expects no dot before the index and places the index at the very end, which will cause the conversion script to fail with a ValueError for expert weights. We should support both standard formats robustly.

        match_middle = re.match(r"mlp.experts\.(\d+)\.(.+)\.weight", rest)
        match_end = re.match(r"mlp.experts\.(.+)\.weight\.(\d+)", rest)
        if match_middle or match_end:
            if match_middle:
                expert_idx, rest = match_middle.groups()
            else:
                rest, expert_idx = match_end.groups()

[gemini-code-assist]

Import torch and reduce_from_tensor_model_parallel_region to support the optimized local RMSNorm calculation.

Suggested change

	from megatron.core import parallel_state
	from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
	from megatron.core.tensor_parallel import (
	gather_from_tensor_model_parallel_region,
	scatter_to_tensor_model_parallel_region,
	)
	import torch
	from megatron.core import parallel_state
	from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
	from megatron.core.tensor_parallel import (
	gather_from_tensor_model_parallel_region,
	reduce_from_tensor_model_parallel_region,
	scatter_to_tensor_model_parallel_region,
	)

[gemini-code-assist]

Initialize the RMSNorm modules with the local hidden size (partitioned across TP ranks) instead of the global hidden size to enable local RMSNorm computation with a single all-reduce on the sum of squares.

Suggested change

	# Create full-dimension norms
	self.q_norm = build_module(
	q_layernorm,
	hidden_size=self.hidden_size_per_attention_head * config.num_attention_heads,
	config=self.config,
	eps=self.config.layernorm_epsilon,
	)
	self.k_norm = build_module(
	k_layernorm,
	hidden_size=self.hidden_size_per_attention_head * config.num_query_groups,
	config=self.config,
	eps=self.config.layernorm_epsilon,
	)
	# Create full-dimension norms partitioned across TP ranks
	tp_size = parallel_state.get_tensor_model_parallel_world_size()
	self.q_norm = build_module(
	q_layernorm,
	hidden_size=self.hidden_size_per_attention_head * config.num_attention_heads // tp_size,
	config=self.config,
	eps=self.config.layernorm_epsilon,
	)
	self.k_norm = build_module(
	k_layernorm,
	hidden_size=self.hidden_size_per_attention_head * config.num_query_groups // tp_size,
	config=self.config,
	eps=self.config.layernorm_epsilon,
	)

[gemini-code-assist]

The current implementation of full-dimension QK Norm performs an all_gather and a scatter of the large activation tensors (query and key) across TP ranks on every single attention layer. This introduces massive communication overhead and will severely bottleneck training. Since RMSNorm only requires the global sum of squares to compute the normalization factor, we can optimize this by computing the local sum of squares on each rank, performing a single all_reduce (sum) of shape [sq, b, 1], and then applying the normalization and partitioned weights locally. This mathematically identical optimization reduces the communication volume by orders of magnitude (e.g., from ~200MB to ~32KB per layer).

    def get_query_key_value_tensors(self, hidden_states, key_value_states=None, *args, **kwargs):
        query, key, value = super().get_query_key_value_tensors(hidden_states, key_value_states, *args, **kwargs)
        # query: [sq, b, num_heads_local, head_dim]
        # key:   [sq, b, num_kv_heads_local, head_dim]

        tp_size = parallel_state.get_tensor_model_parallel_world_size()
        if tp_size > 1:
            # Compute global RMSNorm using all-reduce on sum of squares to avoid expensive all-gather/scatter of activations
            query_sq = query.pow(2).sum(dim=[-2, -1], keepdim=True)
            key_sq = key.pow(2).sum(dim=[-2, -1], keepdim=True)

            query_sq = reduce_from_tensor_model_parallel_region(query_sq)
            key_sq = reduce_from_tensor_model_parallel_region(key_sq)

            global_dim_q = self.config.num_attention_heads * self.hidden_size_per_attention_head
            global_dim_k = self.config.num_query_groups * self.hidden_size_per_attention_head

            query_rms = torch.rsqrt(query_sq / global_dim_q + self.config.layernorm_epsilon)
            key_rms = torch.rsqrt(key_sq / global_dim_k + self.config.layernorm_epsilon)

            query = query * query_rms
            key = key * key_rms
        else:
            # Single GPU / TP=1: standard local RMSNorm
            query_sq = query.pow(2).mean(dim=[-2, -1], keepdim=True)
            key_sq = key.pow(2).mean(dim=[-2, -1], keepdim=True)
            query = query * torch.rsqrt(query_sq + self.config.layernorm_epsilon)
            key = key * torch.rsqrt(key_sq + self.config.layernorm_epsilon)

        # Apply weight scaling (weights are partitioned/local)
        query = query * self.q_norm.weight.view(-1, self.hidden_size_per_attention_head)
        key = key * self.k_norm.weight.view(-1, self.hidden_size_per_attention_head)

        return query, key, value

[gemini-code-assist]

The standard Python environment variable to disable output buffering is PYTHONUNBUFFERED, not PYTHONBUFFERED. Using PYTHONBUFFERED will have no effect, and standard output/error might remain buffered, which can delay logs or cause them to be lost in the event of a crash.

Suggested change

	export PYTHONBUFFERED=16
	export PYTHONUNBUFFERED=1

[gemini-code-assist]

This try-except block can be simplified and made more Pythonic by using getattr with a fallback value.

    head_dim = getattr(args, "kv_channels", None) or (args.hidden_size // args.num_attention_heads)