From 73f45db4f8c4707d0929f319d047682c8f306f6f Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 00:39:20 +0000
Subject: [PATCH 1/7] Support MoVA live weight synchronization

Convert gated dense and routed-value attention tensors into the canonical xLLM inference layout. Keep MoVA value experts on attention TP rather than FFN expert parallelism, synchronize router bias buffers, and reject stale cache-preserving updates.
---
 .../megatron_utils/megatron_to_hf/xllm.py     | 198 +++++++-
 .../megatron_utils/update_weight/common.py    |  41 +-
 .../hf_weight_iterator_direct.py              |   6 +-
 .../update_weight_from_distributed/mixin.py   |  10 +-
 .../update_weight_from_tensor.py              |   5 +-
 .../test_xllm_mova_weight_sync.py             | 426 ++++++++++++++++++
 6 files changed, 662 insertions(+), 24 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py

diff --git a/miles/backends/megatron_utils/megatron_to_hf/xllm.py b/miles/backends/megatron_utils/megatron_to_hf/xllm.py
index 76c340977fa..6a05ae51443 100644
--- a/miles/backends/megatron_utils/megatron_to_hf/xllm.py
+++ b/miles/backends/megatron_utils/megatron_to_hf/xllm.py
@@ -3,8 +3,133 @@
 import torch
 
 
+def _is_mova(args) -> bool:
+    return getattr(args, "mova_num_value_experts", 0) > 0
+
+
+def _attention_geometry(args) -> tuple[int, int, int, int]:
+    hidden_size = args.hidden_size
+    num_attention_heads = args.num_attention_heads
+    num_query_groups = args.num_query_groups
+    kv_channels = getattr(args, "kv_channels", None)
+    head_dim = kv_channels if kv_channels is not None else hidden_size // num_attention_heads
+
+    if num_attention_heads % num_query_groups:
+        raise ValueError(
+            f"num_attention_heads={num_attention_heads} must be divisible by "
+            f"num_query_groups={num_query_groups}"
+        )
+    if head_dim <= 0 or head_dim % 2:
+        raise ValueError(f"xLLM MoVA requires an even positive head dimension, got {head_dim}")
+    return hidden_size, num_attention_heads, num_query_groups, head_dim
+
+
+def _permute_qk_to_hf(
+    weight: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+    hidden_size: int,
+    name: str,
+) -> torch.Tensor:
+    """Convert MCore's adjacent-complex-pair Q/K rows to xLLM HF layout."""
+
+    expected_shape = (num_heads * head_dim, hidden_size)
+    if tuple(weight.shape) != expected_shape:
+        raise ValueError(f"Invalid {name} shape: got {tuple(weight.shape)}, expected {expected_shape}")
+    return (
+        weight.reshape(num_heads, head_dim // 2, 2, hidden_size)
+        .transpose(1, 2)
+        .reshape(expected_shape)
+        .contiguous()
+    )
+
+
+def _unpack_grouped_attention_projection(
+    args,
+    name: str,
+    param: torch.Tensor,
+    *,
+    include_value: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Unpack MCore's per-GQA-group Q/G/K[/V] projection."""
+
+    hidden_size, num_attention_heads, num_query_groups, head_dim = _attention_geometry(args)
+    query_heads_per_group = num_attention_heads // num_query_groups
+    segment_heads = [query_heads_per_group, query_heads_per_group, 1]
+    if include_value:
+        segment_heads.append(1)
+
+    expected_rows = num_query_groups * sum(segment_heads) * head_dim
+    if tuple(param.shape) != (expected_rows, hidden_size):
+        raise ValueError(
+            f"Invalid {name} shape: got {tuple(param.shape)}, expected "
+            f"{(expected_rows, hidden_size)}"
+        )
+
+    packed = param.reshape(num_query_groups, sum(segment_heads), head_dim, hidden_size)
+    chunks = torch.split(packed, segment_heads, dim=1)
+    query = chunks[0].reshape(num_attention_heads * head_dim, hidden_size)
+    gate = chunks[1].reshape(num_attention_heads * head_dim, hidden_size)
+    key = chunks[2].reshape(num_query_groups * head_dim, hidden_size)
+    value = chunks[3].reshape(num_query_groups * head_dim, hidden_size) if include_value else None
+
+    query = _permute_qk_to_hf(
+        query,
+        num_heads=num_attention_heads,
+        head_dim=head_dim,
+        hidden_size=hidden_size,
+        name=f"{name}.query",
+    )
+    key = _permute_qk_to_hf(
+        key,
+        num_heads=num_query_groups,
+        head_dim=head_dim,
+        hidden_size=hidden_size,
+        name=f"{name}.key",
+    )
+    return query, gate, key, value
+
+
+def _convert_grouped_value_experts(
+    args, layer_idx: str, name: str, param: torch.Tensor
+) -> list[tuple[str, torch.Tensor]]:
+    """Convert gathered MCore Wv [expert, hidden, value] to HF [value, hidden]."""
+
+    if param.ndim != 3:
+        raise ValueError(
+            f"Invalid {name} shape: got {tuple(param.shape)}, expected "
+            "[num_value_experts, hidden_size, value_width]"
+        )
+    expected_experts = getattr(args, "mova_num_value_experts", 0)
+    if expected_experts and param.shape[0] != expected_experts:
+        raise ValueError(
+            f"Invalid {name} expert count: got {param.shape[0]}, expected {expected_experts}"
+        )
+    if param.shape[1] != args.hidden_size:
+        raise ValueError(
+            f"Invalid {name} hidden dimension: got {param.shape[1]}, expected {args.hidden_size}. "
+            "The grouped MoVA weight must be gathered across regular attention TP before conversion."
+        )
+
+    return [
+        (
+            f"model.layers.{layer_idx}.self_attn.v_experts.{expert_idx}.weight",
+            expert_weight.transpose(0, 1).contiguous(),
+        )
+        for expert_idx, expert_weight in enumerate(param.unbind(dim=0))
+    ]
+
+
 def convert_xllm_to_hf(args, name, param):
-    """Convert Megatron parameter names/tensors to HuggingFace xLLM format."""
+    """Convert Megatron parameter names/tensors to HuggingFace xLLM format.
+
+    MoVA uses MCore's interleaved Q/G/K[/V] projections and a grouped value
+    weight stored as ``[expert, hidden / TP, value_width]``. The caller first
+    gathers those tensors over regular attention TP; this function then emits
+    canonical xLLM HF names consumed by both SGLang broadcast and P2P loading.
+    """
+
     if name == "module.module.embedding.word_embeddings.weight":
         return [("model.embed_tokens.weight", param)]
     if name == "module.module.output_layer.weight":
@@ -12,11 +137,8 @@ def convert_xllm_to_hf(args, name, param):
     if name == "module.module.decoder.final_layernorm.weight":
         return [("model.norm.weight", param)]
 
-    try:
-        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
-    except AttributeError:
-        head_dim = args.hidden_size // args.num_attention_heads
-    value_num_per_group = args.num_attention_heads // args.num_query_groups
+    hidden_size, num_attention_heads, num_query_groups, head_dim = _attention_geometry(args)
+    query_heads_per_group = num_attention_heads // num_query_groups
 
     decoder_layers_pattern = r"module\.module\.decoder\.layers\.(\d+)\.(.+)"
     match = re.match(decoder_layers_pattern, name)
@@ -53,18 +175,60 @@ def convert_xllm_to_hf(args, name, param):
 
         if rest == "self_attention.linear_proj.weight":
             return [(f"model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
+
         if rest == "self_attention.linear_qkv.weight":
-            param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
-            q_param, k_param, v_param = torch.split(param, [value_num_per_group, 1, 1], dim=1)
-            q_param = q_param.reshape(-1, args.hidden_size)
-            k_param = k_param.reshape(-1, args.hidden_size)
-            v_param = v_param.reshape(-1, args.hidden_size)
+            if _is_mova(args):
+                query, gate, key, value = _unpack_grouped_attention_projection(
+                    args, name, param, include_value=True
+                )
+                assert value is not None
+                return [
+                    (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query),
+                    (f"model.layers.{layer_idx}.self_attn.attn_gate_proj.weight", gate),
+                    (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key),
+                    (f"model.layers.{layer_idx}.self_attn.v_proj.weight", value),
+                ]
+
+            # Preserve the legacy xLLM Q/K/V contract when MoVA is disabled.
+            packed = param.view(num_query_groups, -1, head_dim, hidden_size)
+            query, key, value = torch.split(packed, [query_heads_per_group, 1, 1], dim=1)
             return [
-                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
-                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
-                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query.reshape(-1, hidden_size)),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key.reshape(-1, hidden_size)),
+                (f"model.layers.{layer_idx}.self_attn.v_proj.weight", value.reshape(-1, hidden_size)),
             ]
 
+        if rest == "self_attention.linear_qkg.weight":
+            if not _is_mova(args):
+                raise ValueError(f"Found MoVA Q/K/gate projection while MoVA is disabled: {name}")
+            query, gate, key, value = _unpack_grouped_attention_projection(
+                args, name, param, include_value=False
+            )
+            assert value is None
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_proj.weight", query),
+                (f"model.layers.{layer_idx}.self_attn.attn_gate_proj.weight", gate),
+                (f"model.layers.{layer_idx}.self_attn.k_proj.weight", key),
+            ]
+
+        if rest == "self_attention.value_projection.experts.weight":
+            if not _is_mova(args):
+                raise ValueError(f"Found grouped MoVA value experts while MoVA is disabled: {name}")
+            return _convert_grouped_value_experts(args, layer_idx, name, param)
+
+        sequential_value_expert_pattern = (
+            r"self_attention\.value_projection\.experts\.experts\.(\d+)\.weight"
+        )
+        match = re.match(sequential_value_expert_pattern, rest)
+        if match:
+            expert_idx = match.group(1)
+            return [(f"model.layers.{layer_idx}.self_attn.v_experts.{expert_idx}.weight", param)]
+
+        if rest == "self_attention.value_projection.router.weight":
+            return [(f"model.layers.{layer_idx}.self_attn.v_router.weight", param)]
+        if rest == "self_attention.value_projection.router.expert_bias":
+            return [(f"model.layers.{layer_idx}.self_attn.v_router.bias", param)]
+
         if rest == "mlp.linear_fc1.weight":
             gate_weight, up_weight = param.chunk(2, dim=0)
             return [
@@ -74,7 +238,11 @@ def convert_xllm_to_hf(args, name, param):
         if rest == "mlp.linear_fc2.weight":
             return [(f"model.layers.{layer_idx}.mlp.down_proj.weight", param)]
 
-        if rest in ("self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"):
+        if rest in (
+            "self_attention.linear_qkv.layer_norm_weight",
+            "self_attention.linear_qkg.layer_norm_weight",
+            "input_layernorm.weight",
+        ):
             return [(f"model.layers.{layer_idx}.input_layernorm.weight", param)]
         if rest in ("mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight"):
             return [(f"model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py
index c61f2648766..2f2621f893b 100644
--- a/miles/backends/megatron_utils/update_weight/common.py
+++ b/miles/backends/megatron_utils/update_weight/common.py
@@ -17,6 +17,38 @@
 logger = logging.getLogger(__name__)
 
 
+def is_ffn_expert_parameter(name: str) -> bool:
+    """Return whether ``name`` belongs to the EP-sharded feed-forward experts.
+
+    MoVA value projections also contain an ``experts`` path component, but
+    they are replicated over FFN expert parallelism and sharded over regular
+    attention tensor parallelism. Keep the predicate deliberately tied to the
+    MLP path so the two expert systems cannot be mixed accidentally.
+    """
+
+    return ".mlp.experts." in name
+
+
+def is_mova_model(args: Namespace) -> bool:
+    return getattr(args, "mova_num_value_experts", 0) > 0
+
+
+def validate_weight_update_cache_mode(args: Namespace) -> None:
+    """Reject cache-preserving updates for MoVA.
+
+    MoVA caches the routed value, which depends on both router and value-expert
+    weights. Reusing a KV cache after a policy update therefore mixes two model
+    versions even when keys are unchanged.
+    """
+
+    if is_mova_model(args) and getattr(args, "pause_generation_mode", None) == "in_place":
+        raise ValueError(
+            "MoVA weight updates do not support pause_generation_mode='in_place': "
+            "the routed-value KV cache must be flushed after every policy update. "
+            "Use 'retract' (recommended) or 'abort'."
+        )
+
+
 def _gather_with_stride(
     param_partitions: list[torch.Tensor], partition_dim: int, partition_stride: int
 ) -> torch.Tensor:
@@ -50,7 +82,8 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int,
 def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> torch.Tensor:
     """
     All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data.
-    Uses expert-TP for ".experts.", else regular-TP. Handles strided partitioning via partition_stride.
+    Uses expert-TP for FFN ``.mlp.experts.`` tensors, else regular-TP. Handles
+    strided partitioning via ``partition_stride``.
     """
     if "expert_bias" in name:
         return param
@@ -59,7 +92,7 @@ def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> t
     if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
         return param.data
 
-    if ".experts." in name:
+    if is_ffn_expert_parameter(name):
         tp_size = mpu.get_expert_tensor_parallel_world_size()
         tp_group = mpu.get_expert_tensor_parallel_group()
     else:
@@ -99,7 +132,7 @@ def all_gather_params_async(
             handles.append(None)
         else:
             # Start async all_gather
-            if ".experts." in info.name:
+            if is_ffn_expert_parameter(info.name):
                 tp_size = mpu.get_expert_tensor_parallel_world_size()
                 tp_group = mpu.get_expert_tensor_parallel_group()
             else:
@@ -270,7 +303,7 @@ def collect_named_tensors_for_weight_transfer(
         convert_to_global_name,
         translate_gpu_to_cpu,
     ):
-        if is_expert == (".experts." in name):
+        if is_expert == is_ffn_expert_parameter(name):
             yield name, tensor
 
 
diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
index ecdba3c8c77..152b212e270 100644
--- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
+++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
@@ -12,7 +12,7 @@
 
 from ..megatron_to_hf import convert_to_hf
 from ..sglang import monkey_patch_torch_reductions
-from .common import all_gather_params_async, named_params_and_buffers
+from .common import all_gather_params_async, is_ffn_expert_parameter, named_params_and_buffers
 from .hf_weight_iterator_base import HfWeightIteratorBase
 
 
@@ -83,7 +83,7 @@ def _get_megatron_full_params(
     if ep_size > 1:
         handles = []
         for info, param in zip(megatron_local_param_infos, params, strict=False):
-            if ".experts." in info.name:
+            if is_ffn_expert_parameter(info.name):
                 src_rank = (
                     info.src_rank
                     if info.src_rank in dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
@@ -118,7 +118,7 @@ def _get_megatron_local_param_info_buckets(args: Namespace, model: Sequence[torc
 
     for info in param_infos:
         # Expert params use expert-TP size, others use regular-TP size
-        if ".experts." in info.name:
+        if is_ffn_expert_parameter(info.name):
             tp_size = mpu.get_expert_tensor_parallel_world_size()
         else:
             tp_size = mpu.get_tensor_model_parallel_world_size()
diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py
index 6707993822f..85b931912d8 100644
--- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py
+++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py
@@ -10,7 +10,12 @@
 from miles.utils.distributed_utils import get_gloo_group
 
 from ...megatron_to_hf import convert_to_hf
-from ..common import all_gather_param, collect_named_tensors_for_weight_transfer, post_process_weights
+from ..common import (
+    all_gather_param,
+    collect_named_tensors_for_weight_transfer,
+    post_process_weights,
+    validate_weight_update_cache_mode,
+)
 
 
 class DistBucketedWeightUpdateMixin:
@@ -138,6 +143,9 @@ def _update_expert_bucket_weights(
 
     def _pause_and_prepare_engines(self) -> None:
         """Pause rollout engines, flush cache, and run pre-process if needed."""
+        # Validate on every training rank before rank 0 performs RPCs. Raising
+        # only on rank 0 would strand the other ranks at the following barrier.
+        validate_weight_update_cache_mode(self.args)
         if dist.get_rank() == 0:
             mode = self.args.pause_generation_mode
             ray.get([engine.pause_generation.remote(mode=mode) for engine in self.rollout_engines])
diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py
index 86073d5ab3b..0cc74d94347 100644
--- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py
+++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@@ -15,7 +15,7 @@
 from miles.utils.distributed_utils import get_gloo_group
 
 from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer
-from .common import post_process_weights
+from .common import post_process_weights, validate_weight_update_cache_mode
 from .hf_weight_iterator_base import HfWeightIteratorBase
 from .update_weight_from_distributed.broadcast import (
     connect_rollout_engines_from_distributed,
@@ -172,6 +172,9 @@ def update_weights(self) -> None:
         """
         version++, flush caches, process buckets. Progress on rank 0.
         """
+        # Validate on every rank before any RPC or collective. MoVA caches
+        # routed values, so retaining KV state across a policy update is stale.
+        validate_weight_update_cache_mode(self.args)
         self.weight_version += 1
 
         rank = dist.get_rank()
diff --git a/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py
new file mode 100644
index 00000000000..438e4edd9a9
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py
@@ -0,0 +1,426 @@
+from argparse import Namespace
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from torch import nn
+
+from miles.backends.megatron_utils.megatron_to_hf.xllm import convert_xllm_to_hf
+from miles.backends.megatron_utils.update_weight.common import (
+    all_gather_param,
+    collect_named_tensors_for_weight_transfer,
+    is_ffn_expert_parameter,
+    named_params_and_buffers,
+    validate_weight_update_cache_mode,
+)
+
+
+def _args(*, num_value_experts: int = 3, pause_generation_mode: str = "retract") -> Namespace:
+    return Namespace(
+        hidden_size=8,
+        num_attention_heads=4,
+        num_query_groups=2,
+        kv_channels=2,
+        mova_num_value_experts=num_value_experts,
+        pause_generation_mode=pause_generation_mode,
+    )
+
+
+def _projection(rows: int, hidden_size: int, offset: int) -> torch.Tensor:
+    return torch.arange(
+        offset,
+        offset + rows * hidden_size,
+        dtype=torch.float32,
+    ).reshape(rows, hidden_size)
+
+
+def _pack_by_query_group(*projections: torch.Tensor, args: Namespace) -> torch.Tensor:
+    query_heads_per_group = args.num_attention_heads // args.num_query_groups
+    head_dim = args.kv_channels
+    grouped = []
+    for index, projection in enumerate(projections):
+        heads_per_group = query_heads_per_group if index < 2 else 1
+        grouped.append(
+            projection.reshape(
+                args.num_query_groups,
+                heads_per_group,
+                head_dim,
+                args.hidden_size,
+            )
+        )
+    return torch.cat(grouped, dim=1).reshape(-1, args.hidden_size)
+
+
+def _permute_qk_to_hf(weight: torch.Tensor, num_heads: int, args: Namespace) -> torch.Tensor:
+    return (
+        weight.reshape(num_heads, args.kv_channels // 2, 2, args.hidden_size)
+        .transpose(1, 2)
+        .reshape_as(weight)
+    )
+
+
+def _restore_native_qk(weight: torch.Tensor, num_heads: int, args: Namespace) -> torch.Tensor:
+    """Independent inverse used by the Megatron HF checkpoint reader."""
+
+    return (
+        weight.reshape(num_heads, 2, args.kv_channels // 2, args.hidden_size)
+        .transpose(1, 2)
+        .reshape_as(weight)
+    )
+
+
+def test_dense_mova_qkgv_conversion_preserves_gate_and_value() -> None:
+    args = _args()
+    query = _projection(8, args.hidden_size, 0)
+    gate = _projection(8, args.hidden_size, 1_000)
+    key = _projection(4, args.hidden_size, 2_000)
+    value = _projection(4, args.hidden_size, 3_000)
+    packed = _pack_by_query_group(query, gate, key, value, args=args)
+
+    converted = dict(
+        convert_xllm_to_hf(
+            args,
+            "module.module.decoder.layers.0.self_attention.linear_qkv.weight",
+            packed,
+        )
+    )
+
+    torch.testing.assert_close(
+        converted["model.layers.0.self_attn.q_proj.weight"],
+        _permute_qk_to_hf(query, args.num_attention_heads, args),
+    )
+    torch.testing.assert_close(
+        _restore_native_qk(
+            converted["model.layers.0.self_attn.q_proj.weight"],
+            args.num_attention_heads,
+            args,
+        ),
+        query,
+    )
+    torch.testing.assert_close(converted["model.layers.0.self_attn.attn_gate_proj.weight"], gate)
+    torch.testing.assert_close(
+        converted["model.layers.0.self_attn.k_proj.weight"],
+        _permute_qk_to_hf(key, args.num_query_groups, args),
+    )
+    torch.testing.assert_close(
+        _restore_native_qk(
+            converted["model.layers.0.self_attn.k_proj.weight"],
+            args.num_query_groups,
+            args,
+        ),
+        key,
+    )
+    torch.testing.assert_close(converted["model.layers.0.self_attn.v_proj.weight"], value)
+
+
+def test_sparse_mova_qkg_conversion_emits_no_dense_value() -> None:
+    args = _args()
+    query = _projection(8, args.hidden_size, 0)
+    gate = _projection(8, args.hidden_size, 1_000)
+    key = _projection(4, args.hidden_size, 2_000)
+    packed = _pack_by_query_group(query, gate, key, args=args)
+
+    converted = dict(
+        convert_xllm_to_hf(
+            args,
+            "module.module.decoder.layers.3.self_attention.linear_qkg.weight",
+            packed,
+        )
+    )
+
+    assert set(converted) == {
+        "model.layers.3.self_attn.q_proj.weight",
+        "model.layers.3.self_attn.k_proj.weight",
+        "model.layers.3.self_attn.attn_gate_proj.weight",
+    }
+    torch.testing.assert_close(
+        converted["model.layers.3.self_attn.q_proj.weight"],
+        _permute_qk_to_hf(query, args.num_attention_heads, args),
+    )
+    torch.testing.assert_close(converted["model.layers.3.self_attn.attn_gate_proj.weight"], gate)
+    torch.testing.assert_close(
+        converted["model.layers.3.self_attn.k_proj.weight"],
+        _permute_qk_to_hf(key, args.num_query_groups, args),
+    )
+
+
+def test_grouped_value_experts_convert_input_to_output_sharding_layout() -> None:
+    args = _args()
+    value_width = args.num_query_groups * args.kv_channels
+    # MCore all-gather result: [expert, full hidden input, full value output].
+    grouped_weight = torch.arange(
+        args.mova_num_value_experts * args.hidden_size * value_width,
+        dtype=torch.float32,
+    ).reshape(args.mova_num_value_experts, args.hidden_size, value_width)
+
+    converted = convert_xllm_to_hf(
+        args,
+        "module.module.decoder.layers.3.self_attention.value_projection.experts.weight",
+        grouped_weight,
+    )
+
+    assert [name for name, _ in converted] == [
+        f"model.layers.3.self_attn.v_experts.{expert}.weight"
+        for expert in range(args.mova_num_value_experts)
+    ]
+    for expert, (_, weight) in enumerate(converted):
+        assert weight.is_contiguous()
+        torch.testing.assert_close(weight, grouped_weight[expert].transpose(0, 1))
+
+
+def test_grouped_value_experts_require_regular_tp_gather_first() -> None:
+    args = _args()
+    local_input_shard = torch.empty(args.mova_num_value_experts, args.hidden_size // 2, 4)
+    with pytest.raises(ValueError, match="gathered across regular attention TP"):
+        convert_xllm_to_hf(
+            args,
+            "module.module.decoder.layers.3.self_attention.value_projection.experts.weight",
+            local_input_shard,
+        )
+
+
+@pytest.mark.parametrize(
+    ("megatron_suffix", "hf_suffix"),
+    [
+        ("router.weight", "v_router.weight"),
+        ("router.expert_bias", "v_router.bias"),
+    ],
+)
+def test_value_router_parameter_and_buffer_are_both_synchronized(
+    megatron_suffix: str, hf_suffix: str
+) -> None:
+    args = _args()
+    tensor = torch.randn(args.mova_num_value_experts, args.hidden_size)
+    if megatron_suffix.endswith("expert_bias"):
+        tensor = torch.randn(args.mova_num_value_experts)
+
+    converted = convert_xllm_to_hf(
+        args,
+        f"module.module.decoder.layers.3.self_attention.value_projection.{megatron_suffix}",
+        tensor,
+    )
+    assert len(converted) == 1
+    assert converted[0][0] == f"model.layers.3.self_attn.{hf_suffix}"
+    assert converted[0][1] is tensor
+
+
+def test_legacy_xllm_qkv_conversion_is_unchanged_when_mova_is_disabled() -> None:
+    args = _args(num_value_experts=0)
+    query = _projection(8, args.hidden_size, 0)
+    key = _projection(4, args.hidden_size, 1_000)
+    value = _projection(4, args.hidden_size, 2_000)
+    # Legacy Q/K/V has Q first, so construct its own three-way group packing.
+    query_grouped = query.reshape(2, 2, 2, 8)
+    key_grouped = key.reshape(2, 1, 2, 8)
+    value_grouped = value.reshape(2, 1, 2, 8)
+    packed = torch.cat((query_grouped, key_grouped, value_grouped), dim=1).reshape(-1, 8)
+
+    converted = dict(
+        convert_xllm_to_hf(
+            args,
+            "module.module.decoder.layers.0.self_attention.linear_qkv.weight",
+            packed,
+        )
+    )
+
+    torch.testing.assert_close(converted["model.layers.0.self_attn.q_proj.weight"], query)
+    torch.testing.assert_close(converted["model.layers.0.self_attn.k_proj.weight"], key)
+    torch.testing.assert_close(converted["model.layers.0.self_attn.v_proj.weight"], value)
+    assert "model.layers.0.self_attn.attn_gate_proj.weight" not in converted
+
+
+@pytest.mark.parametrize(
+    ("name", "expected"),
+    [
+        ("module.module.decoder.layers.3.mlp.experts.linear_fc1.weight7", True),
+        ("module.module.mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight7", True),
+        ("module.module.decoder.layers.3.self_attention.value_projection.experts.weight", False),
+        (
+            "module.module.decoder.layers.3.self_attention.value_projection.experts.experts.7.weight",
+            False,
+        ),
+        ("model.layers.3.self_attn.v_experts.7.weight", False),
+    ],
+)
+def test_only_ffn_experts_use_expert_parallelism(name: str, expected: bool) -> None:
+    assert is_ffn_expert_parameter(name) is expected
+
+
+@pytest.mark.parametrize(
+    ("name", "expected_group"),
+    [
+        (
+            "module.module.decoder.layers.3.self_attention.value_projection.experts.weight",
+            "regular-tp",
+        ),
+        ("module.module.decoder.layers.3.mlp.experts.weight", "expert-tp"),
+    ],
+)
+def test_tensor_gather_uses_the_correct_tp_group(name: str, expected_group: str) -> None:
+    param = torch.nn.Parameter(torch.arange(4, dtype=torch.float32).reshape(2, 2))
+    param.tensor_model_parallel = True
+    param.parallel_mode = None
+    param.partition_dim = 0
+    param.partition_stride = 1
+
+    def fake_all_gather(partitions, source, *, group):
+        assert group == expected_group
+        partitions[0].copy_(source)
+
+    with (
+        patch(
+            "miles.backends.megatron_utils.update_weight.common.mpu.get_tensor_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "miles.backends.megatron_utils.update_weight.common.mpu.get_tensor_model_parallel_group",
+            return_value="regular-tp",
+        ),
+        patch(
+            "miles.backends.megatron_utils.update_weight.common.mpu.get_expert_tensor_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "miles.backends.megatron_utils.update_weight.common.mpu.get_expert_tensor_parallel_group",
+            return_value="expert-tp",
+        ),
+        patch(
+            "miles.backends.megatron_utils.update_weight.common.dist.all_gather",
+            side_effect=fake_all_gather,
+        ),
+    ):
+        gathered = all_gather_param(Namespace(swiglu=False), name, param)
+
+    torch.testing.assert_close(gathered, param)
+
+
+def test_weight_transfer_partition_keeps_mova_values_with_non_experts() -> None:
+    tensors = [
+        ("module.module.decoder.layers.3.mlp.experts.linear_fc1.weight7", torch.tensor(1)),
+        (
+            "module.module.decoder.layers.3.self_attention.value_projection.experts.weight",
+            torch.tensor(2),
+        ),
+        ("module.module.decoder.layers.3.self_attention.linear_qkg.weight", torch.tensor(3)),
+    ]
+    with patch(
+        "miles.backends.megatron_utils.update_weight.common.named_params_and_buffers",
+        return_value=iter(tensors),
+    ):
+        regular = list(collect_named_tensors_for_weight_transfer(_args(), [], is_expert=False))
+    with patch(
+        "miles.backends.megatron_utils.update_weight.common.named_params_and_buffers",
+        return_value=iter(tensors),
+    ):
+        experts = list(collect_named_tensors_for_weight_transfer(_args(), [], is_expert=True))
+
+    assert [name for name, _ in regular] == [tensors[1][0], tensors[2][0]]
+    assert [name for name, _ in experts] == [tensors[0][0]]
+
+
+def test_value_and_ffn_expert_bias_buffers_are_enumerated_for_sync() -> None:
+    class ModuleWithRouterBuffers(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.register_buffer("value_projection_router_expert_bias", torch.arange(3.0))
+            self.register_buffer("mlp_router_expert_bias", torch.arange(4.0))
+            self.register_buffer("unrelated_statistics", torch.arange(5.0))
+
+    names = [
+        name
+        for name, _ in named_params_and_buffers(
+            _args(),
+            [ModuleWithRouterBuffers()],
+            convert_to_global_name=False,
+        )
+    ]
+    assert names == [
+        "vp_stages.0.value_projection_router_expert_bias",
+        "vp_stages.0.mlp_router_expert_bias",
+    ]
+
+
+def test_p2p_staging_contract_requires_all_value_expert_shards() -> None:
+    from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p import (
+        UpdateWeightP2P,
+    )
+
+    updater = object.__new__(UpdateWeightP2P)
+    num_value_experts = 3
+    packed_name = "model.layers.3.self_attn.v_experts.weight"
+    updater._shared_params_dict = {packed_name: torch.empty(num_value_experts, 4, 8)}
+    updater._shared_param_mapper = MagicMock()
+    updater._shared_param_mapper.map.return_value = SimpleNamespace(
+        sglang_name=packed_name,
+        num_shards=num_value_experts,
+        num_local_experts=None,
+    )
+    updater._staged_tensors = {}
+    updater._tensor_update_pending = {}
+
+    ready_names = []
+    ready_tensors = []
+    for expert in range(num_value_experts):
+        names, tensors = updater._get_transfer_ready_params(
+            [(f"model.layers.3.self_attn.v_experts.{expert}.weight", torch.full((4, 8), expert))]
+        )
+        ready_names.extend(names)
+        ready_tensors.extend(tensors)
+
+    assert ready_names == [packed_name]
+    assert [name for name, _ in ready_tensors] == [
+        f"model.layers.3.self_attn.v_experts.{expert}.weight"
+        for expert in range(num_value_experts)
+    ]
+    assert updater._staged_tensors == {}
+    assert updater._tensor_update_pending == {}
+
+
+def test_broadcast_path_preserves_all_converted_value_expert_metadata() -> None:
+    from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (
+        update_weights_from_distributed,
+    )
+
+    tensors = [
+        (f"model.layers.3.self_attn.v_experts.{expert}.weight", torch.full((4, 8), expert))
+        for expert in range(3)
+    ]
+    engine = MagicMock()
+    engine.update_weights_from_distributed.remote.return_value = "engine-ref"
+    handle = MagicMock()
+
+    with patch(
+        "miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast.dist.broadcast",
+        return_value=handle,
+    ) as broadcast:
+        refs = update_weights_from_distributed(
+            "mova-test",
+            MagicMock(),
+            7,
+            [engine],
+            tensors,
+        )
+
+    assert refs == ["engine-ref"]
+    kwargs = engine.update_weights_from_distributed.remote.call_args.kwargs
+    assert kwargs["names"] == [name for name, _ in tensors]
+    assert kwargs["dtypes"] == [tensor.dtype for _, tensor in tensors]
+    assert kwargs["shapes"] == [tensor.shape for _, tensor in tensors]
+    assert kwargs["weight_version"] == "7"
+    assert broadcast.call_count == len(tensors)
+    assert handle.wait.call_count == len(tensors)
+
+
+@pytest.mark.parametrize("mode", ["retract", "abort"])
+def test_mova_cache_safe_update_modes(mode: str) -> None:
+    validate_weight_update_cache_mode(_args(pause_generation_mode=mode))
+
+
+def test_mova_rejects_in_place_cache_preservation() -> None:
+    with pytest.raises(ValueError, match="routed-value KV cache must be flushed"):
+        validate_weight_update_cache_mode(_args(pause_generation_mode="in_place"))
+
+
+def test_legacy_model_may_retain_existing_in_place_behavior() -> None:
+    validate_weight_update_cache_mode(_args(num_value_experts=0, pause_generation_mode="in_place"))

From 97cbef7219206c8067f9e7205a01624b3981f164 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 01:32:15 +0000
Subject: [PATCH 2/7] Build native MoVA models in Miles

Register the MoVA and xLLM router CLI contract in Miles and select Megatron's MoVA config and heterogeneous block spec whenever value experts are enabled. Fail early on incompatible provider, checkpoint-conversion, and attention settings while preserving the ordinary Transformer provider unchanged.
---
 .../backends/megatron_utils/model_provider.py |  44 +++-
 miles/utils/arguments.py                      | 119 +++++++++++
 .../test_mova_model_provider.py               | 193 ++++++++++++++++++
 tests/fast/utils/test_arguments.py            | 148 +++++++++++++-
 4 files changed, 500 insertions(+), 4 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_mova_model_provider.py

diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py
index 617068333e4..2c3b120fe06 100644
--- a/miles/backends/megatron_utils/model_provider.py
+++ b/miles/backends/megatron_utils/model_provider.py
@@ -23,6 +23,24 @@
 logger = logging.getLogger(__name__)
 
 
+def _get_mova_model_components():
+    """Import MoVA components only when the architecture is requested.
+
+    Keeping this import lazy preserves compatibility with ordinary Megatron
+    installations that have not yet taken the native MoVA extension.
+    """
+
+    try:
+        from megatron.core.models.gpt.mova_layer_specs import get_mova_gpt_decoder_block_spec
+        from megatron.core.transformer.mova import MoVATransformerConfig
+    except ImportError as error:
+        raise RuntimeError(
+            "Native MoVA was requested, but this Megatron installation does not "
+            "provide MoVATransformerConfig/get_mova_gpt_decoder_block_spec"
+        ) from error
+    return MoVATransformerConfig, get_mova_gpt_decoder_block_spec
+
+
 # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82
 class LinearForLastLayer(torch.nn.Linear):
     def __init__(
@@ -59,6 +77,14 @@ def get_model_provider_func(
     args: argparse.Namespace,
     role: Literal["actor", "critic"] = "actor",
 ):
+    is_mova = getattr(args, "mova_num_value_experts", 0) > 0
+    if is_mova:
+        # Validate before any provider branch so MoVA can never silently build a
+        # custom/ordinary provider or select an unsupported converter.
+        from miles.utils.arguments import validate_mova_args
+
+        validate_mova_args(args)
+
     # Support custom model provider path (similar to --custom-rm-path for reward models)
     if getattr(args, "custom_model_provider_path", None):
 
@@ -147,9 +173,21 @@ def model_provider(
 
         # Experimental loading arguments from yaml
         assert config is None, "miles builds the config from args, so it expects config to be None"
-        config = core_transformer_config_from_args(args)
-
-        if args.spec is not None:
+        if is_mova:
+            mova_config_class, get_mova_block_spec = _get_mova_model_components()
+            config = core_transformer_config_from_args(args, mova_config_class)
+        else:
+            config = core_transformer_config_from_args(args)
+
+        if is_mova:
+            transformer_layer_spec = get_mova_block_spec(
+                config,
+                use_transformer_engine=use_te,
+                moe_grouped_gemm=args.moe_grouped_gemm,
+                moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
+                vp_stage=vp_stage,
+            )
+        elif args.spec is not None:
             transformer_layer_spec = import_module(args.spec)
             # Allow the spec to be a function so that user can use customized Megatron easier.
             if callable(transformer_layer_spec):
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 2d6e43e1384..92b87c3afbd 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -35,6 +35,115 @@ def reset_arg(parser, name, **kwargs):
         parser.add_argument(name, **kwargs)
 
 
+def add_mova_arguments(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    """Register the native Megatron MoVA architecture arguments.
+
+    These names and defaults intentionally mirror Megatron's ``pretrain_mova``
+    entry point.  Miles owns the registration because its train entry points do
+    not invoke that script's ``extra_args_provider``.
+
+    Unlike ``pretrain_mova``, this function must not change shared Transformer
+    defaults globally: the same Miles parser also serves every non-MoVA model.
+    Exact MoVA runs therefore pass the required zero-dropout/base-attention
+    flags explicitly and are checked by :func:`validate_mova_args`.
+    """
+
+    group = parser.add_argument_group(title="mixture-of-value attention")
+    group.add_argument("--mova-num-value-experts", type=int, default=0)
+    group.add_argument("--mova-router-topk", type=int, default=1)
+    group.add_argument(
+        "--mova-router-score-function",
+        choices=("softmax", "sigmoid"),
+        default="sigmoid",
+    )
+    group.add_argument("--mova-router-topk-scaling-factor", type=float, default=1.0)
+    group.add_argument("--mova-router-enable-expert-bias", action="store_true")
+    group.add_argument("--mova-router-bias-update-rate", type=float, default=1.0e-3)
+    group.add_argument("--mova-router-aux-loss-coeff", type=float, default=0.0)
+    group.add_argument(
+        "--mova-router-load-balancing-type",
+        choices=("none", "aux_loss"),
+        default="none",
+    )
+    group.add_argument("--mova-num-dense-layers", type=int, default=0)
+    group.add_argument("--mova-norm-num-groups", type=int, default=1)
+    group.add_argument(
+        "--mova-attention-gate-function",
+        choices=("softplus", "silu"),
+        default="softplus",
+    )
+    group.add_argument(
+        "--mova-value-backend",
+        choices=("sequential", "grouped_gemm"),
+        default="grouped_gemm",
+    )
+    group.add_argument(
+        "--mova-use-torch-rms-norm",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Use PyTorch's native RMSNorm primitive for grouped normalization.",
+    )
+    group.add_argument(
+        "--xllm-router-compatibility",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Match xLLM's BF16 router GEMM followed by FP32 top-k scoring.",
+    )
+    group.add_argument(
+        "--xllm-router-gemm-partitions",
+        type=int,
+        default=1,
+        help="Original xLLM model-parallel partitions used by router GEMMs.",
+    )
+    return parser
+
+
+def validate_mova_args(args) -> None:
+    """Validate the Miles-specific boundary for native Megatron MoVA.
+
+    Shape/routing invariants remain owned by ``MoVATransformerConfig``.  The
+    checks here prevent selecting an incompatible Miles provider, conversion
+    path, or cache/model feature before distributed model construction starts.
+    """
+
+    if getattr(args, "mova_num_value_experts", 0) <= 0:
+        return
+
+    if args.train_backend != "megatron":
+        raise ValueError("MoVA training requires --train-backend megatron")
+    if args.megatron_to_hf_mode != "raw":
+        raise ValueError("Native MoVA requires --megatron-to-hf-mode raw")
+    if getattr(args, "custom_model_provider_path", None) is not None:
+        raise ValueError("Native MoVA owns its model provider; remove --custom-model-provider-path")
+    if getattr(args, "use_legacy_models", False):
+        raise ValueError("MoVA is supported only by Megatron Core models")
+    if getattr(args, "yaml_cfg", None) is not None:
+        raise ValueError("MoVA's initial Miles integration supports CLI configuration only")
+    if getattr(args, "spec", None) is not None:
+        raise ValueError("MoVA owns its heterogeneous layer spec; --spec is not supported")
+    if getattr(args, "mtp_num_layers", None) is not None:
+        raise ValueError("MoVA does not currently support MTP layers")
+    if getattr(args, "multi_latent_attention", False):
+        raise ValueError("MoVA cannot be combined with multi-latent attention")
+    if getattr(args, "heterogeneous_layers_config_path", None) is not None:
+        raise ValueError("MoVA owns its heterogeneous layer layout")
+    if not getattr(args, "attention_output_gate", False):
+        raise ValueError("MoVA requires --attention-output-gate")
+    if not getattr(args, "rotary_interleaved", False):
+        raise ValueError("xLLM MoVA requires --rotary-interleaved")
+    if not getattr(args, "group_query_attention", False):
+        raise ValueError("MoVA requires --group-query-attention")
+    if getattr(args, "num_experts", 0) <= 0:
+        raise ValueError("MoVA sparse layers require --num-experts")
+    if getattr(args, "attention_dropout", None) != 0.0 or getattr(args, "hidden_dropout", None) != 0.0:
+        raise ValueError("xLLM MoVA requires --attention-dropout 0 and --hidden-dropout 0")
+    if (
+        getattr(args, "tensor_model_parallel_size", 1) > 1
+        and not getattr(args, "sequence_parallel", False)
+    ):
+        raise ValueError("MoVA with tensor parallelism requires --sequence-parallel")
+
+
 def get_miles_extra_args_provider(add_custom_arguments=None):
     def add_miles_arguments(parser):
         # Ray
@@ -1362,6 +1471,14 @@ def add_debug_arguments(parser):
                 default="torch",
             )
             parser.add_argument("--check-weight-update-equal", action="store_true")
+            parser.add_argument(
+                "--check-all-engine-weight-versions",
+                action="store_true",
+                help=(
+                    "After every live-weight update, query every rollout engine "
+                    "and require its version to match the updater."
+                ),
+            )
             parser.add_argument(
                 "--env-report",
                 type=str,
@@ -1671,6 +1788,7 @@ def add_sglang_tp_size():
 
         parser = add_cluster_arguments(parser)
         parser = add_train_arguments(parser)
+        parser = add_mova_arguments(parser)
         parser = add_rollout_arguments(parser)
         parser = add_fault_tolerance_arguments(parser)
         parser = add_data_arguments(parser)
@@ -1843,6 +1961,7 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 
 
 def miles_validate_args(args):
+    validate_mova_args(args)
     args.eval_datasets = _resolve_eval_datasets(args)
 
     # Normalize --tito-allowed-append-roles: lowercase + deduplicate.
diff --git a/tests/fast/backends/megatron_utils/test_mova_model_provider.py b/tests/fast/backends/megatron_utils/test_mova_model_provider.py
new file mode 100644
index 00000000000..cfecc01d6d3
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_mova_model_provider.py
@@ -0,0 +1,193 @@
+from argparse import Namespace
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from miles.backends.megatron_utils import model_provider as model_provider_module
+
+
+def _provider_args(*, mova_num_value_experts: int) -> Namespace:
+    return Namespace(
+        custom_model_provider_path=None,
+        mova_num_value_experts=mova_num_value_experts,
+        train_backend="megatron",
+        megatron_to_hf_mode="raw",
+        use_legacy_models=False,
+        yaml_cfg=None,
+        spec=None,
+        mtp_num_layers=None,
+        multi_latent_attention=False,
+        heterogeneous_layers_config_path=None,
+        attention_output_gate=True,
+        rotary_interleaved=True,
+        group_query_attention=True,
+        num_experts=100 if mova_num_value_experts else 0,
+        attention_dropout=0.0,
+        hidden_dropout=0.0,
+        tensor_model_parallel_size=8,
+        sequence_parallel=True,
+        transformer_impl="transformer_engine",
+        moe_grouped_gemm=True,
+        moe_use_legacy_grouped_gemm=False,
+        qk_layernorm=False,
+        fp8_param_gather=False,
+        padded_vocab_size=256,
+        max_position_embeddings=32768,
+        fp16_lm_cross_entropy=False,
+        untie_embeddings_and_output_weights=True,
+        position_embedding_type="rope",
+        rotary_percent=1.0,
+        rotary_base=10_000_000,
+        use_rope_scaling=False,
+        use_rollout_routing_replay=False,
+    )
+
+
+def test_native_mova_provider_uses_mova_config_and_heterogeneous_spec(monkeypatch):
+    args = _provider_args(mova_num_value_experts=64)
+    fake_config_class = type("FakeMoVATransformerConfig", (), {})
+    fake_config = SimpleNamespace(hidden_size=2560)
+    config_builder = MagicMock(return_value=fake_config)
+    spec_builder = MagicMock(return_value="mova-block-spec")
+    gpt_model = MagicMock(return_value=SimpleNamespace(config=fake_config))
+
+    monkeypatch.setattr(model_provider_module, "core_transformer_config_from_args", config_builder)
+    monkeypatch.setattr(
+        model_provider_module,
+        "_get_mova_model_components",
+        lambda: (fake_config_class, spec_builder),
+    )
+    monkeypatch.setattr(model_provider_module, "GPTModel", gpt_model)
+
+    provider = model_provider_module.get_model_provider_func(args, role="actor")
+    model = provider(pre_process=False, post_process=True, vp_stage=2)
+
+    assert model.config is fake_config
+    config_builder.assert_called_once_with(args, fake_config_class)
+    spec_builder.assert_called_once_with(
+        fake_config,
+        use_transformer_engine=True,
+        moe_grouped_gemm=True,
+        moe_use_legacy_grouped_gemm=False,
+        vp_stage=2,
+    )
+    gpt_model.assert_called_once_with(
+        config=fake_config,
+        transformer_layer_spec="mova-block-spec",
+        vocab_size=256,
+        max_sequence_length=32768,
+        pre_process=False,
+        post_process=True,
+        fp16_lm_cross_entropy=False,
+        parallel_output=True,
+        share_embeddings_and_output_weights=False,
+        position_embedding_type="rope",
+        rotary_percent=1.0,
+        rotary_base=10_000_000,
+        rope_scaling=False,
+        vp_stage=2,
+    )
+
+
+@pytest.mark.parametrize("role", ["actor", "critic"])
+def test_actor_and_critic_share_native_mova_provider(monkeypatch, role):
+    args = _provider_args(mova_num_value_experts=64)
+    fake_config_class = type("FakeMoVATransformerConfig", (), {})
+    fake_config = SimpleNamespace(hidden_size=2560, sequence_parallel=True)
+    spec_builder = MagicMock(return_value="mova-block-spec")
+    model = SimpleNamespace(config=fake_config)
+    model.output_layer = "original-output-layer"
+
+    monkeypatch.setattr(
+        model_provider_module,
+        "core_transformer_config_from_args",
+        MagicMock(return_value=fake_config),
+    )
+    monkeypatch.setattr(
+        model_provider_module,
+        "_get_mova_model_components",
+        lambda: (fake_config_class, spec_builder),
+    )
+    monkeypatch.setattr(model_provider_module, "GPTModel", MagicMock(return_value=model))
+    critic_head = MagicMock(return_value="critic-output-layer")
+    monkeypatch.setattr(model_provider_module, "LinearForLastLayer", critic_head)
+
+    provider = model_provider_module.get_model_provider_func(args, role=role)
+    result = provider()
+
+    assert result is model
+    spec_builder.assert_called_once()
+    if role == "critic":
+        critic_head.assert_called_once_with(
+            input_size=2560,
+            output_size=1,
+            config=fake_config,
+        )
+        assert model.output_layer == "critic-output-layer"
+    else:
+        critic_head.assert_not_called()
+        assert model.output_layer == "original-output-layer"
+
+
+def test_non_mova_provider_preserves_standard_transformer_path(monkeypatch):
+    args = _provider_args(mova_num_value_experts=0)
+    fake_config = SimpleNamespace(hidden_size=2560)
+    config_builder = MagicMock(return_value=fake_config)
+    te_spec_builder = MagicMock(return_value="standard-te-layer-spec")
+    gpt_model = MagicMock(return_value=SimpleNamespace(config=fake_config))
+    mova_components = MagicMock(side_effect=AssertionError("legacy path imported MoVA"))
+
+    monkeypatch.setattr(model_provider_module, "core_transformer_config_from_args", config_builder)
+    monkeypatch.setattr(model_provider_module, "get_gpt_layer_with_transformer_engine_spec", te_spec_builder)
+    monkeypatch.setattr(model_provider_module, "_get_mova_model_components", mova_components)
+    monkeypatch.setattr(model_provider_module, "GPTModel", gpt_model)
+
+    provider = model_provider_module.get_model_provider_func(args, role="actor")
+    provider(pre_process=True, post_process=False)
+
+    config_builder.assert_called_once_with(args)
+    mova_components.assert_not_called()
+    te_spec_builder.assert_called_once_with(
+        num_experts=0,
+        moe_grouped_gemm=True,
+        qk_layernorm=False,
+        multi_latent_attention=False,
+        moe_use_legacy_grouped_gemm=False,
+    )
+    assert gpt_model.call_args.kwargs["transformer_layer_spec"] == "standard-te-layer-spec"
+    assert "vp_stage" not in gpt_model.call_args.kwargs
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("megatron_to_hf_mode", "bridge", "raw"),
+        ("custom_model_provider_path", "some.module.provider", "owns its model provider"),
+        ("spec", "some.module.spec", "heterogeneous layer spec"),
+        ("mtp_num_layers", 1, "MTP"),
+    ],
+)
+def test_mova_provider_fails_before_selecting_incompatible_provider(field, value, message):
+    args = _provider_args(mova_num_value_experts=64)
+    setattr(args, field, value)
+
+    with pytest.raises(ValueError, match=message):
+        model_provider_module.get_model_provider_func(args)
+
+
+def test_mova_components_report_incompatible_megatron_cleanly(monkeypatch):
+    real_import = __import__
+
+    def reject_mova(name, *args, **kwargs):
+        if name in {
+            "megatron.core.models.gpt.mova_layer_specs",
+            "megatron.core.transformer.mova",
+        }:
+            raise ImportError("no native MoVA")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr("builtins.__import__", reject_mova)
+
+    with pytest.raises(RuntimeError, match="does not provide MoVATransformerConfig"):
+        model_provider_module._get_mova_model_components()
diff --git a/tests/fast/utils/test_arguments.py b/tests/fast/utils/test_arguments.py
index aa2c35bd311..3bd0efe46c5 100644
--- a/tests/fast/utils/test_arguments.py
+++ b/tests/fast/utils/test_arguments.py
@@ -5,13 +5,159 @@
 
 import pytest
 
-from miles.utils.arguments import _maybe_apply_dumper_overrides, get_miles_extra_args_provider
+from miles.utils.arguments import (
+    _maybe_apply_dumper_overrides,
+    get_miles_extra_args_provider,
+    validate_mova_args,
+)
 from miles.utils.misc import function_registry
 
 PATH_ARGS = ["--rollout-function-path", "--custom-generate-function-path"]
 REQUIRED_ARGS = ["--rollout-batch-size", "64"]
 
 
+class TestMoVAArguments:
+    def _parse(self, flags):
+        with patch.object(sys, "argv", ["test", *flags]):
+            parser = argparse.ArgumentParser()
+            get_miles_extra_args_provider()(parser)
+            return parser.parse_args(flags)
+
+    def test_defaults_match_native_megatron_entry_point(self):
+        args = self._parse(REQUIRED_ARGS)
+
+        assert args.mova_num_value_experts == 0
+        assert args.mova_router_topk == 1
+        assert args.mova_router_score_function == "sigmoid"
+        assert args.mova_router_topk_scaling_factor == 1.0
+        assert args.mova_router_enable_expert_bias is False
+        assert args.mova_router_bias_update_rate == 1.0e-3
+        assert args.mova_router_aux_loss_coeff == 0.0
+        assert args.mova_router_load_balancing_type == "none"
+        assert args.mova_num_dense_layers == 0
+        assert args.mova_norm_num_groups == 1
+        assert args.mova_attention_gate_function == "softplus"
+        assert args.mova_value_backend == "grouped_gemm"
+        assert args.mova_use_torch_rms_norm is False
+        assert args.xllm_router_compatibility is True
+        assert args.xllm_router_gemm_partitions == 1
+
+    def test_generated_k2mova_rl_flags_are_accepted(self):
+        args = self._parse(
+            REQUIRED_ARGS
+            + [
+                "--mova-num-value-experts",
+                "64",
+                "--mova-router-topk",
+                "4",
+                "--mova-router-score-function",
+                "sigmoid",
+                "--mova-router-topk-scaling-factor",
+                "2.5",
+                "--mova-router-enable-expert-bias",
+                "--mova-router-bias-update-rate",
+                "0.001",
+                "--mova-router-aux-loss-coeff",
+                "0",
+                "--mova-router-load-balancing-type",
+                "none",
+                "--mova-num-dense-layers",
+                "3",
+                "--mova-norm-num-groups",
+                "2",
+                "--mova-attention-gate-function",
+                "softplus",
+                "--mova-value-backend",
+                "grouped_gemm",
+                "--no-mova-use-torch-rms-norm",
+                "--xllm-router-compatibility",
+                "--xllm-router-gemm-partitions",
+                "1",
+            ]
+        )
+
+        assert args.mova_num_value_experts == 64
+        assert args.mova_router_topk == 4
+        assert args.mova_router_topk_scaling_factor == 2.5
+        assert args.mova_router_enable_expert_bias is True
+        assert args.mova_num_dense_layers == 3
+        assert args.mova_norm_num_groups == 2
+        assert args.mova_value_backend == "grouped_gemm"
+        assert args.mova_use_torch_rms_norm is False
+        assert args.xllm_router_compatibility is True
+
+    def test_registration_does_not_change_shared_transformer_defaults(self):
+        with patch.object(sys, "argv", ["test", *REQUIRED_ARGS]):
+            parser = argparse.ArgumentParser()
+            parser.add_argument("--attention-dropout", type=float, default=0.1)
+            parser.add_argument("--hidden-dropout", type=float, default=0.1)
+            get_miles_extra_args_provider()(parser)
+            args, _ = parser.parse_known_args(REQUIRED_ARGS)
+
+        assert args.attention_dropout == 0.1
+        assert args.hidden_dropout == 0.1
+
+    def test_native_validation_accepts_exact_k2mova_contract(self):
+        args = SimpleNamespace(
+            mova_num_value_experts=64,
+            train_backend="megatron",
+            megatron_to_hf_mode="raw",
+            custom_model_provider_path=None,
+            use_legacy_models=False,
+            yaml_cfg=None,
+            spec=None,
+            mtp_num_layers=None,
+            multi_latent_attention=False,
+            heterogeneous_layers_config_path=None,
+            attention_output_gate=True,
+            rotary_interleaved=True,
+            group_query_attention=True,
+            num_experts=100,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            tensor_model_parallel_size=8,
+            sequence_parallel=True,
+        )
+
+        validate_mova_args(args)
+
+    @pytest.mark.parametrize(
+        ("field", "value", "message"),
+        [
+            ("megatron_to_hf_mode", "bridge", "raw"),
+            ("attention_output_gate", False, "attention-output-gate"),
+            ("rotary_interleaved", False, "rotary-interleaved"),
+            ("attention_dropout", 0.1, "attention-dropout"),
+            ("sequence_parallel", False, "sequence-parallel"),
+        ],
+    )
+    def test_native_validation_rejects_incompatible_rl_flags(self, field, value, message):
+        args = SimpleNamespace(
+            mova_num_value_experts=64,
+            train_backend="megatron",
+            megatron_to_hf_mode="raw",
+            custom_model_provider_path=None,
+            use_legacy_models=False,
+            yaml_cfg=None,
+            spec=None,
+            mtp_num_layers=None,
+            multi_latent_attention=False,
+            heterogeneous_layers_config_path=None,
+            attention_output_gate=True,
+            rotary_interleaved=True,
+            group_query_attention=True,
+            num_experts=100,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            tensor_model_parallel_size=8,
+            sequence_parallel=True,
+        )
+        setattr(args, field, value)
+
+        with pytest.raises(ValueError, match=message):
+            validate_mova_args(args)
+
+
 def make_class_with_add_arguments():
     class MyFn:
         @classmethod

From 436736266044dbcc2bd5e6a863b1ab9f12d6cbc4 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 01:32:25 +0000
Subject: [PATCH 3/7] Verify every rollout weight version

Add a cheap opt-in acceptance check that queries every rollout engine after each synchronization. Keep the existing random CI sample as the default and avoid full tensor comparisons.
---
 .../backends/experimental/fsdp_utils/actor.py |  9 +++-
 miles/backends/megatron_utils/actor.py        | 19 ++++++---
 miles/backends/training_utils/ci_utils.py     | 25 +++++++++++
 .../test_weight_version_validation.py         | 41 +++++++++++++++++++
 4 files changed, 86 insertions(+), 8 deletions(-)
 create mode 100644 tests/fast/backends/training_utils/test_weight_version_validation.py

diff --git a/miles/backends/experimental/fsdp_utils/actor.py b/miles/backends/experimental/fsdp_utils/actor.py
index 47c2540f98d..5e878129cac 100644
--- a/miles/backends/experimental/fsdp_utils/actor.py
+++ b/miles/backends/experimental/fsdp_utils/actor.py
@@ -21,7 +21,7 @@
 from miles.utils.tracking_utils import init_tracking
 
 from ....utils.profile_utils import TrainProfiler
-from ...training_utils.ci_utils import check_grad_norm
+from ...training_utils.ci_utils import assert_rollout_engine_weight_versions, check_grad_norm
 from ...training_utils.data import DataIterator, get_batch, get_data_iterator, get_rollout_data
 from ...training_utils.log_utils import (
     aggregate_forward_results,
@@ -569,7 +569,12 @@ def update_weights(self) -> None:  # type: ignore[override]
 
         self.weight_updater.update_weights()
 
-        if self.args.ci_test and len(rollout_engines) > 0:
+        if getattr(self.args, "check_all_engine_weight_versions", False):
+            assert_rollout_engine_weight_versions(
+                rollout_engines,
+                self.weight_updater.weight_version,
+            )
+        elif self.args.ci_test and len(rollout_engines) > 0:
             engine = random.choice(rollout_engines)
             engine_version = ray.get(engine.get_weight_version.remote())
             if str(engine_version) != str(self.weight_updater.weight_version):
diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index c138ea87460..08f771614aa 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -29,6 +29,7 @@
 from ...utils.profile_utils import TrainProfiler
 from ...utils.tensor_backper import TensorBackuper
 from ..training_utils.cp_utils import slice_with_cp
+from ..training_utils.ci_utils import assert_rollout_engine_weight_versions
 from ..training_utils.data import DataIterator, get_data_iterator, get_rollout_data, sync_actor_critic_data
 from ..training_utils.log_utils import log_cpu_memory, log_perf_data, log_rollout_data
 from ..training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values
@@ -1013,13 +1014,19 @@ def update_weights(self) -> None:
             self.weight_updater.update_weights()
             print_memory("after update_weights")
 
-            if self.args.ci_test and len(rollout_engines) > 0 and not is_lora_enabled(self.args):
-                engine = random.choice(rollout_engines)
-                engine_version = ray.get(engine.get_weight_version.remote())
-                if str(engine_version) != str(self.weight_updater.weight_version):
-                    raise RuntimeError(
-                        f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+            if not is_lora_enabled(self.args):
+                if getattr(self.args, "check_all_engine_weight_versions", False):
+                    assert_rollout_engine_weight_versions(
+                        rollout_engines,
+                        self.weight_updater.weight_version,
                     )
+                elif self.args.ci_test and len(rollout_engines) > 0:
+                    engine = random.choice(rollout_engines)
+                    engine_version = ray.get(engine.get_weight_version.remote())
+                    if str(engine_version) != str(self.weight_updater.weight_version):
+                        raise RuntimeError(
+                            f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                        )
 
             if getattr(self.args, "keep_old_actor", False):
                 if self.args.update_weights_interval == 1:
diff --git a/miles/backends/training_utils/ci_utils.py b/miles/backends/training_utils/ci_utils.py
index 9124889beda..3bd4eaf0775 100644
--- a/miles/backends/training_utils/ci_utils.py
+++ b/miles/backends/training_utils/ci_utils.py
@@ -9,6 +9,31 @@
 logger = logging.getLogger(__name__)
 
 
+def assert_rollout_engine_weight_versions(rollout_engines, expected_version) -> None:
+    """Assert that every rollout engine exposes the synchronized version.
+
+    This is intentionally metadata-only: it is cheap enough to enable for
+    acceptance runs and avoids the cost of full tensor equality checks.
+    """
+
+    if not rollout_engines:
+        return
+
+    import ray
+
+    versions = ray.get([engine.get_weight_version.remote() for engine in rollout_engines])
+    mismatches = [
+        (index, version)
+        for index, version in enumerate(versions)
+        if str(version) != str(expected_version)
+    ]
+    if mismatches:
+        raise RuntimeError(
+            "Rollout engine weight-version mismatch: "
+            f"expected {expected_version}, mismatches {mismatches}"
+        )
+
+
 def check_kl(args: Namespace, log_dict: dict[str, float], step_id: int, accumulated_step_id: int) -> None:
     if step_id == 0 and "train/ppo_kl" in log_dict and "train/pg_clipfrac" in log_dict:
         if args.multi_latent_attention:
diff --git a/tests/fast/backends/training_utils/test_weight_version_validation.py b/tests/fast/backends/training_utils/test_weight_version_validation.py
new file mode 100644
index 00000000000..5a5ed6f2516
--- /dev/null
+++ b/tests/fast/backends/training_utils/test_weight_version_validation.py
@@ -0,0 +1,41 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from miles.backends.training_utils import ci_utils
+
+
+def _engine(version):
+    remote = MagicMock(return_value=f"version-ref-{version}")
+    return SimpleNamespace(get_weight_version=SimpleNamespace(remote=remote))
+
+
+def test_all_rollout_engine_versions_are_queried(monkeypatch):
+    engines = [_engine("7"), _engine("7"), _engine("7")]
+    ray_get = MagicMock(return_value=["7", 7, "7"])
+    monkeypatch.setattr("ray.get", ray_get)
+
+    ci_utils.assert_rollout_engine_weight_versions(engines, expected_version=7)
+
+    assert [engine.get_weight_version.remote.call_count for engine in engines] == [1, 1, 1]
+    ray_get.assert_called_once_with(
+        ["version-ref-7", "version-ref-7", "version-ref-7"]
+    )
+
+
+def test_all_rollout_engine_versions_report_every_mismatch(monkeypatch):
+    engines = [_engine("7"), _engine("6"), _engine("stale")]
+    monkeypatch.setattr("ray.get", MagicMock(return_value=["7", "6", "stale"]))
+
+    with pytest.raises(RuntimeError, match=r"expected 7.*\(1, '6'\).*\(2, 'stale'\)"):
+        ci_utils.assert_rollout_engine_weight_versions(engines, expected_version=7)
+
+
+def test_all_rollout_engine_versions_accepts_empty_engine_set(monkeypatch):
+    ray_get = MagicMock()
+    monkeypatch.setattr("ray.get", ray_get)
+
+    ci_utils.assert_rollout_engine_weight_versions([], expected_version=7)
+
+    ray_get.assert_not_called()

From 1206e623659f72610df7296a38b05a4452525b87 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 03:55:48 +0000
Subject: [PATCH 4/7] Fail closed on P2P transfer errors

Weight-version metadata must never advance after a failed or timed-out RDMA write. Drain every submitted transfer, clear the queue reliably, and surface aggregate failures to abort the update before rollout resumes.
---
 .../p2p_transfer_utils.py                     | 34 +++++++---
 .../test_p2p_transfer_manager.py              | 66 +++++++++++++++++++
 2 files changed, 92 insertions(+), 8 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py

diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py
index 804d73ce043..5d75a35b1c9 100644
--- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py
+++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p_transfer_utils.py
@@ -172,14 +172,32 @@ def submit_returning_future(self, fn: Callable, *args) -> torch.Future:
         return future
 
     def wait_transfers(self) -> None:
-        """Wait for all submitted tasks to complete."""
-        for future in self.transfer_futures:
-            try:
-                future.result(timeout=self.transfer_timeout)
-            except Exception as e:
-                logger.error(f"[P2P] Transfer future failed: {e}")
-
-        self.transfer_futures.clear()
+        """Drain submitted tasks and fail the update if any transfer failed."""
+        failures: list[tuple[int, Exception]] = []
+        futures = tuple(self.transfer_futures)
+        try:
+            for index, future in enumerate(futures):
+                try:
+                    future.result(timeout=self.transfer_timeout)
+                except Exception as error:
+                    failures.append((index, error))
+                    logger.error(
+                        "[P2P] Transfer task %d/%d failed: %s",
+                        index + 1,
+                        len(futures),
+                        error,
+                    )
+        finally:
+            self.transfer_futures.clear()
+
+        if failures:
+            details = "; ".join(
+                f"task {index + 1}: {type(error).__name__}: {error}" for index, error in failures
+            )
+            first_error = failures[0][1]
+            raise RuntimeError(
+                f"{len(failures)} of {len(futures)} P2P weight transfers failed ({details})"
+            ) from first_error
 
 
 def create_server_args_from_dict(data_dict: dict) -> ServerArgs:
diff --git a/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py b/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py
new file mode 100644
index 00000000000..a01dcc3627d
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_p2p_transfer_manager.py
@@ -0,0 +1,66 @@
+from concurrent.futures import Future
+from unittest.mock import MagicMock
+
+import pytest
+
+from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p_transfer_utils import (
+    P2PTransferManager,
+)
+
+
+def _completed_future() -> Future:
+    future = Future()
+    future.set_result(None)
+    return future
+
+
+def test_wait_transfers_drains_successful_tasks_and_clears_queue() -> None:
+    manager = P2PTransferManager(transfer_timeout=7.0)
+    first = MagicMock(wraps=_completed_future())
+    second = MagicMock(wraps=_completed_future())
+    manager.transfer_futures = [first, second]
+
+    manager.wait_transfers()
+
+    first.result.assert_called_once_with(timeout=7.0)
+    second.result.assert_called_once_with(timeout=7.0)
+    assert manager.transfer_futures == []
+
+
+def test_wait_transfers_drains_all_tasks_then_raises_aggregate_error() -> None:
+    manager = P2PTransferManager(transfer_timeout=3.0)
+    first = MagicMock()
+    second = MagicMock()
+    third = MagicMock()
+    first.result.side_effect = RuntimeError("session-a failed")
+    second.result.return_value = None
+    third.result.side_effect = ValueError("session-c failed")
+    manager.transfer_futures = [first, second, third]
+
+    with pytest.raises(RuntimeError, match=r"2 of 3 P2P weight transfers failed") as exc_info:
+        manager.wait_transfers()
+
+    first.result.assert_called_once_with(timeout=3.0)
+    second.result.assert_called_once_with(timeout=3.0)
+    third.result.assert_called_once_with(timeout=3.0)
+    assert "session-a failed" in str(exc_info.value)
+    assert "session-c failed" in str(exc_info.value)
+    assert isinstance(exc_info.value.__cause__, RuntimeError)
+    assert manager.transfer_futures == []
+
+
+def test_wait_transfers_propagates_timeout_and_clears_queue() -> None:
+    manager = P2PTransferManager(transfer_timeout=0.25)
+    timed_out = MagicMock()
+    completed = MagicMock()
+    timed_out.result.side_effect = TimeoutError("transfer timed out")
+    completed.result.return_value = None
+    manager.transfer_futures = [timed_out, completed]
+
+    with pytest.raises(RuntimeError, match=r"1 of 2 P2P weight transfers failed") as exc_info:
+        manager.wait_transfers()
+
+    timed_out.result.assert_called_once_with(timeout=0.25)
+    completed.result.assert_called_once_with(timeout=0.25)
+    assert isinstance(exc_info.value.__cause__, TimeoutError)
+    assert manager.transfer_futures == []

From 816f92af8fe77a70be76093289c7f31644790959 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 16:56:45 +0000
Subject: [PATCH 5/7] Reject zero rollout temperature during training

Training recomputes policy log probabilities by dividing logits by the rollout temperature, so greedy temperature zero creates NaNs only after an expensive rollout. Validate before runtime setup while retaining zero-temperature support for rollout-only generation.
---
 miles/utils/arguments.py                     |  2 ++
 miles/utils/rollout_temperature.py           | 13 +++++++
 tests/fast/utils/test_rollout_temperature.py | 38 ++++++++++++++++++++
 3 files changed, 53 insertions(+)
 create mode 100644 miles/utils/rollout_temperature.py
 create mode 100644 tests/fast/utils/test_rollout_temperature.py

diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index ca683c17af2..a1a3507ad15 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -15,6 +15,7 @@
 from miles.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list
 from miles.utils.logging_utils import configure_logger
 from miles.utils.misc import load_function
+from miles.utils.rollout_temperature import validate_rollout_temperature
 
 logger = logging.getLogger(__name__)
 
@@ -1973,6 +1974,7 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 
 def miles_validate_args(args):
     validate_mova_args(args)
+    validate_rollout_temperature(args)
     args.eval_datasets = _resolve_eval_datasets(args)
 
     # Normalize --tito-allowed-append-roles: lowercase + deduplicate.
diff --git a/miles/utils/rollout_temperature.py b/miles/utils/rollout_temperature.py
new file mode 100644
index 00000000000..c9ea00ea249
--- /dev/null
+++ b/miles/utils/rollout_temperature.py
@@ -0,0 +1,13 @@
+from argparse import Namespace
+
+
+def validate_rollout_temperature(args: Namespace) -> None:
+    """Reject temperatures that make training log-prob recomputation undefined."""
+    if args.rollout_temperature < 0:
+        raise ValueError("--rollout-temperature must be non-negative for generation.")
+    if args.rollout_temperature == 0 and not args.debug_rollout_only:
+        raise ValueError(
+            "--rollout-temperature must be greater than 0 for training because "
+            "Miles divides policy logits by it during log-probability recomputation. "
+            "Temperature 0 is only supported with --debug-rollout-only for greedy generation."
+        )
diff --git a/tests/fast/utils/test_rollout_temperature.py b/tests/fast/utils/test_rollout_temperature.py
new file mode 100644
index 00000000000..fb57b1a7ca8
--- /dev/null
+++ b/tests/fast/utils/test_rollout_temperature.py
@@ -0,0 +1,38 @@
+import unittest
+from argparse import Namespace
+
+from miles.utils.rollout_temperature import validate_rollout_temperature
+
+
+class TestRolloutTemperature(unittest.TestCase):
+    def test_training_rejects_zero(self) -> None:
+        args = Namespace(rollout_temperature=0.0, debug_rollout_only=False)
+
+        with self.assertRaisesRegex(ValueError, "greater than 0 for training"):
+            validate_rollout_temperature(args)
+
+    def test_training_rejects_negative(self) -> None:
+        args = Namespace(rollout_temperature=-1.0, debug_rollout_only=False)
+
+        with self.assertRaisesRegex(ValueError, "non-negative for generation"):
+            validate_rollout_temperature(args)
+
+    def test_rollout_only_allows_greedy_zero(self) -> None:
+        args = Namespace(rollout_temperature=0.0, debug_rollout_only=True)
+
+        validate_rollout_temperature(args)
+
+    def test_rollout_only_rejects_negative(self) -> None:
+        args = Namespace(rollout_temperature=-1.0, debug_rollout_only=True)
+
+        with self.assertRaisesRegex(ValueError, "non-negative for generation"):
+            validate_rollout_temperature(args)
+
+    def test_training_accepts_positive(self) -> None:
+        args = Namespace(rollout_temperature=1.0, debug_rollout_only=False)
+
+        validate_rollout_temperature(args)
+
+
+if __name__ == "__main__":
+    unittest.main()

From 009cd8a47be97d0888cce0eee2a85472c7c1c128 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 19:16:04 +0000
Subject: [PATCH 6/7] Propagate rollout dtype to P2P CPU replicas

P2P builds a sharded SGLang model locally before transferring weights. Passing only the model path discarded the rollout server dtype, so float32 HF metadata silently resolved auto to fp16 even when live engines explicitly used bf16. Preserve the queried rollout dtype and pin both explicit and auto behavior in regression coverage.
---
 .../update_weight_from_distributed/p2p.py     |  2 +-
 .../test_xllm_mova_weight_sync.py             | 55 +++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py
index ba6cebde7e1..60b17cfc188 100644
--- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py
+++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py
@@ -267,7 +267,7 @@ def create_cpu_replica(
         initialize_fp4_gemm_config(server_args)
         with ParallelismContext(parallelism_config):
             model = get_model(
-                model_config=ModelConfig(model_path),
+                model_config=ModelConfig(model_path, dtype=server_args.dtype),
                 load_config=load_config,
                 device_config=DeviceConfig(),
             )
diff --git a/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py
index 438e4edd9a9..6e78558577c 100644
--- a/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py
+++ b/tests/fast/backends/megatron_utils/test_xllm_mova_weight_sync.py
@@ -377,6 +377,61 @@ def test_p2p_staging_contract_requires_all_value_expert_shards() -> None:
     assert updater._tensor_update_pending == {}
 
 
+@pytest.mark.parametrize(
+    ("runtime_dtype", "expected_dtype"),
+    [
+        ("bfloat16", torch.bfloat16),
+        ("auto", torch.float16),
+    ],
+)
+def test_p2p_cpu_replica_uses_runtime_rollout_dtype(runtime_dtype: str, expected_dtype: torch.dtype) -> None:
+    from contextlib import nullcontext
+
+    from sglang.srt.configs.model_config import _get_and_verify_dtype
+
+    from miles.backends.megatron_utils.update_weight.update_weight_from_distributed import p2p
+
+    updater = object.__new__(p2p.UpdateWeightP2P)
+    updater._shared_params_dict = {}
+    server_args = SimpleNamespace(dtype=runtime_dtype, rl_quant_profile=None)
+    observed = {}
+
+    def fake_model_config(model_path: str, dtype: str = "auto") -> SimpleNamespace:
+        # The MoVA HF artifact declares float32. SGLang resolves that to fp16
+        # for auto, but must honor an explicit bfloat16 rollout dtype.
+        hf_config = {"model_type": "xllm", "torch_dtype": "float32"}
+        observed["model_path"] = model_path
+        observed["requested_dtype"] = dtype
+        return SimpleNamespace(dtype=_get_and_verify_dtype(hf_config, dtype))
+
+    def fake_get_model(*, model_config, load_config, device_config) -> nn.Module:
+        observed["validation_dtype"] = model_config.dtype
+        return nn.Module()
+
+    with (
+        patch.object(p2p, "ModelConfig", side_effect=fake_model_config),
+        patch.object(p2p, "LoadConfig", return_value=object()),
+        patch.object(p2p, "DeviceConfig", return_value=object()),
+        patch.object(p2p, "ParallelismContext", side_effect=lambda _: nullcontext()),
+        patch.object(p2p, "get_model", side_effect=fake_get_model),
+        patch.object(p2p, "initialize_moe_config"),
+        patch.object(p2p, "initialize_fp8_gemm_config"),
+        patch.object(p2p, "initialize_fp4_gemm_config"),
+        patch.object(p2p.torch.cuda, "empty_cache"),
+    ):
+        updater.create_cpu_replica(
+            parallelism_config=object(),
+            model_path="/models/mova-float32-config",
+            server_args=server_args,
+        )
+
+    assert observed == {
+        "model_path": "/models/mova-float32-config",
+        "requested_dtype": runtime_dtype,
+        "validation_dtype": expected_dtype,
+    }
+
+
 def test_broadcast_path_preserves_all_converted_value_expert_metadata() -> None:
     from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (
         update_weights_from_distributed,

From 34460b4fa528a8b8015f252fe2fd738a87e606a0 Mon Sep 17 00:00:00 2001
From: Yash Akhauri <akhauri.yash@gmail.com>
Date: Thu, 13 Aug 2026 21:01:52 +0000
Subject: [PATCH 7/7] Preserve rollout filter mask provenance

Rollout filters intentionally zero loss masks, but the fail-fast trainer validator could not distinguish those samples from corrupted data. Carry the filter decision through DP sharding, accept only explicitly filtered all-zero masks, and keep unexplained zeros fail-closed.
---
 docs/en/get_started/customization.md          |   6 +-
 miles/backends/megatron_utils/actor.py        |  69 +++++--
 miles/ray/rollout.py                          |   5 +
 .../test_rollout_filter_mask_validation.py    | 185 ++++++++++++++++++
 4 files changed, 253 insertions(+), 12 deletions(-)
 create mode 100644 tests/fast/backends/megatron_utils/test_rollout_filter_mask_validation.py

diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md
index 8aa63c23fbd..7e0b0909dae 100644
--- a/docs/en/get_started/customization.md
+++ b/docs/en/get_started/customization.md
@@ -170,6 +170,11 @@ def filter_function(args, samples: list[Sample]) -> None
 
 **Note**: This function should directly modify the `remove_sample` attribute of each `Sample` object.
 
+The default sample-to-training-data converter propagates this decision as
+`removed_by_filter`. A custom `--custom-convert-samples-to-train-data-path`
+that emits an all-zero loss mask must also emit a parallel boolean
+`removed_by_filter` field. The trainer rejects unexplained all-zero masks.
+
 **Use Cases**:
 - Filtering samples based on response quality
 - Implementing selective training strategies
@@ -439,4 +444,3 @@ For detailed explanation of R3 and MilesRouter, see [Miles Router](../advanced/m
 def custom_model_provider(pre_process: bool, post_process: bool, vp_stage: int | None = None) -> GPTModel
 ```
 
-
diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index 08f771614aa..1e861604ba1 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -191,6 +191,13 @@ def _sum_float(x):
             return float(sum(float(v) for v in x))
         return float(x)
 
+    def _all_zero(x):
+        if torch.is_tensor(x):
+            return bool(torch.all(x == 0).item())
+        if isinstance(x, (list, tuple)):
+            return all(float(v) == 0 for v in x)
+        return float(x) == 0
+
     def _summarize_vector_list(key, limit=3):
         if not _present(key):
             return f"{key}=MISSING"
@@ -228,6 +235,7 @@ def _basic_batch_summary():
             "rewards",
             "response_lengths",
             "total_lengths",
+            "removed_by_filter",
             "loss_masks",
             "tokens",
             "input_ids",
@@ -242,6 +250,16 @@ def _basic_batch_summary():
             lines.append(_summarize_vector_list(key))
 
         # Numeric aggregate summary.
+        try:
+            if _present("removed_by_filter") and _is_seq(rollout_data["removed_by_filter"]):
+                flags = rollout_data["removed_by_filter"]
+                lines.append(
+                    f"removed_by_filter: count={len(flags)} "
+                    f"removed={sum(bool(flag) for flag in flags)}"
+                )
+        except Exception as e:
+            lines.append(f"removed_by_filter aggregate failed: {type(e).__name__}: {e}")
+
         try:
             if _present("response_lengths") and _is_seq(rollout_data["response_lengths"]):
                 rs = [int(x) for x in rollout_data["response_lengths"]]
@@ -329,6 +347,28 @@ def _basic_batch_summary():
         if got != n:
             _add_error(f"{key!r} length mismatch: got {got}, expected {n}")
 
+    removed_by_filter_flags = [False] * n
+    if _present("removed_by_filter"):
+        candidate_flags = rollout_data["removed_by_filter"]
+        candidate_flags_valid = True
+        if not _is_seq(candidate_flags):
+            _add_error(
+                f"'removed_by_filter' must be list/tuple, got {type(candidate_flags).__name__}"
+            )
+            candidate_flags_valid = False
+        elif len(candidate_flags) != n:
+            _add_error(f"'removed_by_filter' length mismatch: got {len(candidate_flags)}, expected {n}")
+            candidate_flags_valid = False
+        else:
+            for i, removed in enumerate(candidate_flags):
+                if not isinstance(removed, bool):
+                    _add_error(
+                        f"removed_by_filter[{i}] must be bool, got {type(removed).__name__}"
+                    )
+                    candidate_flags_valid = False
+            if candidate_flags_valid:
+                removed_by_filter_flags = list(candidate_flags)
+
     token_key = None
     if _present("tokens"):
         token_key = "tokens"
@@ -411,8 +451,24 @@ def _basic_batch_summary():
             continue
 
         mask_sum = _sum_float(mask)
-        if mask_sum <= 0:
-            _add_error(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}")
+        removed_by_filter = removed_by_filter_flags[i]
+        mask_is_all_zero = _all_zero(mask)
+        if removed_by_filter:
+            if mask_is_all_zero:
+                _add_warning(
+                    f"loss_masks[{i}] has no active tokens because removed_by_filter=True, "
+                    f"sum={mask_sum}, response_len={resp}"
+                )
+            else:
+                _add_error(
+                    f"loss_masks[{i}] is not all-zero despite removed_by_filter=True, "
+                    f"sum={mask_sum}, response_len={resp}"
+                )
+        elif mask_sum <= 0:
+            _add_error(
+                f"loss_masks[{i}] has no active tokens without removed_by_filter=True, "
+                f"sum={mask_sum}, response_len={resp}"
+            )
         if mask_sum > resp:
             # Warning-only: float/weighted masks can legitimately have sum > resp.
             _add_warning(
@@ -484,15 +540,6 @@ def check_vector_list(key, expected_lengths):
         check_vector_list(key, response_lengths)
     check_vector_list("rollout_log_probs", [] if cp_size > 1 else response_lengths)
 
-    # GRPO grouping diagnostics. Warning only because dynamic filtering can alter counts.
-    n_samples_per_prompt = int(getattr(args, "n_samples_per_prompt", 0) or 0)
-    if n_samples_per_prompt > 0 and n % n_samples_per_prompt != 0:
-        _add_warning(f"sample count {n} not divisible by n_samples_per_prompt={n_samples_per_prompt}")
-
-    grpo_group_size = int(getattr(args, "grpo_group_size", 0) or 0)
-    if grpo_group_size > 0 and n % grpo_group_size != 0:
-        _add_warning(f"sample count {n} not divisible by grpo_group_size={grpo_group_size}")
-
     # This is important for your failure mode:
     # If compute_advantages_and_returns will normalize, every rank that reaches
     # it must have log_probs/values in the same structural state.
diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py
index baec979375d..382031ba71c 100644
--- a/miles/ray/rollout.py
+++ b/miles/ray/rollout.py
@@ -725,6 +725,10 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
             "rewards": rewards,
             "raw_reward": raw_rewards,
             "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples],
+            # Preserve why an all-zero loss mask is intentional. Rollout sample
+            # filters mark Sample.remove_sample; the trainer validator must be
+            # able to distinguish that supported state from a corrupted mask.
+            "removed_by_filter": [bool(sample.remove_sample) for sample in samples],
             "sample_indices": [sample.index for sample in samples],
         }
 
@@ -880,6 +884,7 @@ def _stat(xs):
                 "response_lengths",
                 "rewards",
                 "truncated",
+                "removed_by_filter",
                 "loss_masks",
                 "round_number",
                 "sample_indices",
diff --git a/tests/fast/backends/megatron_utils/test_rollout_filter_mask_validation.py b/tests/fast/backends/megatron_utils/test_rollout_filter_mask_validation.py
new file mode 100644
index 00000000000..d947a04c766
--- /dev/null
+++ b/tests/fast/backends/megatron_utils/test_rollout_filter_mask_validation.py
@@ -0,0 +1,185 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from miles.backends.megatron_utils import actor
+from miles.ray import rollout
+from miles.utils.types import Sample
+
+
+def _validator_args():
+    return SimpleNamespace(
+        advantage_estimator="grpo",
+        normalize_advantages=False,
+        use_rollout_logprobs=False,
+        use_critic=False,
+        qkv_format="thd",
+        compute_advantages_and_returns=True,
+        n_samples_per_prompt=8,
+        grpo_group_size=2,
+        generate_multi_samples=None,
+        global_batch_size=32,
+        data_parallel_size=8,
+        context_parallel_size=1,
+        tensor_model_parallel_size=2,
+        pipeline_model_parallel_size=2,
+        expert_model_parallel_size=4,
+    )
+
+
+def _rollout_batch(*, mask, removed_by_filter=None):
+    response_length = len(mask)
+    batch = {
+        "rewards": [0.25],
+        "response_lengths": [response_length],
+        "total_lengths": [response_length + 2],
+        "loss_masks": [torch.tensor(mask, dtype=torch.int32)],
+        "tokens": [torch.arange(response_length + 2)],
+        "truncated": [0],
+    }
+    if removed_by_filter is not None:
+        batch["removed_by_filter"] = [removed_by_filter]
+    return batch
+
+
+def test_validator_allows_only_explicitly_filtered_zero_mask(monkeypatch):
+    monkeypatch.setattr(actor, "get_parallel_state", lambda: None)
+    logger = MagicMock()
+
+    actor.validate_rollout_for_grpo_training_step(
+        _validator_args(),
+        _rollout_batch(mask=[0, 0, 0], removed_by_filter=True),
+        logger=logger,
+    )
+
+    warnings = "\n".join(call.args[0] for call in logger.warning.call_args_list)
+    assert "removed_by_filter=True" in warnings
+
+
+def test_validator_keeps_backward_compatibility_for_active_mask_without_filter_field(monkeypatch):
+    monkeypatch.setattr(actor, "get_parallel_state", lambda: None)
+
+    actor.validate_rollout_for_grpo_training_step(
+        _validator_args(),
+        _rollout_batch(mask=[1, 1, 1]),
+        logger=MagicMock(),
+    )
+
+
+@pytest.mark.parametrize(
+    "removed_by_filter",
+    [None, False],
+    ids=["missing-provenance", "explicitly-not-filtered"],
+)
+def test_validator_rejects_unexplained_zero_mask(monkeypatch, removed_by_filter):
+    monkeypatch.setattr(actor, "get_parallel_state", lambda: None)
+    logger = MagicMock()
+
+    with pytest.raises(ValueError, match="rollout validation failed"):
+        actor.validate_rollout_for_grpo_training_step(
+            _validator_args(),
+            _rollout_batch(mask=[0, 0, 0], removed_by_filter=removed_by_filter),
+            logger=logger,
+        )
+
+    errors = "\n".join(call.args[0] for call in logger.error.call_args_list)
+    assert "without removed_by_filter=True" in errors
+
+
+@pytest.mark.parametrize("mask", [[1, 1, 1], [1, -1, 0]], ids=["active", "zero-sum-nonzero"])
+def test_validator_rejects_filtered_sample_with_nonzero_mask(monkeypatch, mask):
+    monkeypatch.setattr(actor, "get_parallel_state", lambda: None)
+    logger = MagicMock()
+
+    with pytest.raises(ValueError, match="rollout validation failed"):
+        actor.validate_rollout_for_grpo_training_step(
+            _validator_args(),
+            _rollout_batch(mask=mask, removed_by_filter=True),
+            logger=logger,
+        )
+
+    errors = "\n".join(call.args[0] for call in logger.error.call_args_list)
+    assert "not all-zero despite removed_by_filter=True" in errors
+
+
+@pytest.mark.parametrize(
+    ("bad_provenance", "expected_error"),
+    [
+        ("yes", "must be list/tuple"),
+        ([], "length mismatch"),
+        ([1], "must be bool"),
+    ],
+    ids=["bad-type", "short-list", "non-bool"],
+)
+def test_validator_rejects_malformed_filter_provenance(monkeypatch, bad_provenance, expected_error):
+    monkeypatch.setattr(actor, "get_parallel_state", lambda: None)
+    logger = MagicMock()
+    batch = _rollout_batch(mask=[0, 0, 0])
+    batch["removed_by_filter"] = bad_provenance
+
+    with pytest.raises(ValueError, match="rollout validation failed"):
+        actor.validate_rollout_for_grpo_training_step(
+            _validator_args(),
+            batch,
+            logger=logger,
+        )
+
+    errors = "\n".join(call.args[0] for call in logger.error.call_args_list)
+    assert expected_error in errors
+
+
+def test_non_truncated_filter_provenance_reaches_each_dp_shard(monkeypatch):
+    manager_class = rollout.RolloutManager.__ray_metadata__.modified_class
+    manager = manager_class.__new__(manager_class)
+    manager.custom_convert_samples_to_train_data_func = None
+    manager.custom_reward_post_process_func = None
+    manager.args = SimpleNamespace(
+        reward_key=None,
+        advantage_estimator="ppo",
+        rewards_normalization=False,
+        use_dynamic_global_batch_size=False,
+        balance_by_flops=False,
+        balance_data=False,
+    )
+
+    filtered = Sample(
+        index=0,
+        tokens=[10, 11, 12, 13],
+        response_length=3,
+        reward=0.0,
+        status=Sample.Status.COMPLETED,
+        remove_sample=True,
+    )
+    kept = Sample(
+        index=1,
+        tokens=[20, 21, 22, 23],
+        response_length=3,
+        reward=1.0,
+        status=Sample.Status.COMPLETED,
+    )
+
+    train_data = manager._convert_samples_to_train_data([filtered, kept])
+    assert train_data["truncated"] == [0, 0]
+    assert train_data["removed_by_filter"] == [True, False]
+    assert train_data["loss_masks"] == [[0, 0, 0], [1, 1, 1]]
+
+    payloads = []
+
+    class FakeObjectRef:
+        def __init__(self, index):
+            self.index = index
+
+        def hex(self):
+            return f"ref-{self.index}"
+
+    def fake_put(payload):
+        payloads.append(payload)
+        return FakeObjectRef(len(payloads) - 1)
+
+    monkeypatch.setattr(rollout.ray, "put", fake_put)
+    manager._split_train_data_by_dp(train_data, dp_size=2)
+
+    assert [payload["removed_by_filter"] for payload in payloads] == [[True], [False]]
+    assert [payload["loss_masks"] for payload in payloads] == [[[0, 0, 0]], [[1, 1, 1]]]