[Feature] Freeze native MTP parameters while preserving load, save and rollout publication

docs/en/rl/advanced_tutorial/index.rst

@@ -9,3 +9,4 @@ Reinforcement Learning
    rl_trainer.md
    loss.md
    on_policy_distillation.md
+   native_frozen_mtp.md

docs/en/rl/advanced_tutorial/native_frozen_mtp.md

@@ -0,0 +1,13 @@
+# Native frozen MTP (experimental)
+
+Use `MTPConfig(num_layers=1, frozen=True)` to retain dedicated MTP parameters for native loading, saving and rollout weight publication while excluding them from actor training. The main model and its shared embedding/output head remain trainable. Frozen MTP does not execute its training forward or produce MTP training losses.
+
+This is an alternative to the external `WorkerConfig.frozen_mtp` adapter proposed in XTuner #2289. Do not enable both routes. This standalone change does not depend on the Agent lifecycle changes in #2257.
+
+For self-contained DCP checkpoints, explicitly set the outer model configuration `dcp_ignore_frozen_params=False`; its existing default can omit frozen parameters. Use the native training engine/Trainer saving and loading paths. An HF export or a worker model/optimizer checkpoint alone does not validate full Trainer cold restoration.
+
+Qwen3.5 source checkpoints may contain fused expert tensors under `mtp.layers.*.mlp.experts.gate_up_proj` and `down_proj`. The loader resolves these into native per-expert matrices for unquantized BF16, FP16 and FP32 safetensors. Direct indexed keys retain priority. Other key families and bin checkpoints retain their previous behavior. Invalid fused shapes, expert indices and unsupported dtypes fail explicitly. Quantized fused-source aliases are not supported.
+
+Inference speculation still requires the existing backend configuration: LMDeploy PyTorch `lmdeploy_speculative_algorithm="qwen3_5_mtp"` and a positive `lmdeploy_speculative_num_draft_tokens`. Freezing alone does not enable speculative decoding. Backend/topology support follows native weight publication; this change does not add a new transport.
+
+The feature is experimental. Component GPU evidence and current integration validation are described in the PR, with exact source scopes. Short smoke tests do not establish learning quality, speedup or general distributed fault recovery.

tests/model/test_glm53_shared_mtp_depth.py

@@ -29,6 +29,85 @@ def test_glm53_builder_and_hf_mapping_keep_one_physical_layer(depth):
     ]
 
 
+class TestFrozenMTP:
+    @pytest.mark.parametrize("frozen", [False, True])
+    def test_builder_retains_parameters_and_load_specs(self, frozen):
+        cfg = Glm53TextMoEConfig(
+            compile_cfg=False,
+            ep_size=1,
+            dispatcher=None,
+            mtp_config=MTPConfig(num_layers=1, share_weights=True, frozen=frozen),
+        )
+        with patch("torch.cuda.Stream"), torch.device("meta"):
+            model = cfg.build()
+
+        mtp_parameters = dict(model.mtp_block.named_parameters())
+        assert mtp_parameters
+        if frozen:
+            assert not any(parameter.requires_grad for parameter in mtp_parameters.values())
+        else:
+            assert any(parameter.requires_grad for parameter in mtp_parameters.values())
+        assert model._mtp_training_enabled is not frozen
+        assert any(name.startswith("mtp_block.") for name, _ in model.trainable_parameters()) is not frozen
+        assert any(spec.name.startswith("mtp_block.") for _, spec in model._load_spec_params())
+
+    def test_micro_batch_forward_skips_frozen_mtp(self):
+        class LossContext:
+            @classmethod
+            def cat(cls, contexts):
+                return cls()
+
+        cfg = Glm53TextMoEConfig(
+            compile_cfg=False,
+            ep_size=1,
+            dispatcher=None,
+            mtp_config=MTPConfig(num_layers=1, share_weights=True, frozen=True),
+        )
+        with patch("torch.cuda.Stream"), torch.device("meta"):
+            model = cfg.build()
+        seq_ctx = SequenceContext.from_input_ids((torch.arange(4).view(1, -1),), device="cpu")
+        loss_ctx = {
+            "lm": LossContext(),
+            "balancing": None,
+            "z_loss": None,
+            "mtp": None,
+            "mtp_e2e_tv": None,
+        }
+
+        with (
+            patch.object(model.embed_tokens, "forward", lambda ids: ids.float().unsqueeze(-1)),
+            patch.object(
+                model.rotary_emb,
+                "forward",
+                lambda hidden, positions: (torch.zeros_like(hidden), torch.zeros_like(hidden)),
+            ),
+            patch.object(model, "_micro_batch_decoder_stack", side_effect=lambda **kwargs: kwargs["hidden_states_list"]),
+            patch.object(model, "norm", nn.Identity()),
+            patch.object(model.lm_head, "forward", return_value=(torch.tensor(1.0), (None, None))),
+            patch.object(model.aux_loss, "finalize", return_value=(None, None, torch.zeros(1))),
+            patch.object(model.mtp_block, "forward") as mtp_forward,
+        ):
+            output = model._micro_batch_forward([seq_ctx, seq_ctx], [loss_ctx, loss_ctx])
+
+        mtp_forward.assert_not_called()
+        assert output.mtp_loss is None
+
+    def test_bias_update_consumes_only_main_router_rows(self):
+        cfg = Glm53TextMoEConfig(
+            compile_cfg=False,
+            ep_size=1,
+            dispatcher=None,
+            mtp_config=MTPConfig(num_layers=1, share_weights=True, frozen=True),
+        )
+        with patch("torch.cuda.Stream"), torch.device("meta"):
+            model = cfg.build()
+
+        num_main_routers = cfg.num_hidden_layers - cfg.first_k_dense_replace
+        expert_counts = torch.zeros(num_main_routers, cfg.n_routed_experts)
+        with patch.object(torch.Tensor, "add_", lambda tensor, other: tensor):
+            model.update_bias(expert_counts, expert_counts.float().mean(dim=1))
+
+
 class _RecurrentLayer(nn.Module):
     def __init__(self):
         super().__init__()

tests/rl/test_weight_iterator.py

@@ -96,6 +96,34 @@ def _single_rank_rollout_info(*, backend: RolloutBackend) -> RolloutWeightUpdate
     )
 
 
+class FrozenMTPModel(MixedDtypeModel):
+    def __init__(self) -> None:
+        super().__init__()
+        self.mtp_block = nn.Module()
+        self.mtp_block.weight = nn.Parameter(
+            torch.ones(2, device=get_device(), dtype=torch.bfloat16), requires_grad=False
+        )
+        self._init_load_spec()
+
+
+class TestFrozenMTPWeightUpdate:
+    def test_frozen_mtp_parameters_remain_in_native_batches(self) -> None:
+        model = FrozenMTPModel()
+        iterator = WeightIterator(
+            config=SimpleNamespace(update_weight_bucket_size_in_gb=1, model_cfg=None),
+            engine=SimpleNamespace(model=model),
+            rollout_info=_single_rank_rollout_info(backend="pytorch"),
+            global_hf_keys_mapping_cache={},
+        )
+
+        batches = list(iterator.iter_hf_batches(final_update=True))
+        state_dict = {name: tensor for batch in batches for name, tensor in batch.state_dict.items()}
+
+        assert not model.mtp_block.weight.requires_grad
+        assert set(state_dict) == {"bf16_weight", "fp32_weight", "mtp_block.weight"}
+        assert batches[-1].finished
+
+
 def test_hf_weight_update_batches_have_one_dtype() -> None:
     model = MixedDtypeModel()
     rollout_info = _single_rank_rollout_info(backend="pytorch")

tests/utils/test_fused_mtp_loader.py

@@ -0,0 +1,77 @@
+"""Actual safetensors I/O and native loader; no allocator or model stubs."""
+
+import json
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from xtuner.v1.utils.loader import HFCheckpointLoader
+
+
+PREFIX = "mtp.layers.0.mlp.experts"
+
+
+def checkpoint(path, tensors, indexed=False):
+    save_file(tensors, str(path / "model.safetensors"))
+    if indexed:
+        (path / "model.safetensors.index.json").write_text(
+            json.dumps({"weight_map": dict.fromkeys(tensors, "model.safetensors")})
+        )
+    return HFCheckpointLoader(str(path))
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("indexed", [False, True])
+def test_fused_mtp_expert_reads(tmp_path, dtype, indexed):
+    gate_up = torch.arange(24).reshape(2, 4, 3).to(dtype)
+    down = torch.arange(12).reshape(2, 3, 2).to(dtype)
+    loader = checkpoint(tmp_path, {PREFIX + ".gate_up_proj": gate_up, PREFIX + ".down_proj": down}, indexed)
+    for expert in range(2):
+        for part, expected in [("gate", gate_up[expert, :2]), ("up", gate_up[expert, 2:]), ("down", down[expert])]:
+            key = f"{PREFIX}.{expert}.{part}_proj.weight"
+            assert loader.is_key_exist(key)
+            torch.testing.assert_close(loader.load(key), expected, rtol=0, atol=0)
+    assert not loader.is_key_exist(f"{PREFIX}.0.gate_proj.weight_scale_inv")
+
+
+def test_direct_export_keys_take_precedence(tmp_path):
+    direct = f"{PREFIX}.0.gate_proj.weight"
+    loader = checkpoint(tmp_path, {PREFIX + ".gate_up_proj": torch.zeros(2, 4, 3), direct: torch.ones(2, 3)})
+    torch.testing.assert_close(loader.load(direct), torch.ones(2, 3), rtol=0, atol=0)
+
+
+def test_missing_and_non_mtp_keys_preserve_behavior(tmp_path):
+    loader = checkpoint(tmp_path, {PREFIX + ".gate_up_proj": torch.zeros(2, 4, 3)})
+    for key in [
+        "model.layers.0.mlp.experts.0.gate_proj.weight",
+        f"{PREFIX}.0.unknown_proj.weight",
+        "mtp.layers.1.mlp.experts.0.gate_proj.weight",
+    ]:
+        assert not loader.is_key_exist(key)
+        assert loader.load(key) is None
+
+
+def test_bin_checkpoint_direct_path_preserved(tmp_path):
+    value = torch.arange(6).reshape(2, 3)
+    torch.save({"ordinary.weight": value}, tmp_path / "part.bin")
+    (tmp_path / "model.bin.index.json").write_text(json.dumps({"weight_map": {"ordinary.weight": "part.bin"}}))
+    loader = HFCheckpointLoader(str(tmp_path))
+    assert loader.is_key_exist("ordinary.weight")
+    torch.testing.assert_close(loader.load("ordinary.weight"), value)
+    assert loader.load(f"{PREFIX}.0.gate_proj.weight") is None
+
+
+@pytest.mark.parametrize(
+    "tensor,expert,message",
+    [
+        (torch.zeros(2, 4), 0, "shape/index"),
+        (torch.zeros(2, 4, 3), 2, "shape/index"),
+        (torch.zeros(2, 5, 3), 0, "even"),
+        (torch.zeros(2, 4, 3, dtype=torch.int64), 0, "unquantized"),
+    ],
+)
+def test_invalid_fused_inputs_fail_explicitly(tmp_path, tensor, expert, message):
+    loader = checkpoint(tmp_path, {PREFIX + ".gate_up_proj": tensor})
+    with pytest.raises(ValueError, match=message):
+        loader.load(f"{PREFIX}.{expert}.gate_proj.weight")

xtuner/v1/model/moe/moe.py

@@ -273,6 +273,8 @@ def __init__(self, config: MoEConfig):
         self.rotary_emb = self.build_rotary_embedding(config)
         self.embed_tokens = self.build_embeddings(config)
         self.mtp_block = self.build_mtp_block(config) if config.mtp_config is not None else None
+        if self.mtp_block is not None and not self._mtp_training_enabled:
+            self.mtp_block.requires_grad_(False)
         self._configure_model_specific_layers()
         # Per-dispatch grouped-GEMM row counts of the current train step; None outside one.
         self._ep_recv_tokens: list[int] | None = None
@@ -306,6 +308,10 @@ def _maybe_offload_router(self, tensor: torch.Tensor) -> torch.Tensor:
     def _configure_model_specific_layers(self) -> None:
         return
 
+    @property
+    def _mtp_training_enabled(self) -> bool:
+        return self.mtp_block is not None and self.config.mtp_config is not None and not self.config.mtp_config.frozen
+
     def _saved_tensors_offload_ctx(
         self,
         block_idx: int,
@@ -413,7 +419,8 @@ def update_bias(self, total_expert_counts_pre_iter, expected_loads):
             )
             stats_idx += 1
 
-        if self.mtp_block is not None and self.config.mtp_config is not None:
+        if self._mtp_training_enabled:
+            assert self.mtp_block is not None and self.config.mtp_config is not None
             mtp_config = self.config.mtp_config
             if mtp_config.share_weights:
                 end = stats_idx + mtp_config.num_layers
@@ -484,17 +491,18 @@ def build_loss_ctx_batch(  # type: ignore[override]
         self._add_auxiliary_loss("balancing", self.config.balancing_loss_cfg, _data_batch, res)
         self._add_auxiliary_loss("z_loss", self.config.z_loss_cfg, _data_batch, res)
 
-        # Add MTP loss contexts if MTP is enabled.
-        if self.config.mtp_config is not None:
+        # Add MTP loss contexts if MTP training is enabled.
+        if self._mtp_training_enabled:
             mtp_config = self.config.mtp_config
+            assert mtp_config is not None
             if mtp_config.loss_type == "e2e_tv":
                 tv_loss_cfg = MTPE2ETVLossConfig(
                     ignore_idx=self.config.lm_loss_cfg.ignore_idx,
                     mode="chunk",
                     chunk_size=mtp_config.tv_loss_chunk_size,
                     loss_reduction=self.config.lm_loss_cfg.loss_reduction,
                     num_steps=mtp_config.num_layers,
-                    detach_mtp_lm_head_weight=self.config.mtp_config.detach_mtp_lm_head_weight,
+                    detach_mtp_lm_head_weight=mtp_config.detach_mtp_lm_head_weight,
                 )
                 tv_loss_ctx_list = self._build_loss_ctx(tv_loss_cfg, _data_batch, sp_mesh)
                 if tv_loss_ctx_list is not None:
@@ -739,8 +747,8 @@ def _micro_batch_forward(
 
         assert hidden_states_list, "XTuner Internal Error, found empty hidden states for domino EP"
 
-        if self.mtp_block is not None:
-            assert self.config.mtp_config is not None
+        if self._mtp_training_enabled:
+            assert self.mtp_block is not None and self.config.mtp_config is not None
 
             # Build a per-microbatch SequenceContext clone for MTP. We always run the MTP
             # block on every micro-batch so domino EP can overlap dispatch/combine across
@@ -1025,10 +1033,11 @@ def _forward(
         mtp_loss_ctx_list = loss_ctx.get("mtp") if loss_ctx is not None else None
         mtp_tv_loss_ctx = loss_ctx.get("mtp_e2e_tv") if loss_ctx is not None else None
         if (
-            self.mtp_block is not None
+            self._mtp_training_enabled
             and loss_ctx is not None
             and (mtp_loss_ctx_list is not None or mtp_tv_loss_ctx is not None)
         ):
+            assert self.mtp_block is not None
             mtp_seq_ctx = seq_ctx.copy(
                 input_ids=input_ids.clone() if input_ids is not None else None,
                 position_ids=position_ids.clone(),

xtuner/v1/module/mtp/config.py

@@ -23,6 +23,8 @@ class MTPConfig(BaseModel):
         share_weights (bool): Whether to share the weights of the MTP layers.
             If True, the weights of the MTP layers are shared across all layers.
             Default: False.
+        frozen (bool): Whether to retain MTP parameters for export and rollout weight updates without
+            training them or executing the MTP training forward pass. Default: False.
         detach_mtp_lm_head_weight (bool): Whether to detach the LM head weight.
             This is used in RL training. Default is False.
         detach_mtp_inputs (bool): Whether to detach the input embeddings and hidden states.
@@ -58,6 +60,7 @@ class MTPConfig(BaseModel):
 
     num_layers: Annotated[int, Parameter(group="model")]
     share_weights: Annotated[bool, Parameter(group="model")] = False
+    frozen: Annotated[bool, Parameter(group="model")] = False
     detach_mtp_lm_head_weight: Annotated[bool, Parameter(group="model")] = False
     detach_mtp_inputs: Annotated[bool, Parameter(group="model")] = False
     loss_scaling_factor: Annotated[float, Parameter(group="model")] = 0.1

xtuner/v1/rl/trainer/actor_worker.py

@@ -119,6 +119,8 @@ def __init__(self, worker_cfg: WorkerConfig, ctx: RoleContext) -> None:
                 self.mtp_config = worker_cfg.model_cfg.text_config.mtp_config
         elif hasattr(worker_cfg.model_cfg, "mtp_config"):
             self.mtp_config = worker_cfg.model_cfg.mtp_config
+        if self.mtp_config is not None and self.mtp_config.frozen:
+            self.mtp_config = None
         self.mtp_cfgs: list[MTPLossConfig] = []
         if self.mtp_config is not None:
             for depth in range(self.mtp_config.num_layers):

xtuner/v1/utils/loader.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 from typing import Literal
 
 import torch
@@ -78,11 +79,48 @@ def __init__(
         self.buffer = None
 
     def is_key_exist(self, key: str) -> bool:
-        return key in self.weight_map
+        return key in self.weight_map or self._resolve_fused_mtp_expert(key) is not None
+
+    def _resolve_fused_mtp_expert(self, key: str) -> tuple[str, int, str] | None:
+        """Resolve only MTP expert aliases in unquantized safetensors sources.
+
+        Qwen3.5 checkpoints can store [experts, 2*intermediate, hidden]
+        gate/up and [experts, hidden, intermediate] down tensors while the
+        native MTP load spec addresses per-expert matrices.
+        """
+        if not self.use_safetensors:
+            return None
+        match = re.fullmatch(r"(mtp\.layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight", key)
+        if match is None:
+            return None
+        prefix, expert, part = match.groups()
+        fused = prefix + (".down_proj" if part == "down" else ".gate_up_proj")
+        return (fused, int(expert), part) if fused in self.weight_map else None
+
+    def _load_fused_mtp_expert(self, key: str) -> torch.Tensor | None:
+        resolved = self._resolve_fused_mtp_expert(key)
+        if resolved is None:
+            return None
+        fused, expert, part = resolved
+        path = os.path.join(self.model_path, self.weight_map[fused])
+        if path not in self._safetensor_cache:
+            self._safetensor_cache[path] = safe_open(path, framework="pt")
+        view = self._safetensor_cache[path].get_slice(fused)
+        shape = view.get_shape()
+        if len(shape) != 3 or expert >= shape[0]:
+            raise ValueError(f"Invalid fused MTP expert shape/index for {key}: {shape}")
+        if view.get_dtype() not in {"BF16", "F16", "F32"}:
+            raise ValueError(f"Fused MTP aliases require unquantized floating weights: {fused}")
+        if part == "down":
+            return view[expert]
+        if shape[1] % 2:
+            raise ValueError(f"Fused MTP gate/up dimension must be even: {fused}")
+        half = shape[1] // 2
+        return view[expert, :half] if part == "gate" else view[expert, half:]
 
     def load(self, key) -> torch.Tensor | None:
         if key not in self.weight_map:
-            return None
+            return self._load_fused_mtp_expert(key)
 
         _file = self.weight_map[key]