feat(speechlm2): distill MTP heads with LK loss

[Slyne]

What changed

• add an opt-in mtp.loss_type: lk objective for head-only MTP fine-tuning
• use the frozen backbone distribution as the teacher and optimize
  lambda * KL(p || q) + (1 - lambda) * TV(p, q)
• align each teacher distribution with the corresponding future MTP depth while preserving padding and packed-sequence masks
• keep the frozen perception/backbone/LM head in eval mode while the MTP head remains trainable
• support context parallelism by exchanging only requested teacher-logit rows, avoiding a full sequence-logit all-gather
• preserve dataset-label cross-entropy as the default MTP objective

This implements the fixed-mixture LK objective described in LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding.

Why

When the speech encoder and backbone LLM are frozen and only the MTP module is trained as a speculative-decoding draft model, one-hot dataset labels are not the best target. Matching the verifier distribution directly trains the draft head toward higher speculative acceptance.

Configuration

mtp:
  enabled: true
  training_mode: head_only
  loss_type: lk
  lk_lambda: 0.5

lk_lambda is effective only when training_mode: head_only and loss_type: lk; LK loss is rejected for other training modes. The default cross-entropy path is unchanged.

Dependency

#16071 has merged and main now pins Automodel to 7b15c9bf. This PR uses that pin and contains no Automodel source changes.

Validation

• all GitHub checks passing, including lint, formatting, install, ASR, DCO, and CodeQL
• black --check and isort --check-only on modified Python files
• python -m py_compile on modified implementation files
• 73 passing focused tests:
  • tests/collections/speechlm2/test_salm_automodel_mtp.py
  • tests/collections/speechlm2/test_salm_packed_sequences.py
• real two-rank Gloo context-parallel regression covering TE-style int32 partition indices
• real two-rank FSDP2 regression with uneven/all-ignore labels, validating rank-symmetric frozen LM-head projection calls
• IAD GPU smoke job 12262461 completed 0:0 in 5m35s on one 8xA100 node with CP=2, EP=4, TP=1, packed THD, TE attention, head-only LK, and lk_lambda=0.5:
  • one forward/backward/optimizer step completed with finite loss=5.780 and mtp_loss=0.962
  • one validation batch completed
  • complete step=1-last.ckpt DCP checkpoint with all eight rank shards plus metadata
  • no traceback, non-finite value, distributed hang, or NCCL failure

The IAD smoke used TORCH_NCCL_USE_COMM_NONBLOCKING=0, the established workaround for the cluster's Mamba context-parallel collective issue. No checked-in configuration was changed for the smoke.

[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.