Add TV (total-variation) loss option for the MTP draft head

[ElliotXinqiWang]

PR: Add TV (total-variation) loss option for the MTP draft head

Repo: radixark/Megatron-LM Base: miles-main @ 36ceb4fc0 Branch: feat/mtp-tv-loss
Files: megatron/core/models/gpt/gpt_model.py (+106), megatron/core/transformer/transformer_config.py (+13)

Motivation

The MTP draft head is trained with cross-entropy. But for speculative decoding the
rejection-sampling acceptance rate is accept_rate = 1 - d_TV(p_main, q_draft), so
total variation is the direct objective; CE only bounds TV via Pinsker. Optimizing TV
directly heals the draft head faster (esp. when the main model has drifted from the head's
training distribution — e.g. an official MTP head grafted onto an SFT'd / RL'd main model).

Ref: "Breaking Entropy Bounds: Accelerating RL Training via MTP with Rejection Sampling" (Bebop).

What it adds

• TransformerConfig.mtp_loss_type ('ce' default | 'tv') and mtp_tv_chunk: int = 0.
  Auto-exposed as --mtp-loss-type / --mtp-tv-chunk (ArgumentGroupFactory) and mapped back
  via core_transformer_config_from_args — no manual arg wiring.
• GPTModel: when mtp_loss_type == 'tv', the MTP-loss branch computes per token
  1 - Σ_v min(softmax(p_main).detach(), softmax(q_draft)). p_main is the main model's
  next-token distribution (fully detached; main hidden states are rolled to align with the
  MTP target, mirroring the existing label shift). Produces mtp_loss in the same [b, s]
  layout as the CE branch, so the existing MTPLossAutoScaler / scaling / logging path is
  unchanged — it's a drop-in alternative to the CE branch.
• mtp_tv_chunk > 0 chunks the vocab-sum over the sequence dim and gradient-checkpoints
  each chunk (recomputes softmax in backward), so backward memory is ~one chunk instead of
  O(T·V) — this is what lets TV fit alongside full-backbone (joint) training.

Limitation

TV needs the full vocabulary per rank for Σ_v min(p,q), so it requires
tensor-model-parallel size 1 (raises NotImplementedError for TP>1). Vocab-parallel TV
(all-reduce of the min-overlap) is possible future work.

Validation

Trained the Qwen3.5-35B-A3B grafted MTP head with --mtp-loss-type tv --mtp-tv-chunk 512
(freeze + joint), 300 steps each, on the miles RL stack; mtp_loss (= mean TV) ~0.13–0.15,
no OOM with chunking, spec accept-rate tracked as expected. CE path unchanged (default).

Notes for reviewers

• Default behavior is identical (mtp_loss_type='ce').
• --mtp-tv-chunk only affects the TV path; ignored for CE.
• Companion (separate) change in miles: none required — it forwards Megatron args; an optional
  reset_arg could surface --mtp-loss-type explicitly.