[feat][SFT] Ingest pretokenized datasets from local or object storage (S3/GCS)

[avigyabb]

What

Adds a pretokenized data path to the SFT trainer: point pretokenized_dataset_paths at store(s) of already-tokenized rows — on local disk, S3, or GCS — 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) and publishes the result to an object store.

How it works

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

• Storage: local path, s3://, gs://, or gcs:// URI, resolved through the shared fsspec io layer (skyrl.backends.skyrl_train.utils.io) with its S3 retry/credential-refresh handling. Cloud downloads are cached under cache_dir keyed by remote path (atomic tmp-dir-then-rename, NFS-safe); disable_cache / force_recache are honored.
• 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 actionable 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 rather than truncated (mirrors 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).
• 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: 35 new tests in tests/train/test_sft_pretokenized.py covering format detection, schema validation/normalization, truncation, VLM pass-through + collation, S3 download/cache/force-redownload (monkeypatched io), multi-store concatenation + mixing-sampler wiring, and config validation. Full SFT suite (test_sft_config / test_sft_dataloader / test_sft_callbacks / test_max_training_steps): 170 passed.
• GPU e2e (Qwen2.5-0.5B-Instruct, FSDP on 1×L4, ingesting parquet shards from S3; alpaca tokenized offline with this repo's own tokenizer helpers):
  • Multi-store run on the final merged code — two S3 stores concatenated (180+20 rows) + pretokenized S3 eval store: https://wandb.ai/sky-posttraining-uc-berkeley/skyrl_sft/runs/gx6a6xny
  • Single-store run (fresh S3 download, full-sequence mask schema): https://wandb.ai/sky-posttraining-uc-berkeley/skyrl_sft/runs/cpsetz30
  • Loss parity: pretokenized runs reproduce the online-tokenization runs' loss curves at identical seeds (e.g. step-5 train loss 1.0803 vs 1.0813).

Notes / limitations

• The offline pipeline must apply the same chat template as the trained model — token ids can't be verified at train time.
• The download cache is keyed by remote path only; replacing a store's contents in-place requires force_recache=true once.
• HF-datasets-style split syntax does not apply to pretokenized stores (the store is the split); streaming ingestion is out of scope.
• VLM ingestion is covered at unit/collation level; no GPU VLM e2e in this PR.

🤖 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.

[gemini-code-assist]

In multi-node or multi-process distributed training environments, concurrent processes/ranks will attempt to download the same dataset at the same time. Using a deterministic temporary directory name like target_dir + ".tmp" causes a severe race condition where one process can delete or overwrite files currently being downloaded by another process. Using a unique temporary directory via tempfile.mkdtemp in the same parent directory, followed by an atomic os.rename, completely avoids this conflict and ensures safe concurrent downloads.

def _download_to(path: str, target_dir: str) -> None:
    """Download a cloud file or directory into ``target_dir``.

    A cloud directory's contents land directly in ``target_dir``; a single
    cloud file lands at ``target_dir/<basename>``. The download goes to a
    unique temporary directory in the same parent directory first and is atomically
    renamed so concurrent readers (e.g. multi-node runs sharing an NFS cache)
    never see a partial download or conflict with each other.
    """
    parent_dir = os.path.dirname(target_dir)
    os.makedirs(parent_dir, exist_ok=True)
    temp_dir = tempfile.mkdtemp(prefix=".download_tmp_", dir=parent_dir)

    try:
        if io.isdir(path):
            io.download_directory(path, temp_dir)
        else:
            io.download_file(path, os.path.join(temp_dir, os.path.basename(path.rstrip("/"))))

        try:
            os.rename(temp_dir, target_dir)
        except OSError:
            if os.path.isdir(target_dir):
                # Another process completed the download and renamed first
                shutil.rmtree(temp_dir, ignore_errors=True)
            else:
                raise
    except Exception:
        if os.path.isdir(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)
        raise

[gemini-code-assist]

When recursively collecting data files under a directory, os.walk will traverse hidden directories (such as .ipynb_checkpoints or .git) and collect hidden or temporary files (such as macOS AppleDouble ._* files). This can lead to unexpected files being loaded or causing parsing/loading errors. It is safer to explicitly skip hidden files and directories during traversal.

Suggested change

	def _collect_data_files(root: str) -> dict[str, list[str]]:
	"""Recursively collect data files under ``root``, grouped by format."""
	groups: dict[str, list[str]] = {"parquet": [], "json": [], "arrow": []}
	for dirpath, _, filenames in os.walk(root):
	for name in sorted(filenames):
	full = os.path.join(dirpath, name)
	lower = name.lower()
	if lower.endswith(_PARQUET_EXTS):
	groups["parquet"].append(full)
	elif lower.endswith(_JSON_EXTS):
	groups["json"].append(full)
	elif lower.endswith(_ARROW_EXTS):
	groups["arrow"].append(full)
	return groups
	def _collect_data_files(root: str) -> dict[str, list[str]]:
	"""Recursively collect data files under ``root``, grouped by format."""
	groups: dict[str, list[str]] = {"parquet": [], "json": [], "arrow": []}
	for dirpath, dirnames, filenames in os.walk(root):
	# Skip hidden directories in-place
	dirnames[:] = [d for d in dirnames if not d.startswith(".")]
	for name in sorted(filenames):
	# Skip hidden or temporary files
	if name.startswith("."):
	continue
	full = os.path.join(dirpath, name)
	lower = name.lower()
	if lower.endswith(_PARQUET_EXTS):
	groups["parquet"].append(full)
	elif lower.endswith(_JSON_EXTS):
	groups["json"].append(full)
	elif lower.endswith(_ARROW_EXTS):
	groups["arrow"].append(full)
	return groups