[feat][SFT] Stream tokenization to arrow, bounding memory for large text datasets - #1971
Draft
avigyabb wants to merge 3 commits into
Draft
[feat][SFT] Stream tokenization to arrow, bounding memory for large text datasets#1971avigyabb wants to merge 3 commits into
avigyabb wants to merge 3 commits into
Conversation
The tokenized cache is already an arrow-backed HF Dataset on disk in the trainer's internal row form, but _load_from_cache materialized it back into a list[dict] (O(dataset) RAM, re-pickled into every spawn dataloader worker). Serve it through the same map-style mmap wrapper as pretokenized stores instead, with no transform attached (cached rows are already normalized): - _load_from_cache returns PretokenizedDataset(load_from_disk(...), lengths), with per-row lengths from arrow offsets via the new sequence_lengths_from_arrow helper (chunked, no row materialization). - Fresh tokenization round-trips through the cache in both the sequential and parallel paths, so cold runs also train memory-mapped. - disable_cache=True keeps the in-memory list (wrapped in TextDataset); with no arrow file on disk there is nothing to map. Tokenization cost is unchanged; only the residency of the results changes: the text path now has the same memory profile as pretokenized stores (O(page cache), workers pickle a file reference). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
Since the tokenized-dataset cache is now served through the same class, the old name was wrong for one of its two roles: the class is a map-style view over any validated arrow store in (or normalizable to) the trainer's internal row form. Rename before anything external depends on the name (NovaSky-AI#1961 merged it one release ago). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
…ext datasets Cache-miss tokenization accumulated every tokenized row in a Python list before writing the arrow cache (and on the parallel path, pickled each worker's slice back to the controller), so first-run peak memory was O(dataset) even though training then served the cache memory-mapped. Tokenized rows now stream into arrow files in bounded batches via HF's ArrowWriter (flush granularity: _TOKENIZE_WRITER_BATCH_ROWS = 1000): - Sequential path: tokenize as a generator -> _write_tokenized_arrow to a temp shard -> _save_to_cache(Dataset.from_file(shard)) -> serve memory-mapped. The tokenized dataset is never resident; size is bounded by disk, not RAM. - Parallel path: each spawn worker streams its slice into a per-worker arrow shard and returns only its row count; the controller concatenates the memory-mapped shards (a view) into the cache. This also removes the pickle-results round-trip. - _save_to_cache accepts an arrow-backed Dataset directly (save_to_disk writes through in batches) in addition to list[dict]. - Materializing fallbacks, documented: disable_cache=True (no arrow file to stream to or serve from) and VLM datasets (image tensors only round-trip through Dataset.from_list). New test forces multiple ArrowWriter flush boundaries and checks row parity against the materialized path; the existing parallel-vs-serial parity test now exercises the shard-writing workers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Closes the transient flagged in #1970's known-limitations note: cache-miss tokenization accumulated every tokenized row in a Python list before writing the arrow cache (and on the parallel path, pickled each worker's slice back to the controller), so first-run peak memory was still O(dataset) even though training then served the cache memory-mapped. With this PR, the tokenize-on-load path never holds the tokenized dataset in memory — text datasets are bounded by disk, not RAM, end to end.
How
Tokenized rows stream into arrow files in bounded batches via HF's
ArrowWriter(flush every_TOKENIZE_WRITER_BATCH_ROWS = 1000examples):_write_tokenized_arrow(tokenize → bounded arrow writes to a temp shard →_save_to_cache(Dataset.from_file(shard))→ serve memory-mapped). Peak memory during tokenization is O(writer batch)._save_to_cachenow accepts an arrow-backedDatasetdirectly (save_to_diskwrites through in batches) in addition tolist[dict].Materializing fallbacks, documented:
disable_cache=True(no arrow file to stream to or serve from — keeps the in-memory list behavior) and VLM datasets (image tensors only round-trip throughDataset.from_list; VLM already forces sequential tokenization).Memory profile of the text path, before → after this stack
Tests
_TOKENIZE_WRITER_BATCH_ROWSmonkeypatched to 2, forcing multiple flush boundaries.tests/trainsuite green: 884 passed (+35 tokenization tests).🤖 Generated with Claude Code