SFT: Truncate during dataset preparation, not collation

[qgallouedec]

Move sequence truncation in SFTTrainer from the data collator into _prepare_dataset.

Why

Truncation is a pure per-example slice. Doing it once during (cached) dataset preparation is cleaner than recomputing it in the collator on every batch, and keeps all dataset shaping (tokenize → build labels → truncate → pack) in one place.

It will also allow to drop rows with no trainable rewards, see #6025

Changes

• _prepare_dataset now truncates input_ids and labels to max_length (respecting truncation_mode), right after labels are built. Skipped when packing (packing already chunks to max_length).
• build_labels now drops the completion_mask / assistant_masks columns: they're fully baked into labels.
• DataCollatorForLanguageModeling no longer truncates: the max_length and truncation_mode arguments are removed.

⚠️ Behavior change

With skip_prepare_dataset=True, preparation (and therefore truncation) is skipped. The dataset must already be truncated.

---

Note

Medium Risk
Changes core SFT data pipeline and loss-relevant truncation order (labels before truncate), with a documented break for skip_prepare_dataset=True users who relied on collator truncation.

Overview
SFT sequence truncation now runs in cached _prepare_dataset (after label building, before packing) instead of on every batch in DataCollatorForLanguageModeling. input_ids and labels are sliced to max_length using truncation_mode (keep_start / keep_end); packing skips this step.

DataCollatorForLanguageModeling only pads now—max_length and truncation_mode were removed from the collator and from how the trainer constructs it.

When labels are built from masks, assistant_masks / completion_mask are dropped after being folded into labels.

Behavior change: with skip_prepare_dataset=True, truncation no longer happens anywhere; examples must already fit max_length. Tests were updated (collator truncation tests removed; preparation/truncation and #3927 assistant-only cases assert truncation at prepare time).

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0579163751

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".