[feat][SFT] Serve the tokenized-dataset cache memory-mapped

[avigyabb]

What

The tokenize-on-load path (train_datasets) no longer materializes the tokenized dataset in memory for training. Follow-up to #1961 flagged in review: "TextDataset currently materializes the entire dataset in memory, but it doesn't have to be."

How

One observation makes this small: the tokenized-dataset cache is already an arrow-backed HF Dataset on disk in the trainer's internal row form -- it is, in effect, a pretokenized store. But _load_from_cache immediately materialized it back into a list[dict] via .to_list() (O(dataset) RAM, re-pickled into every spawn dataloader worker). This PR serves it through the same map-style mmap wrapper the pretokenized path uses, with no transform attached (cached rows are already normalized):

• _load_from_cache returns MemoryMappedDataset(Dataset.load_from_disk(...), lengths) -- memory-mapped, rows page in per accessed batch through the existing dataloader prefetch. Per-row lengths come from arrow offsets (sequence_lengths_from_arrow, chunked scan, nothing materialized).
• Fresh tokenization round-trips through the cache in both the sequential and parallel paths (_save_to_cache -> reload mmap'd), so cold runs also train from the mmap. Tokenization cost is unchanged -- it was always one-time; only the residency of the results changes.
• disable_cache=True keeps today's behavior (in-memory list wrapped in TextDataset): with no arrow file on disk there is nothing to map.

This also resolves the old TODO at the cache-save site ("migrate to HF datasets + dataloader so we don't materialize the full dataset in memory").

Rename: PretokenizedDataset -> MemoryMappedDataset

Since the class now serves two roles (pretokenized stores and the tokenized cache), the old name was wrong for one of them: it is a map-style view over any validated arrow store in (or normalizable to) the trainer's internal row form. Renamed now, before anything external depends on the name (#1961 merged it one release ago). The module keeps the pretokenized.py name for diff locality; moving the class to sft_dataset.py could be a later cleanup.

Known limitation

Cache-miss tokenization still materializes the full list[dict] transiently (tokenize -> save -> reload mmap'd); the list dies at function exit, so this bounds first-run peak memory, not steady state. Closing the transient too would mean chunked tokenize-and-append arrow writing -- a possible follow-up.

Effect

With caching enabled (the default), the text path now has the same memory profile as pretokenized stores: resident memory is the OS page cache (reclaimable) instead of the dataset, and dataloader workers pickle a file reference instead of receiving a full copy -- the old peak was (dataloader_num_workers + 1) x the dataset.

Tests

• New: cache-backed _load_and_tokenize returns a memory-mapped SFTDataset with arrow-derived lengths, row-parity against the materialized (disable_cache=True) path.
• Existing cache round-trip / parallel-vs-serial parity tests pass unchanged (they compare rows by iteration, which the map-style dataset satisfies).
• Full tests/train suite green: 884 passed.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request optimizes memory usage during dataset loading and tokenization by serving cached tokenized datasets as memory-mapped PretokenizedDataset instances instead of fully materializing them in memory as lists. This is achieved by extracting sequence lengths directly from arrow offsets and returning the memory-mapped dataset when caching is enabled. A new test has been added to verify this behavior. I have no feedback to provide.

[gemini-code-assist]

Code Review

This pull request refactors the tokenization caching mechanism to serve cached datasets as memory-mapped MemoryMappedDataset (formerly PretokenizedDataset) instances instead of fully materializing them in memory as lists of dicts, significantly reducing memory overhead. Feedback suggests optimizing the sequence length extraction in sequence_lengths_from_arrow by specifying only the "input_ids" column when calling with_format on the dataset to avoid loading unnecessary columns.

[gemini-code-assist]

To optimize performance and reduce memory overhead, specify columns=["input_ids"] when calling with_format. This prevents Hugging Face Dataset from loading, formatting, or copying other columns (such as large VLM image tensors like pixel_values) when we only need to scan the sequence lengths of input_ids.

Suggested change

	arrow_ds = dataset.with_format("arrow")
	arrow_ds = dataset.with_format("arrow", columns=["input_ids"])