[DistillationTrainer refactor] Remove messages-format support and prompt-length config

[qgallouedec]

Item 16 (Group A) — stacked on PR 15. Part of #6449.

The trainer is on-policy prompt-only now (item 15), so the collator's messages branch and the prompt-truncation config are dead weight. GRPO has neither — only max_completion_length.

• Collator: drop the messages branch, the completion-tokenization/trimming path, and prompt truncation. Now prompt-only: tokenize example["prompt"], no completion. __init__ loses max_length/max_prompt_length/messages_key.
• Config: remove max_length and max_prompt_length (+ their __post_init__ logic). Keep max_completion_length.
• Tests: delete the messages-deprecation test; _make_local_trainer → prompt-only; test_loss_normalizes hand-builds a batch with completion tokens (collator no longer supplies them). Server test → prompt-only dataset + drop removed kwargs (server is on-policy, inherits generation).

Verified: pytest tests/experimental/test_distillation_trainer.py (39) + test_server_distillation_trainer.py (12) pass; ruff clean.

---

Note

Medium Risk
Breaking change for callers still using messages datasets or max_length/max_prompt_length; training behavior for prompt-only on-policy paths should be unchanged.

Overview
DistillationTrainer is aligned with on-policy, prompt-only training (same dataset shape as GRPO). The collator no longer supports conversational messages datasets, no longer tokenizes or truncates completions from the dataset, and only formats the prompt column for generation.

DistillationConfig drops max_length and max_prompt_length (and related __post_init__ validation/auto-default). Sequence budgeting for training is max_completion_length only, matching generation limits.

Tests switch to conversational_prompt_only, remove the messages deprecation test, and hand-build batches where loss tests need completion tokens (since the collator no longer supplies them). Server distillation tests use prompt-only fixtures and omit the removed config kwargs.

^{Reviewed by Cursor Bugbot for commit 76115bf. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cmpatino]

One small nit, but it's non-blocking.

Should we add a comment about the change that removes max_prompt_length? I think some users with long prompts may start seeing OOM errors that weren't showing up before, and it would be good to let them know this might be the cause.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 5aac5aa. Configure here.}

[cursor]

Eval skips on-policy generation

High Severity

Removing the messages path leaves only prompt-only batches (all labels -100), while _prepare_inputs still returns the raw batch whenever the model is not training. Evaluation therefore never generates completions and reports a vacuous eval_loss of zero. GRPO generates during eval for the same prompt-only setup.

Additional Locations (1)

• trl/experimental/distillation/distillation_trainer.py#L635-L637

 

^{Triggered by project rule: ../.ai/AGENTS.md}

^{Reviewed by Cursor Bugbot for commit 5aac5aa. Configure here.}