[feat][SFT] Cloud-path (S3/GCS) dataset ingestion for pretokenized and text datasets - #1933
Draft
avigyabb wants to merge 1 commit into
Draft
[feat][SFT] Cloud-path (S3/GCS) dataset ingestion for pretokenized and text datasets#1933avigyabb wants to merge 1 commit into
avigyabb wants to merge 1 commit into
Conversation
SumanthRH
pushed a commit
that referenced
this pull request
Jul 22, 2026
## 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 #1933. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Avi Basnet <avigyabb@stanford.edu> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…atasets Follow-up to NovaSky-AI#1927 (merged): factor cloud materialization into skyrl/train/dataset/remote_storage.py and use it from both SFT data paths: - pretokenized stores: pretokenized_dataset_paths / eval_pretokenized_dataset_paths accept s3:// gs:// gcs:// URIs (context-managed download; rows are materialized in memory before an uncached temp download is cleaned up). - text-format datasets: train_datasets / eval_datasets entries may be cloud URIs; they are resolved to a persistent local copy before load_dataset (persistent because parallel tokenization workers re-open the dataset by path). The original URI stays the tokenization cache key so caches are stable across nodes. Downloads cache under cache_dir keyed by remote path with atomic tmp-then-rename (NFS-safe); disable_cache / force_recache honored. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
avigyabb
force-pushed
the
pretokenized-cloud-paths
branch
from
July 22, 2026 18:57
df2d9a6 to
f0fe2d8
Compare
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
Follow-up to #1927 (now merged and rebased onto — this PR is a single commit on top of main), per review: cloud-path support was split out of the pretokenized PR so it can land once for both SFT data paths.
After this PR,
s3://,gs://, andgcs://URIs are accepted by:pretokenized_dataset_paths/eval_pretokenized_dataset_paths(pretokenized stores), andtrain_datasets/eval_datasets(raw text-format datasets fed to online tokenization).How
New shared module
skyrl/train/dataset/remote_storage.py, built on the existing fsspec io layer (skyrl.backends.skyrl_train.utils.io, with its S3 retry/credential-refresh handling):materialize_local(...)— context-managed download for the pretokenized loader (rows are materialized in memory inside the context, so an uncached temp download is cleaned up at context exit).materialize_remote_dataset(...)— persistent resolution for the text tokenization path: the parallel tokenization workers re-open the dataset by path, so the local copy must outlive the call. The original URI (not the resolved local path) remains the tokenization cache key, keeping caches stable across nodes/runs.cache_dirkeyed by remote path, with atomic tmp-dir-then-rename (NFS-safe for multi-node runs);disable_cache/force_recacheare honored. Note the cache is keyed by remote path only — replacing a store's contents in place requiresforce_recache=trueonce.Testing
materialize_remote_dataset(local passthrough, missing-path error, cache reuse). 148 passed across the pretokenized + SFT config/dataloader suites.🤖 Generated with Claude Code