feat(speechlm2): distill MTP heads with LK loss

[Slyne]

Summary

• add opt-in MTP head training to SALMAutomodel, including strict head_only freezing and recurrent/repeated MTP depth support
• add an LK distillation objective that combines forward KL and total variation against the frozen backbone distribution
• preserve packed-sequence boundaries and exchange only the required future teacher rows under context parallelism
• bound vocabulary-projection memory and keep projection/checkpoint collectives rank-symmetric for uneven or empty supervision
• add per-depth LK metrics, validation agreement metrics, configuration documentation, and focused distributed tests

Configuration

model:
  mtp:
    enabled: true
    training_mode: head_only
    loss_type: lk
    lk_lambda: 0.5
    num_nextn_predict_layers: 2
    use_repeated_layer: true

loss_type: lk requires training_mode: head_only, keeping the backbone deterministic and frozen as the teacher. Cross-entropy MTP remains the backward-compatible default.

Verification

• all 51 focused SpeechLM2 MTP tests passed, including two-rank Gloo FSDP uneven-row and CP-reference cases
• Black formatting check and git diff --check passed
• end-to-end Hero5c CP=1/EP=8 head-only LK training completed five 12,288-token optimizer steps on eight H100s (Slurm 6351147)

A final CP=2/EP=4 CUDA smoke is intentionally retained as a draft-PR gate because it exercises the NCCL/TE context-parallel teacher exchange.

Related

• Automodel empty-EP-rank autograd fix: fix(moe): preserve DeepEP autograd for empty expert ranks Automodel#3680

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.