[Feat] Add Dense MTP and Qwen3.5 VL Dense 4B/27B configs

[PengchengShi00]

Summary

• Enable Multi-Token Prediction on dense models: mtp_config lives on TransformerConfig, Dense builds/forwards an MTP block, and TrainEngine can pick up mtp_loss from ModelOutputs.
• Make MTP layer outputs model-agnostic (MTPDepthOutput): hidden_states is required; router fields are optional so dense MTP no longer pretends to be a MoE decoder.
• Add Qwen3.5 VL Dense 4B / 27B compose + text configs, HF mtp_block.* → mtp.* mapping, and update model.md.

Test plan

• Default path: mtp_config=None dense / MoE training still works (no MTP modules, no extra loss).
• Dense MTP: set text_config.mtp_config = MTPConfig(...), confirm mtp_loss is logged and included in backward.
• Qwen3.5 VL Dense 4B/27B: from_hf / save_hf — mtp.* keys load when MTP is on; unused mtp.* keys stay unexpected when MTP is off.
• MoE MTP: router logits/weights/topk still flow into aux loss (asserts pass).

Usage

model_cfg = Qwen3_5_VLDense27BConfig(freeze_vision=True, freeze_projector=True)
model_cfg.text_config.mtp_config = MTPConfig(
    num_layers=mtp_num_layers,
    share_weights=True,
    detach_mtp_lm_head_weight=True,
    detach_mtp_inputs=True,
    loss_scaling_factor=1.0,

[YanhuiDua]

@claude review

[github-actions]

Claude:

Summary

本 PR 把 MTP 从 MoE 下沉到 TransformerConfig，让 Dense 也能构建/前向 MTP block，并把 MTP 层输出从 MoEDecoderLayerOutput 别名换成独立的 MTPDepthOutput（router 字段 NotRequired），同时新增 Qwen3.5 VL Dense 4B/27B 配置与 mtp_block.* → mtp.* 的 HF key 映射。抽象方向是正确的：dense MTP 不再假装自己是 MoE decoder，MTPBlock/MTPLayer 的 Interface 变小而承载的行为更多。发现 1 个 Critical（仓库内现有 import 被破坏）和 2 个 Warning。

Main Flowchart after this PR

flowchart TD
    A["Dense.__init__"] -->|mtp_config 非 None| B["build_mtp_block()"]
    C["Dense.forward"] --> D["decoder layers"]
    D --> E["norm + lm_head → loss"]
    E --> F["_maybe_forward_mtp()"]
    F --> G["MTPBlock → list[MTPDepthOutput]"]
    G --> H["output['mtp_loss'] = mean(depth losses) × loss_scaling_factor"]

    I["Dense.fully_shard"] --> J["decoder 分支：layer_idx < int(num_hidden_layers × recompute_ratio)"]
    I --> K["MTP 分支：_should_recompute(mtp_idx=...)"]
    J -. 同一规则两套口径 .- K

    classDef changed fill:#d5f5d5,stroke:#2e7d32
    classDef problem fill:#ffd6d6,stroke:#c62828
    class B,F,G,H,K changed
    class J problem

Loading

核心原理实现与单测

• Dense MTP 的核心路径（build_mtp_block → build_loss_ctx_batch 追加 mtp 上下文 → _maybe_forward_mtp → mtp_loss）由 TestQwen3_8MTPSequenceParallel::test_mtp_loss_and_gradients_match_full_sequence 通过 public API（build_loss_ctx_batch + model(...) + backward）真实覆盖，并对齐 SP2 与全序列基线的 loss 与参数梯度，未 mock 项目内模块。
• MTPDepthOutput 改造后的 HF key 映射由 test_mtp_block_and_hf_key_mapping 在 meta device 上经 to_hf_key_list 覆盖，并验证 mtp_config=None 时 mtp_block is None。
• MoE 侧把原本无条件读取的 router 字段改为 assert "router_logits" in ... 后再读，语义与既有 MoE MTP 测试保持一致；init_weight.py 新增的 dt_bias/A_log 初始化被上述 SP 测试经 init_weights() 的 missing-param 校验间接覆盖，uniform_(0.01, 16) 相比 HF 的 uniform_(0, 16) 还规避了 log(0)。
• 未覆盖的核心行为见下方「单测建议」。

抽象与信息隐藏评估

• Warning — xtuner/v1/model/dense/dense.py:361-370：新增的 _should_recompute 只在 MTP 分支被调用，主 decoder 循环仍用旧的 int(num_hidden_layers * recompute_ratio)，同一条 recompute 规则散落在两处且与该方法 docstring 宣称的「decoder + MTP 视为同一序列」不符（MoE 侧两个分支都走 _should_recompute）。

单测建议

• Warning — xtuner/v1/model/dense/dense.py:394-425：Dense 的 MTP FSDP 分片与 recompute 决策（含 share_weights 强制 checkpoint、prefetch 链接）没有任何测试覆盖，上面的规则不一致因此无法被回归发现。

其他 Issues

• Critical — xtuner/v1/model/moe/moe.py:76：MTPConfig 不再从 moe.py 导入，但仓库内 ci/config/qwen3_5_moe_35BA3.py:14、examples/v1/config/sft_qwen3p5.py:9、tests/model/test_qwen3_5.py:15 仍在 from xtuner.v1.model.moe.moe import MTPConfig，会直接 ImportError。

Verdict

REQUEST_CHANGES

[claude]

Claude: [兼容性] 这里移除了 MTPConfig 的导入，但仓库内 ci/config/qwen3_5_moe_35BA3.py:14、examples/v1/config/sft_qwen3p5.py:9、tests/model/test_qwen3_5.py:15 仍在 from xtuner.v1.model.moe.moe import MTPConfig，会直接 ImportError。请同步改为从 xtuner.v1.module.mtp 导入，或在 moe.py 保留 re-export。

[claude]

Claude: [设计] 新增的 _should_recompute 目前只被 MTP 分支调用（本文件 400 行），主 decoder 循环仍沿用旧口径 num_recompute_layers = int(self.config.num_hidden_layers * self.fsdp_config.recompute_ratio)（361 行）配合 layer_idx < num_recompute_layers（370 行）判断。

Problem：同一条 recompute 规则在 Dense.fully_shard 中存在两套口径，与这里 docstring 宣称的「decoder 层与 MTP 层视为同一序列，recompute_ratio 作用于 num_hidden_layers + mtp_layers」不符。开启 MTP 后两者分母不同，decoder 段实际重算层数会多于本方法定义的全局规则，decoder 段也永远不享受「序列最后一层不 checkpoint」的规则。对照 MoE 侧，decoder 循环与 MTP 循环都统一走 self._should_recompute(...)（xtuner/v1/model/moe/moe.py:1324 与 :1380），Dense 这里属于迁移不完整。

Solution：让 decoder 循环也经由同一个 Seam 决策，并删除局部的 num_recompute_layers：

if self._should_recompute(layer_idx=layer_idx, mtp_idx=None):
    layer = apply_activation_checkpointing(...)

Benefits：recompute 规则收敛到 _should_recompute 一处（Locality），Dense 与 MoE 行为一致，后续调整 ratio 语义只需改一处实现，该规则也可通过单一入口被测试。

[claude]

Claude: [测试] Dense 的 MTP 分片路径（share_weights 强制 checkpoint、reshard_after_forward、prefetch 链接）无任何测试覆盖，上面 recompute 规则的双口径问题因此无法被回归发现。建议补一个 FSDP 用例断言各 MTP layer 的 checkpoint 与分片决策。