[Feature] Add DistillationLoss for token-level LLM knowledge distillation

[theap06]

Description

Implements DistillationLoss, the on-policy KL knowledge-distillation objective for LLM policies proposed in #3921, under torchrl/objectives/llm/.

A student policy is distilled toward a (typically stronger, frozen) teacher by minimizing a KL divergence estimated from the per-token log-probabilities the two models assign to the same tokens. It mirrors SFTLoss and reuses the same teacher log-prob plumbing as GRPOLoss/SFTLoss:

• student log-probs are computed by running actor_network on the input history;
• teacher log-probs are consumed from the tensordict, as written by RetrieveLogProb(teacher, log_probs_full_key=("teacher_log_probs", "full")) or an offline scoring pass, and are detached inside the loss;
• the KL is estimated per token with the k3 estimator, exposed as reverse_kl_token_estimate (the same expm1(d) - d approximation the GRPO/SFT KL regularizers inline);
• masking follows the SFT convention: assistant-and-attention masks from ("masks", ...) when present, tokenizer fallback otherwise; masked positions are zeroed on both log-prob tensors before the estimator so extreme masked values cannot poison gradients.

Both directions are supported via kl_direction:

• "reverse" (default): KL(student || teacher) for on-policy distillation on student-generated tokens (mode-seeking, GKD-style);
• "forward": KL(teacher || student) for off-policy distillation on teacher/corpus data (mass-covering).

Plus reduction (mean/sum/none), normalize_by_seq_length, assistant_only=False for plain-corpus distillation, and set_keys() support for custom key layouts.

Example

transform = RetrieveLogProb(
    teacher,
    log_probs_full_key=("teacher_log_probs", "full"),
    assistant_only=True,
    tokenizer=tokenizer,
)
with torch.no_grad():
    data = transform(data)

loss = DistillationLoss(actor_network=student, tokenizer=tokenizer, kl_direction="reverse")
loss_vals = loss(data)
loss_vals.loss_distill.backward()

The class docstring contains a complete runnable example (verified end-to-end against a randomly initialized OPT pair).

Tests

TestDistillation in test/llm/test_llm_objectives.py, following the TestSFT pattern (tiny OPT models + RetrieveLogProb):

• estimator: exact zero at equality, nonnegativity, closed-form equivalence, shape validation, and a gradient-descent check confirming a step on the loss moves the student toward the teacher (all transformers-free);
• loss in both directions: output type/shapes, finiteness, gradients reach the student;
• exact-zero loss when the teacher is the student (both directions);
• reduction semantics (mean/sum consistent with none) and length normalization;
• custom teacher key via set_keys(); missing-teacher KeyError; invalid kl_direction/reduction validation.

Docs entries added to llms.rst and llms_objectives.rst.

Notes

• No Hydra config companion: no LLM loss currently has one (GRPOLoss/SFTLoss included), so this follows the existing pattern.
• A sota-style teacher/student recipe is planned as a follow-up PR, mirroring how SFT/GRPO recipes are sequenced.

Closes #3921

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4038

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 15 Awaiting Approval

As of commit eb849fa with merge base ae421b9 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.