[DistillationTrainer refactor] Pin signature columns to ["prompt", "image", "images"]

[qgallouedec]

Item 17 (Group A, last) — stacked on PR 16. Part of #6449.

Align _set_signature_columns_if_needed with GRPO/RLOO verbatim: pin to ["prompt", "image", "images"] (with their exact comment). Drops the bespoke ["prompt", "prompts", "prompt_attention_mask", "messages", "chat_template_kwargs", "tools"] list and the super() merge — the collator/generation only read the prompt column now (messages gone at item 16), and prompts/prompt_attention_mask were internal batch keys, never dataset columns.

Runtime no-op (remove_unused_columns is forced False); pure code alignment.

Verified: pytest tests/experimental/test_distillation_trainer.py — 39 passed; ruff clean.

---

Note

Low Risk
No intended runtime change because column stripping is disabled; only the signature-column override logic changes.

Overview
DistillationTrainer._set_signature_columns_if_needed is rewritten to match GRPO and RLOO: it no longer calls super() or merges a long custom column list, and instead sets _signature_columns to ["prompt", "image", "images"] when unset, with the same rationale comment as those trainers.

The dropped names (prompts, prompt_attention_mask, messages, chat_template_kwargs, tools, etc.) reflected old dataset/chat paths or collator batch keys, not columns the trainer pipeline still needs after the prior refactor. remove_unused_columns stays False, so training behavior should be unchanged—this is alignment and cleanup.

^{Reviewed by Cursor Bugbot for commit 4f91a90. 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.