Allow document boundaries to come from the metadata file (fixes silent token loss on non-EOS-terminated documents) - #843
abhishekraok wants to merge 7 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58d8db6ef6
ℹ️ 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".
| if self._use_array_if_local is not None: | ||
| fields = fields + ("use_array_if_local",) |
There was a problem hiding this comment.
When the same source and work_dir have previously been prepared with another use_array_if_local value, changing this field only changes the dataset fingerprint; the three packing-cache paths are generated by _get_indices_path() from the source, size, long-document strategy, and dtype, without consulting that fingerprint. Consequently _pack_all_documents_into_instances() reuses the old files, so setting this to False can silently retain the EOS-inferred boundaries and the token loss this change is intended to prevent. Add this setting to the extra_ids used by all three packing-cache paths.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in 1722654 — this was a real P1 and the fix as submitted would have silently no-opped.
Verified the mechanism against source rather than taking it on faith: _get_indices_path() hashes the source paths, their sizes, and extra_ids, and never touches fingerprint, so adding the field to fingerprint_fields had no effect on the three packing-cache paths.
Fix: factored the shared tuple into _packing_cache_extra_ids so the three paths cannot drift apart, and appended the setting there. Following the existing source_group_size precedent it is only appended when not None, so caches built before the option existed are still found and the hardcoded fingerprint assertions in numpy_dataset_test.py still pass.
Added a regression test that prepares the same work_dir twice — first with the inferred boundaries, then with the metadata boundaries — and asserts the second run is not served from the first run's cache. I checked it genuinely catches the bug by reverting the fix: it fails, with the log line Reusing cached packing results, which is exactly the failure mode you described.
Data suite: 46 passed. ruff check clean. (ruff format wants to rewrap three pre-existing assert x, (...) statements at lines 470/555/576 of the test file; that is a black-version style drift unrelated to this change, so I left it alone.)
Adding `use_array_if_local` to `fingerprint_fields` was not sufficient. The three packing caches are looked up by path via `_get_indices_path()`, keyed on the source paths, sizes, long-document strategy and indices dtype -- it never consults `fingerprint`. So a work_dir already prepared with the EOS-inferred boundaries would serve `_pack_all_documents_into_ instances()` from the stale cache, silently keeping the boundaries and the token loss this change exists to prevent. Factor the shared `extra_ids` tuple into `_packing_cache_extra_ids` so the three paths cannot drift, and append the setting there. As with `source_group_size`, it is only appended when not `None`, so caches built before the option existed are still found and existing fingerprints are unchanged. Add a regression test that prepares the same work_dir twice, first with the inferred boundaries and then with the metadata boundaries, and asserts the second run is not served from the first run's cache. Without this fix it fails with "Reusing cached packing results". Reported by Codex review on #843. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1722654f9a
ℹ️ 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".
| dtype=self.dtype, | ||
| indices_dtype=self.indices_dtype, | ||
| long_doc_strategy=self._long_doc_strategy, | ||
| use_array_if_local=self._use_array_if_local, |
There was a problem hiding this comment.
When use_array_if_local=False is combined with generate_doc_lengths=True, packing now honors metadata boundaries, but __getitem__ still calls get_document_lengths() on the packed tokens and reconstructs boundaries from EOS tokens. For the non-EOS-terminated documents this option is meant to support, that merges the document with the next packed document in doc_lens, so intra-document attention masking allows attention across the real boundary and silently changes training semantics. Generate doc_lens from the selected cached document indices instead of rescanning input_ids.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Claude Code · session 6b836cd5 (claude.ai/code), posted on behalf of @abhishekraok
Confirmed and fixed in 6ef0280. This was a real P1 and I should have caught it when I fixed the first one — thanks.
Verified the merge before fixing it, on [1, 2, 3, 4] (terminator lost to truncation) followed by [5, 6, 7, 0] at sequence_length=8:
use_array_if_local=False, before -> [[8]] # one span; attention crosses the boundary
use_array_if_local=False, after -> [[4, 4]]
use_array_if_local=None -> [[8]] # unchanged, correct: the EOS scan defined these documents
__getitem__ already loads the documents individually, so the lengths are exact — no re-derivation needed. Trailing padding stays a final segment to match get_document_lengths, which reports it that way. The regression test asserts all three rows above and fails with assert [[8]] == [[4, 4]] without the fix.
I gated this on use_array_if_local is False rather than always using the exact lengths. On the inferred path every document ends in EOS by construction, so the two agree there, and keeping the scan means this PR cannot change default behavior. Always using the exact lengths is arguably cleaner and I'd be happy to do it, but it interacts with the bos_token_id branch and would be a default-behavior change — which is really the same question already open in the PR description about whether iter_document_indices should prefer the metadata file whenever one exists. Maintainer's call.
… to never stop (allenai#1876) * Add over_length_strategy for truncated SFT conversations `truncation_side` is `right`, so a conversation longer than `max_seq_length` loses its trailing EOS along with the excess. What remains has no terminator anywhere in it -- supervision to keep generating forever -- and because the cut lands mid-turn the trainable tail is a partial answer. Measured against the authoritative row boundaries in `token_ids_part_*.csv.gz`: on the Dolci-Think 32768 cache 491 of 24,292 rows (2.02%) end on a non-EOS token and 487 (2.00%) have that final token trainable; 490 of them sit at exactly 32,768, i.e. they are precisely the rows truncation cut. The 65536 cache has one such row. `over_length_strategy` selects what happens to those rows: `keep` is today's behavior, `terminate` replaces the final token with a trainable EOS so the model learns a long generation ends, and `drop` masks the row so `sft_tulu_filter_v1` removes it. Rows that fit are never touched. The default stays `keep`, and `sft_tokenize_fn_args()` omits the key at that default. The cache hash comes from `DATASET_CACHE_VERSION` plus a JSON encoding of these args, not from the tokenization code, so changing behavior at the default would leave every existing cache hash-valid but semantically stale. Omitting the key keeps default runs on their current caches while an opted-in run gets a distinct hash, so it can never be served a cache tokenized under another strategy. Two interactions worth noting. `terminate` must not unmask a row that `_tokenize_row_or_mask_out` masked out for underivable assistant spans, since one unmasked label would carry it past the filter and into training; that is guarded. And the `FileNotFoundError` in `olmo_core_finetune.py` prints the strategy in its re-tokenize command, because omitting a hash input there would point the reader at a different cache than the run wants. Reported in allenai#1859 as distinct from the token loss fixed by allenai/OLMo-core#843, which does not address it: correct document boundaries do not help a conversation that is longer than the window. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Add the PR URL to the CHANGELOG entry * Detect truncation properly and never synthesize EOS in a masked span Two review findings from Codex on allenai#1876, both real. Inferring truncation from "at the cap and not ending in EOS" is wrong in both directions. A conversation can render to exactly `max_seq_length` tokens on its own, or come from a template that does not end in EOS, and would then be rewritten by `terminate` or discarded by `drop` despite being intact. Conversely an over-length render cut exactly on some earlier turn's EOS ends in EOS and was skipped even though it lost text. Replace the inference with `_was_truncated()`, which asks the offset mapping whether any token reaches the end of the rendered string -- already computed on the main path, so it costs nothing. The fallback path now requests offsets too. `_tokenize_tulu_sft_with_assistant_labels` reports the flag rather than having callers guess. Second, `terminate` could write a trainable EOS into a masked span. When the cut lands in a later user/system/tool turn the final label is masked while earlier assistant turns keep the all-masked guard satisfied, so the old code turned a prompt token into a trainable EOS -- teaching the model to stop partway through its own input. Nothing is trained on those positions, so there is no unterminated supervision to repair there: `terminate` now leaves such a row alone. This also subsumes the previous all-masked guard, since an all-masked row necessarily has a masked final label. Verified end to end on a multi-turn row cut inside a later user turn: `terminate` now leaves it byte-identical to `keep`, where it previously wrote a trainable EOS. Tests cover both directions of the truncation signal and both branches of the masked/trainable split. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Widen the _was_truncated offsets annotation `offset_mapping` arrives via `Tensor.tolist()`, which yields `list[list[int]]`, not `list[tuple[int, int]]`. Annotate the parameter as `Sequence[Sequence[int]]`. Caught by `ty` in CI, not locally: the `ty` constraint is `>=0.0.1a13` with no upper bound, so CI resolves a newer and stricter version than the 0.0.12 installed here. Local `ty check` passes either way. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Narrow rendered to str in the label-derivation fallback The real cause of the CI type error, which the previous commit misread. `apply_chat_template` returns a broad union; the main path narrows it with `assert isinstance(rendered, str)` immediately after, and the fallback branch added here did not, so passing `rendered` to `_was_truncated` failed. The caret in the diagnostic was on `rendered`, not on the offsets argument I widened. The widened `Sequence[Sequence[int]]` annotation stays: `Tensor.tolist()` really does yield `list[list[int]]`, so it was wrong before, just not what CI was complaining about. Verified with `uvx ty@latest`: no diagnostics remain on `_was_truncated`, `_apply_over_length_strategy` or `sft_tokenize_fn_args`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Trim over_length_strategy comments and docstrings Keep the reasoning a reader of the code needs; drop measurements and history that the PR already records. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Kevin Farhat <kevinfarhat@allenai.org>
`iter_document_indices` infers document boundaries by scanning the token array for EOS whenever `eos_token_id` and `dtype` are given for a local path, which is every call from the FSL dataset classes. That is only correct if every document is EOS-terminated. When a producer truncates documents to a maximum length and the terminator goes with the tail, the affected document merges with the one that follows it. `LongDocStrategy.truncate` then keeps only the head of the merged span, so the following document is never yielded and never reaches training. Nothing in the pipeline reports this. Measured on an SFT cache whose producer writes a boundaries metadata file: the inferred path yielded 23,801 documents covering 97.98% of tokens, the metadata path 24,292 documents covering 100.00%, with an identical maximum document length of 32,768. Reading the gzipped metadata was also faster than scanning 268M tokens. Plumb `use_array_if_local` through `pack_documents_into_instances`, `segment_documents_into_instances`, `NumpyPackedFSLDataset` and `NumpyPackedFSLDatasetConfig` so callers can select the authoritative source, include it in the dataset fingerprint so a cached index built with the other source is not reused, and document the hazard on `iter_document_indices`. Default behavior is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Adding `use_array_if_local` to `fingerprint_fields` was not sufficient. The three packing caches are looked up by path via `_get_indices_path()`, keyed on the source paths, sizes, long-document strategy and indices dtype -- it never consults `fingerprint`. So a work_dir already prepared with the EOS-inferred boundaries would serve `_pack_all_documents_into_ instances()` from the stale cache, silently keeping the boundaries and the token loss this change exists to prevent. Factor the shared `extra_ids` tuple into `_packing_cache_extra_ids` so the three paths cannot drift, and append the setting there. As with `source_group_size`, it is only appended when not `None`, so caches built before the option existed are still found and existing fingerprints are unchanged. Add a regression test that prepares the same work_dir twice, first with the inferred boundaries and then with the metadata boundaries, and asserts the second run is not served from the first run's cache. Without this fix it fails with "Reusing cached packing results". Reported by Codex review on #843. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1722654 to
0796e1e
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
With use_array_if_local=False, packing honored the metadata boundaries but __getitem__ still rebuilt doc_lens by scanning the packed tokens for EOS, so a document that lost its terminator merged with the next one. doc_lens drives the block-diagonal attention mask, so tokens could attend across a real document boundary. The loaded per-document tensors give the lengths exactly; trailing padding stays a final segment to match get_document_lengths. The inferred path is untouched, where every document ends in EOS by construction. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6ef0280904
ℹ️ 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".
| if (padding := input_ids.numel() - sum(doc_lens)) > 0: | ||
| # `get_document_lengths` reports trailing padding as a final segment. | ||
| doc_lens.append(padding) |
There was a problem hiding this comment.
When use_array_if_local=False, generate_doc_lengths=True, and pad_token_id == eos_token_id (as in TokenizerConfig.gpt2() and many HF tokenizers without a distinct pad token), appending all trailing padding as one document does not match get_document_lengths(), which treats each EOS padding token as a separate segment. A sparsely filled packed instance can therefore report a padding document almost as long as sequence_length; this inflates max_doc_lens and makes varlen attention process the padding as one large quadratic segment, potentially causing severe slowdowns or OOMs. Preserve the existing EOS-based segmentation for the padding region while deriving only real-document boundaries from metadata.
Useful? React with 👍 / 👎.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 713eedc357
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| out["doc_lens"] = get_document_lengths( | ||
| input_ids, self.eos_token_id, bos_token_id=self.bos_token_id | ||
| ) | ||
| if self._use_array_if_local is False: |
There was a problem hiding this comment.
When use_array_if_local=True and a source path is remote, iter_document_indices() still reads the metadata because its array path requires not is_url(data_path), but this condition assumes that only an explicit False used metadata. With generate_doc_lengths=True, a remote document lacking EOS is therefore packed on its metadata boundary and then merged with its successor in doc_lens, allowing attention across the real boundary. This also affects mixed local/remote source groups; select the doc_lens strategy based on the effective boundary source rather than the raw option.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Claude Code · session 6b836cd5 (claude.ai/code), posted on behalf of @abhishekraok
Confirmed and fixed in cba9ff7. Right on both counts: the gate is use_array_if_local and not is_url(data_path), so a remote source reads the metadata whether the option is True or None, and my doc_lens gate only tested is False. Mixed local/remote groups had the same hole.
Now gated on the effective boundary source via a _packed_from_metadata_boundaries(source_paths) helper, asserted directly since a real remote read isn't available in CI:
(True, ["s3://..."]) -> True
(None, ["s3://..."]) -> True
(True, [local, "s3://..."]) -> True # mixed group
(True, [local]) -> False
(False, [local]) -> True
Reverting the is_url term fails that test.
Verifying this also turned up a bug of my own: get_document_lengths folds trailing padding into the final document when bos_token_id is set, but reports it as its own segment otherwise — [4, 12] vs [4, 4, 8] on the same instance. My exact-length path always did the latter, so it would have disagreed with the scan on bos-delimited data. Fixed and pinned with its own test.
AkshitaB
left a comment
AkshitaB
left a comment
There was a problem hiding this comment.
LGTM, except lets fix the style check.
Regarding default - let's keep it as it is for now, so our existing set of experiments remain comparable.
iter_document_indices reads the metadata file for any URL source whatever use_array_if_local says, so gating doc_lens on the raw option missed remote sources: they were packed on metadata boundaries but doc_lens still rescanned for EOS, merging an unterminated document into the next one. Gate on the effective source instead, via a _packed_from_metadata_boundaries helper. Also match get_document_lengths' padding convention, which folds trailing padding into the final document when bos_token_id is set rather than reporting it as its own segment. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Claude Code · session 6b836cd5 (claude.ai/code), posted on behalf of @abhishekraok Thanks — and thanks for merging main in. Style is fixed in cba9ff7. The failure was a merge artifact: your merge left one blank line between my last test and the next function, and black 23.12.1 wants two. Default: agreed, leaving One more thing in that same commit: Codex flagged a real gap on the re-review, and it was right. Net change since your review is one commit: the style fix plus that gate. Nothing about default behavior moved. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cba9ff7994
ℹ️ 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".
| "document-indices", | ||
| *source_paths, | ||
| extra_ids=(self._long_doc_strategy, self.indices_dtype.__name__), | ||
| extra_ids=self._packing_cache_extra_ids, |
There was a problem hiding this comment.
When use_array_if_local=False, packing depends on each .csv.gz sidecar, but these cache paths still hash only the token-array path/size and the option value. If a producer corrects or regenerates the boundary metadata while leaving the .npy unchanged, preparing again in the same work_dir silently reuses all three old packing files and ignores the new authoritative boundaries. Include the metadata sidecar identity (at least its path and size/version) in the cache key for the metadata-backed mode.
Useful? React with 👍 / 👎.
…ies-from-metadata # Conflicts: # CHANGELOG.md
2 participants
The problem
iter_document_indicesinfers document boundaries by scanning the token array foreos_token_id. That path auto-enables whenevereos_token_idanddtypeare supplied for a local path:Every call from the FSL dataset classes supplies both, so the scan is always used and the metadata file is never read.
Inferring is only correct if every document is EOS-terminated. When a producer truncates documents to a maximum length and the terminator goes with the discarded tail, the affected document merges with the one that follows it.
LongDocStrategy.truncatethen keeps only the head of the merged span — which is just the truncated document again — and never yields the remainder, so the following document does not reach training. Nothing in the pipeline reports this: no warning, no count, and the resulting instances look normal.The victim is not the long document. It is the ordinary-length document that happened to come after one.
Measurement
On an SFT cache produced by open-instruct (which writes a
token_ids_part_*.csv.gzboundaries file alongside each.npy), throughpack_documents_into_instances— the entry pointNumpyPackedFSLDatasetuses:use_array_if_local=False)Maximum document length is 32,768 on both paths, so this changes only how documents are delimited, not their content and not truncation semantics. Roughly 2% of tokens were silently dropped, consistently across sampled shards (1.82–2.05%).
The discrepancy is exactly the mechanism: 24,292 − 23,801 = 491 merges, against 490 documents sitting at precisely the producer's truncation limit.
Reading the gzipped metadata also took ~0.03s versus ~0.5–0.9s to scan 268M tokens, so the correct path is typically faster.
The change
use_array_if_localplumbed throughpack_documents_into_instances,segment_documents_into_instances,NumpyPackedFSLDataset, andNumpyPackedFSLDatasetConfig.NumpyPackedFSLDataset.fingerprint_fieldswhen set, so a cached instance index built with the other boundary source is not silently reused. Follows the existingsource_group_sizebackwards-compat pattern, so fingerprints are unchanged when the option is left atNone... warning::oniter_document_indices, where a reader would look.pack_documents_into_instancesandsegment_documents_into_instances.Default behavior is unchanged — this only makes the correct source reachable.
Resolved
Whether
iter_document_indicesshould instead prefer the metadata file whenever one exists was raised here and answered by @AkshitaB: keep the default as it is for now, so the existing set of experiments stays comparable. Any such change would be separate and deliberate.Testing
pytest src/test/data/utils_test.py src/test/data/numpy_dataset_test.py— 49 passed.black --check .andisort --check-only .clean across the repo (557 files).pack_documents_into_instances, as above.Not tested: a full
NumpyPackedFSLDatasetConfig.build()+prepare()across a multi-shard glob with the flag set. I verified the boundary and packing functions per shard, not an end-to-end dataset build.Context and the full investigation: allenai/open-instruct#1855.
🤖 Generated with Claude Code