Add MoVA training, weight sync, and rollout validation

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
 ```
 
-

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):

miles/backends/megatron_utils/actor.py

@@ -28,6 +28,7 @@
 
 from ...utils.profile_utils import TrainProfiler
 from ...utils.tensor_backper import TensorBackuper
+from ..training_utils.ci_utils import assert_rollout_engine_weight_versions
 from ..training_utils.cp_utils import slice_with_cp
 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
@@ -190,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"
@@ -227,6 +235,7 @@ def _basic_batch_summary():
             "rewards",
             "response_lengths",
             "total_lengths",
+            "removed_by_filter",
             "loss_masks",
             "tokens",
             "input_ids",
@@ -241,6 +250,13 @@ 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"]]
@@ -328,6 +344,24 @@ 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"
@@ -410,10 +444,24 @@ def _basic_batch_summary():
             continue
 
         mask_sum = _sum_float(mask)
-        if mask_sum <= 0:
-            # Warning-only: rollout sample filters zero a sample's loss_mask
-            # (remove_sample) to drop it from the gradient; the reducer clamps.
-            _add_warning(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(
@@ -485,15 +533,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.
@@ -1016,13 +1055,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:

miles/backends/megatron_utils/megatron_to_hf/xllm.py

@@ -3,20 +3,135 @@
 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":
         return [("lm_head.weight", 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 +168,54 @@ 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 +225,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)]