[feat][SFT] Ingest pretokenized datasets

[avigyabb]

What

Adds a pretokenized data path to the SFT trainer: point pretokenized_dataset_paths at local store(s) of already-tokenized rows and train directly, skipping online tokenization (tokenize_chat_example / tokenize_sft_example) entirely. Useful when a data pipeline tokenizes offline (e.g. on a Spark/Ray data cluster).

Scope note (per review): this PR supports local paths only. Cloud-path (S3/GCS) ingestion is split into the follow-up draft #1933, where it lands once for both pretokenized stores and text-format datasets.

How it works

New module skyrl/train/dataset/pretokenized.py — load_from_pretokenized(path, max_length), the pretokenized counterpart of SFTTrainer._load_and_tokenize (same list[dict] return shape, so everything downstream — collators, sequence packing, samplers, checkpoint/resume — works unchanged):

• Formats (auto-detected, single file or directory of shards): Parquet, JSON-lines, raw Arrow IPC, or a HF Dataset.save_to_disk directory.
• Row schema: unpadded input_ids plus a full-sequence 0/1 loss_mask (same length). num_actions is inferred from the first nonzero mask entry; window-form masks, num_actions columns, and HF-style labels are rejected with clear errors. attention_mask is optional and must be all-ones (padding stays collation-internal). The mask form covers instruction-following (1s on the response) and multi-turn conversational data (1s on every assistant turn).
• VLM ingestion: rows carrying pixel_values + image_grid_thw pass through to the collator's TensorList path. Over-length VLM rows are dropped with a warning rather than truncated (consistent with the online VLM path); mixed text+VLM stores handle the null image columns parquet materializes on text rows.
• max_length truncation mirrors the online path (prompt prefix kept, action window shrinks, empty-loss rows dropped).

Config (SFTConfig) — integrated with multi-dataset SFT (#1883):

• pretokenized_dataset_paths: List[str] — multiple stores are concatenated and mixed per train_dataset_weights via DataMixingSampler, exactly like train_datasets. Exclusive with train_datasets/train_dataset_splits (explicit error, no silent precedence): a pretokenized store is the output of a preprocessing job, so the user is assumed to have pre-split/filtered it — HF split syntax does not apply.
• eval_pretokenized_dataset_paths: List[str] — each store becomes one eval set logged under eval/{name}/, named by eval_dataset_names (default: path basenames; collisions error).
• eval_interval / eval_before_train accept pretokenized eval stores.

Testing

• CPU: 32 tests in tests/train/test_sft_pretokenized.py covering format detection, schema validation/normalization, truncation, VLM pass-through + collation, multi-store concatenation + mixing-sampler wiring, and config validation. Full SFT suite passes alongside.
• GPU e2e (Qwen2.5-0.5B-Instruct, FSDP on 1×L4, alpaca tokenized offline with this repo's own tokenizer helpers into parquet, ingested from a local directory of shards + a separate local eval store): https://wandb.ai/sky-posttraining-uc-berkeley/skyrl_sft/runs/lkpv0ip1
• Loss parity: pretokenized runs reproduce the online-tokenization runs' loss curves at identical seeds (e.g. step-5 train loss 1.084 vs 1.081).

Notes / limitations

• The offline pipeline must apply the same chat template as the trained model — token ids can't be verified at train time.
• Streaming ingestion is out of scope (non-requirement).
• VLM ingestion is covered at unit/collation level; no GPU VLM e2e in this PR.
• Cloud paths (S3/GCS): follow-up draft [feat][SFT] Cloud-path (S3/GCS) dataset ingestion for pretokenized and text datasets #1933.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for pretokenized datasets (local or S3/GCS) in SFT training, allowing users to bypass online tokenization by pointing directly to pretokenized stores in Parquet, JSONL, Arrow, or HuggingFace formats. The changes include a new dataset loading module, updated configuration options in SFTConfig, integration in SFTTrainer, and comprehensive unit tests. The review feedback highlights two important improvements: addressing a potential race condition in multi-node environments during concurrent dataset downloads by using unique temporary directories, and skipping hidden files and directories during dataset file collection to prevent loading errors.

[SumanthRH]

Overall comments:

1. Adding support for cloud paths can be a separate PR - we can actually support this for text datasets and pretokenized datasets in that PR.
2. The implementation for handling pretokenized dataset seems good! There are can be more tests for handling different formats.

[SumanthRH]

Nit: Why should we have special handling for a labels field here?

[SumanthRH]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for loading pretokenized SFT datasets offline, skipping online tokenization entirely. It adds new configuration options (pretokenized_dataset_paths and eval_pretokenized_dataset_paths), a new pretokenized.py dataset module to handle loading and normalizing various on-disk formats (Parquet, JSONL, Arrow, HF Dataset), and integrates this loading mechanism into the SFTTrainer. Comprehensive tests are also added. The review feedback highlights two key improvements: first, validating that loss_mask values are strictly numeric 0 or 1 before converting to integers to prevent silent truncation of fractional values; second, changing the individual over-length VLM row warning to a debug log to avoid flooding the console during dataset loading.