[feat][SFT] Stream tokenization to arrow, bounding memory for large text datasets

[avigyabb]

Stacked on #1970 -- builds on the memory-mapped cache serving introduced there; the diff collapses once #1970 merges.

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). With this PR, the tokenize-on-load path never holds the tokenized dataset in memory -- text and VLM 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 = 1000 examples):

• Sequential path: tokenization becomes a generator feeding _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).
• 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, nothing loaded) into the cache. This also removes the old pickle-the-results round-trip from workers to controller.
• _save_to_cache accepts an arrow-backed Dataset directly (save_to_disk writes through in batches) in addition to list[dict].
• VLM streams too: validated against real Qwen3-VL processor output -- image tensors (pixel_values, image_grid_thw) round-trip through the streaming writer with full parity to Dataset.from_list (test_vlm_rows_stream_to_arrow, run under the vllm marker with the fsdp extras like the other VLM tests). VLM rows are the megabyte-sized ones, so they benefit most.
• Materializing fallback: disable_cache=True only (no arrow file to stream to or serve from).

Behavior change: inconsistent row keys now raise

The VLM validation surfaced a pre-existing hazard: datasets mixing text-only and multimodal rows silently lose their image columns to arrow schema inference (locked from the first example) -- identically under the old Dataset.from_list path, so training would proceed without images and no one would know. _write_tokenized_arrow now raises an explicit error on inconsistent row keys. Runs that previously trained silently-wrong on such datasets now fail loudly at tokenization.

Measured (cache-miss tokenization; peak anonymous heap = the memory that can OOM, vs reclaimable file-backed pages)

Rows	Streamed heap	Materialized heap (old)
60k	720 MB	1,569 MB
120k	722 MB (flat)	2,585 MB (linear)

The ~720 MB floor is torch/transformers/tokenizer imports; doubling the dataset moved the streamed heap by 2 MB. Streamed growth appears only in file-backed page cache (the arrow file being written/read), which the OS reclaims under pressure.

Memory profile of the tokenize-on-load path, across the stack

Phase	main (pre #1961)	after #1970	after this PR
Tokenization (cache miss)	O(dataset)	O(dataset), transient	O(writer batch)
Training (serving)	O(dataset) x (workers+1)	O(page cache)	O(page cache)

Tests

• Streamed-vs-materialized row parity with the flush size forced to 2 (multiple batch boundaries).
• Existing parallel-vs-serial parity test exercises the shard-writing workers end to end (real spawn pool, cache enabled).
• VLM tensor round-trip parity (vllm-marked); inconsistent-keys guard test.
• Full tests/train suite green: 884 (+36 tokenization, +3 vllm-marked).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request optimizes the tokenization pipeline by streaming tokenized rows into memory-mapped Arrow files in bounded batches instead of materializing the entire dataset in memory, significantly reducing RAM usage. Feedback on the changes highlights a potential issue in the parallel tokenization path where active references to memory-mapped Arrow files in the temporary shard directory might prevent the directory from being successfully cleaned up in the finally block, especially on Windows, and suggests explicitly deleting these references before returning.

[gemini-code-assist]

In the parallel tokenization path, shards and tokenized_ds hold active references to the memory-mapped Arrow files in shard_dir. Because the finally block is executed before the function returns, these references remain active when shutil.rmtree(shard_dir) is called. On Windows (and potentially some Unix environments), this will cause a PermissionError or prevent the files from being deleted, leading to disk space leaks. Explicitly deleting the references (del shards, tokenized_ds) before returning ensures the file handles are closed and the temporary directory can be successfully cleaned up.

            if self.sft_cfg.disable_cache:\n                # No cache directory to serve the mmap from (the shard dir is\n                # deleted below), so keep the materialized-list behavior.\n                res = tokenized_ds.to_list()\n                del shards, tokenized_ds\n                return res\n\n            _save_to_cache(cache_path, tokenized_ds)\n            del shards, tokenized_ds\n            mmapped = _load_from_cache(cache_path)\n            if mmapped is None:\n                raise RuntimeError(f\"Failed to reload tokenized dataset cache at {cache_path}\")\n            return mmapped