Stop enable_padding_free_metadata writing seq_lengths into the caller's examples

[vineethsaivs]

The padding-free collator wrapper adds seq_lengths to the caller's examples while deriving metadata. Pass copies containing the derived lengths to the wrapped collator, leaving the original examples untouched.

Rebased onto current main. The two CPU source-harness regressions pass for missing, explicit and null lengths, and changed-file pre-commit passes. Full package tests cannot import on this host and remain for upstream CI. The old CI logs have expired, so this refresh requests a fresh run.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b88595cb64

ℹ️ 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".

[chatgpt-codex-connector]

Preserve derived lengths for the wrapped collator

When a padding-free batch has a seq_lengths key whose value is None (for example, a nullable/custom metadata column), this fallback now only fixes the local packed_seq_lengths tensor and still passes the original None through to original_torch_call. The TRL padding-free collator checks for the key on the first example and then consumes example["seq_lengths"], so these rows can now error or produce inconsistent position_ids even though this wrapper already derived a valid fallback. Please keep caller data immutable by passing a copied/normalized examples list to the underlying collator instead of dropping the derived metadata.

Useful? React with 👍 / 👎.

[vineethsaivs]

That case is already covered. example.get("seq_lengths") returns None both when the key is absent and when its value is None, so a null column takes the same derive branch:

lengths = example.get("seq_lengths")   # None, either way
if lengths is None:
    ids = example.get("input_ids")
    ...
    lengths = [len(ids)]
    collated[index] = {**example, "seq_lengths": lengths}

{**example, "seq_lengths": lengths} overwrites the null, so the wrapped collator sees the derived list and never the None.

The one row that still reaches it untouched is a row with no input_ids at all, which hits the continue. That is unchanged from before this PR, where the write-back was also skipped on that path.

[vineethsaivs]

You are right, and the reason is bigger than a nullable column, so thanks for pushing on it. Fixed in 00450464b.

I had framed the write-back as caching, which was wrong. It was load-bearing. TRL's DataCollatorForLanguageModeling decides whether to use seq_lengths at all from the first row (sft_trainer.py line 464: [example["seq_lengths"] for example in examples] if "seq_lengths" in examples[0] else None) and only then reads it off every row, so removing the key changes what it builds.

Two things break, not one. Your null case reaches sum(None). The one I had not seen is a batch mixing an unpacked row with a packed one, where the key is missing on row 0 so TRL ignores seq_lengths entirely and falls back to torch.arange:

with derived lengths : [0, 1, 2, 0, 1, 0, 1]
key missing on row 0 : [0, 1, 2, 0, 1, 2, 3]

The packed row's second document starts at position 2 instead of 0, so it attends across the document boundary. That is silent, and it is exactly what padding-free metadata exists to prevent.

Your suggested shape is what I did: pass a normalized copy to the wrapped collator, built only when something was actually derived, so the common path still hands the caller's own list straight through with no per-row copies. Ran it three ways against upstream/main, this branch before the change and after: main passes everything except leaving seq_lengths on the caller's rows, the previous commit fixes that and breaks the null case and the mixed batch, and the new commit passes all five. The test now asserts what reaches the collator rather than only what comes back.

[vineethsaivs]

All three red checks are pre-existing, and I checked rather than assumed. Repo tests (CPU) fails the same six tests/studio/ assertions with the identical 6 failed / 5193 passed / 129 skipped counts on #8198, which is someone else's PR touching unrelated code; that job also passes --ignore=tests/utils, so it never sees this file. The two Core jobs do run tests/utils/test_packing.py, and their only failures are the two test_generated_fullgraph_fallback.py ones that were also red on #8117 when it merged. The tests here pass in both Core jobs.

[vineethsaivs]

Addressed already, flagging it so the thread is not left looking open.

This was raised against b88595cb and is a fair reading of that commit: not writing the derived lengths back onto the example must not mean the wrapped collator stops seeing them, and TRL's padding-free collator decides whether to use seq_lengths at all from examples[0] and then reads it off every row.

00450464 covers exactly that, with test_enable_padding_free_metadata_still_hands_derived_lengths_to_the_collator, which asserts the wrapped collator still receives the lengths for every row while the caller's examples stay unmutated. The sibling test pins the no-mutation half, including the explicit-seq_lengths path.

[danielhanchen]

Confirmed the write-back is still there in unsloth/utils/packing.py and that the follow-up commit keeps the derived lengths reaching the wrapped collator via a copy. The three Core jobs are green on main now, so could you push a refresh and check they pass here before I review?

[vineethsaivs]

Rebased on current main. Both CPU source-harness regressions and pre-commit pass; the expired CI logs are unavailable, so this push requests fresh Core runs.