[feat][SFT] Memory-map pretokenized stores instead of materializing rows

[avigyabb]

What

load_from_pretokenized now returns a map-style, arrow-backed PretokenizedDataset instead of a fully materialized list[dict]. Rows live in memory-mapped arrow files on disk; accessing row i pages in only the ~4 KB pages that hold it, so resident memory is the OS page cache (reclaimable) rather than the dataset.

This removes the two scaling walls of the materialized path for large pretokenized stores (#1927):

• RAM = O(dataset): every tokenized row held as Python objects for the whole run.
• x (dataloader_num_workers + 1): spawn-based dataloader workers each receive a full pickled copy of the list. The mmap dataset pickles as a file reference — workers re-map the same files, sharing pages instead of duplicating them.

How

• Validation stays eager, becomes vectorized: chunked arrow scans (bounded memory, 200k rows per chunk) over lengths / mask values / attention_mask / VLM pairs — same fail-fast-at-load behavior and same error messages as before, without materializing rows.
• Row dropping (empty loss window, over-length VLM) becomes a vectorized keep-mask + dataset.select().
• Normalization (num_actions inference, truncation, all-ones attention_mask) runs lazily per accessed batch via a picklable transform, so dataloader workers do it off the training critical path.
• sequence_lengths come from arrow offsets — _log_dataset_stats no longer materializes the store. Multi-store concat is a torch ConcatDataset view.
• Map-style contract is preserved, so samplers (random / DataMixingSampler / custom), weighted mixing, and the data.pt StatefulDataLoader resume flow are all unchanged — mid-epoch resume verified bit-exact, including with spawn workers.

Measured (100k rows x 256 tokens, real FSDP training on 1xL4, ~3.6 s steps)

	materialized list[dict]	mmap (this PR)
Load time	20.8 s	0.8 s
Resident memory	O(dataset), x(workers+1)	O(validation chunk); page cache thereafter
Data cost per train step	2.5 ms (w0)	11 ms (w0), 0.8 ms (w2)

Row content is parity-tested identical to the materialized path, and identical-seed training reproduces its loss curve.

Tests

tests/train/test_sft_pretokenized.py — full suite passes; covers format detection, schema validation/row-dropping parity, VLM rows, collation, StatefulDataLoader mid-epoch resume, and spawn-worker pickling.

Follow-ups (stacked on this PR)

This is Part 1 of a trillion-token SFT dataloading plan: Part 2 adds cloud row-group random access (RowGroupPretokenizedDataset, fetches only the row groups a batch touches — no download), Part 3 adds a storage-aligned shuffle sampler (1x read amplification) and a disk cache tier. Complements #1933 (whole-store download path).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request transitions pretokenized dataset loading to a memory-mapped approach rather than loading the entire dataset into RAM. It introduces eager, vectorized validation and row filtering over Arrow columns, alongside lazy row normalization at access time. The review feedback highlights two important improvements: first, correcting the loss mask validation logic to strictly reject fractional values (such as 0.5) which could cause downstream crashes; second, reducing the validation chunk size from 200,000 to 50,000 rows to prevent excessive memory usage and potential OOM crashes when handling datasets with long sequence lengths.

[gemini-code-assist]

The check (flat_mask < 0) | (flat_mask > 1) only ensures that the values are within the range [0, 1], which incorrectly allows fractional values like 0.5. This violates the requirement that loss_mask must contain only 0s and 1s (as stated in the error message on line 199) and can cause ValueError crashes later in _NormalizeTransform when calling loss_mask.index(1).\n\nTo strictly enforce that the mask contains only 0s and 1s, use (flat_mask != 0) & (flat_mask != 1).

Suggested change

	invalid = (flat_mask < 0) | (flat_mask > 1)
	invalid = (flat_mask != 0) & (flat_mask != 1)

[gemini-code-assist]

At _VALIDATION_CHUNK_ROWS = 200_000, if the dataset has long sequence lengths (e.g., 4096 or 8192 tokens), the temporary NumPy arrays created during vectorized validation (such as positions, window, and np.repeat outputs on lines 219-220) can consume tens of gigabytes of CPU RAM, potentially leading to Out-Of-Memory (OOM) crashes.\n\nReducing the chunk size to 50_000 significantly bounds the peak memory usage while keeping the Python loop overhead negligible and maintaining high vectorized performance.

Suggested change

	_VALIDATION_CHUNK_ROWS = 200_000
	_VALIDATION_CHUNK_ROWS = 50_000

[SumanthRH]

Nice! Only major comment is on providing a better API for text / pretokenized datasets to simplify logic in the SFTTrainer

class SFTDataset(torch.utils.data.Dataset, abc.ABC):
		"""Base Dataset abstraction for SFT training with SkyRL"""

       @property
       @abc.abstractmethod
       def sequence_lengths(self) -> List[int]:
           raise NotImplementedError

class TextDataset(SFTDataset):
    pass

class PretokenizedDataset(SFTDataset):
    pass

class ConcatSFTDataset(SFTDataset, torch.utils.data.ConcatDataset):

    @property
    def dataset_lengths(self) -> List[int]:
        return [len(dataset) for dataset in self.datasets]
    
    @property
    def sequence_lengths(self) -> List[int]:
        """Sequence lengths from all the samples in the dataset"""
        seq_lengths = []
        for ds in self.datasets:
            seq_lengths.extend(ds.sequence_lengths)
        return seq_lengths

class SFTTrainer:
    def load_dataset(self) -> SFTDataset:
        ...

    def build_train_sampler(self, dataset: SFTDataset):
        is_multi_dataset = isinstance(dataset, ConcatSFTDataset)
        ...

[SumanthRH]

Given that NormalizedTransform doesn't vectorize the processing, how much better is this batched fetch as opposed to row-wise?

Does ConcatDataset automatically use batched fetch on each dataset? Or will it default to row-wise?

[avigyabb]

You're right, for this PR (part 1) since the transform isn't vectorized the batched fetch is only 1.2x better than row-wise when the dataset is mmap'd. However, if the data is in S3, batched fetches issue much fewer S3 requests as opposed to one S3 request per row in the batch.

ConcatDataset does not automatically use batched fetch, since it didn't have the __getitems__() method. I added in my latest commit.

[SumanthRH]

@avigyabb can you fix the test failures?

https://github.com/NovaSky-AI/SkyRL/actions/runs/30581508311/job/91002847129?pr=1961

=========================== short test summary info ============================
FAILED tests/train/test_async_batch_collation.py::test_async_collation_matches_serial_default_collator - AttributeError: 'tuple' object has no attribute 'sequence_lengths'
FAILED tests/train/test_async_batch_collation.py::test_async_collation_matches_serial_packed_collator - AttributeError: 'tuple' object has no attribute 'sequence_lengths'
FAILED tests/train/test_async_batch_collation.py::test_async_collation_matches_serial_with_partial_tail - AttributeError: 'tuple' object has no attribute 'sequence_lengths'
FAILED tests/train/test_async_batch_collation.py::test_checkpoint_state_excludes_collated_ahead_batch - AttributeError: 'tuple' object has no attribute 'sequence_lengths'