[Feature] Add end-to-end TV loss for MTP SFT

[JT-Ushio]

Motivation

The official GLM-5.2 technical blog explicitly states that GLM-5.2 uses end-to-end TV loss training together with rejection sampling for its recurrent MTP design. Its coding ablation reports an acceptance length of 4.56 for the baseline and 5.47 (+20%) for the final configuration that includes end-to-end TV loss.

The blog cites Bebop, whose Eq. 13 directly optimizes normalized multi-step rejection-sampling acceptance:

tv_alpha_i = sum_v min(p_i(v), q_i(v))
L_e2e = 1 - (1 / gamma) * sum_j prod_{i=1}^j tv_alpha_i

In the current XTuner SFT path, MTP only constructs independent per-depth cross-entropy losses (or KL when rollout log-probabilities are supplied). There is no selectable joint end-to-end TV objective, so the GLM-5.2 training objective cannot currently be reproduced.

Changes

• Add MTPConfig.loss_type with ce and e2e_tv options; ce remains the default for backward compatibility.
• Implement the exact full-vocabulary Eq. 13 objective with a stop-gradient target distribution.
• Compute all recurrent MTP depths jointly, including fixed-horizon masking for packed sequences and existing token/sample/global-batch loss calibration.
• Support sequence-parallel target-state halos without gathering the full sequence in the common case.
• Bound peak activation memory through token chunking and non-reentrant checkpointing; no top-k TV approximation is used.
• Integrate the joint objective into both regular and intra-layer-microbatch MoE paths while preserving MTP router auxiliary losses.
• Expose opt-in settings in examples/v1/config/sft_glm5p2.py.
• Add numerical and gradient-flow unit tests.

Example:

MTP_LOSS_TYPE=e2e_tv \
MTP_LOSS_WEIGHT=0.1 \
MTP_TV_LOSS_CHUNK_SIZE=128 \
...

Scope and compatibility

• Existing CE behavior is unchanged by default.
• This PR intentionally does not change the physical-vs-recurrent MTP depth mapping or shared-weight configuration; GLM-5.2's seven recurrent shared steps are handled separately.
• The target distribution p is detached as required by the TV objective. Existing detach_mtp_inputs and detach_mtp_lm_head_weight switches continue to control whether the draft branch reaches backbone inputs and the shared LM head.
• The GLM-5.2 blog does not publish the loss coefficient or exact trainable-parameter boundary. The example value 0.1 preserves XTuner's existing MTP loss scale and is not claimed to be an official GLM-5.2 value.

Validation

• Ruff lint and format checks for all changed files
• git diff --check
• Exact Eq. 13 numerical comparison on packed samples
• Gradient tests for detached teacher targets, trainable draft states, and optional detached shared LM-head weights

Full CUDA/Triton integration testing is left to CI.

[Copilot]

🔵 Needs a closer look

It introduces a new distributed, full-vocabulary loss path with nontrivial masking/SP behavior that warrants final human validation in real training runs beyond unit tests.

Pull request overview

This PR adds an opt-in end-to-end total-variation (TV) loss for Multi-Token Prediction (MTP) in XTuner SFT, enabling joint optimization of multi-step rejection-sampling acceptance (per Eq. 13 in Bebop / referenced by the GLM-5.2 blog) while preserving existing per-depth CE behavior by default.

Changes:

• Add MTPConfig.loss_type ("ce" default, "e2e_tv" opt-in) and tv_loss_chunk_size to control the new objective and its memory/perf tradeoff.
• Implement the exact full-vocabulary e2e TV objective as a new MTPE2ETVLossConfig/Context, including packed-sequence masking and SP-aware target shifting.
• Integrate the joint loss path into MoE training (regular + intra-layer microbatch) and add focused unit tests plus an example config toggle via env vars.

File summaries

File	Description
xtuner/v1/module/mtp/config.py	Adds config knobs for selecting CE vs e2e TV loss and chunk sizing.
xtuner/v1/module/lm_head/lm_head.py	Extends LMHead typing to accept tuple-shaped hidden states for the joint TV loss context.
xtuner/v1/model/moe/moe.py	Wires loss_type into loss-context building and computes either per-depth CE or joint e2e TV MTP loss.
xtuner/v1/loss/mtp_loss.py	Implements the end-to-end TV loss config/context, including SP halo shifting and token-chunked exact TV overlap.
xtuner/v1/loss/init.py	Exports MTPE2ETVLossContext.
tests/loss/test_mtp_e2e_tv_loss.py	Adds numerical and gradient-flow tests for the e2e TV loss (teacher detachment + optional head detachment).
examples/v1/config/sft_glm5p2.py	Exposes MTP_LOSS_TYPE, MTP_LOSS_WEIGHT, and MTP_TV_LOSS_CHUNK_SIZE environment toggles.

Review details

• Files reviewed: 7/7 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.