cosmos-framework/docs/custom_dataset.md

Custom Datasets for Generator and Reasoner Training

This guide explains how to bring your own dataset into Cosmos training using DataPackerDataLoader and JointDataPackerDataLoader — the OSS-facing data layer that works without any internal infrastructure (no WebDataset, no object-store credentials).

---

Contents

1. Overview
2. DataPackerDataLoader
   • Step 1 — Prepare your data source
   • Step 2 — Write your DataPacker
   • Step 3 — Wire into an experiment config
   • Key parameters
   • Shuffle and stateful checkpoint/resume
   • Data-parallel sharding
3. JointDataPackerDataLoader
   • When to use it
   • How to wire it up
   • Stateful checkpoint/resume
4. Real-world examples
5. Checklist

---

Overview

The data pipeline has two parts you control:

Your dataset (Dataset or IterableDataset)
        │
        ▼
┌────────────────────────────────────────────────┐
│           DataPackerDataLoader                 │
│                                                │
│  map-style Dataset (any shuffle setting):      │
│  ┌──────────────────────────────────────────┐  │
│  │    _ShuffledMapIterableDataset           │  │
│  │  • per-epoch randperm (shuffle=True)     │  │
│  │    or sequential (shuffle=False)         │  │
│  │  • DP × worker sharding                  │  │
│  │  • position metadata for stateful resume │  │
│  └──────────────────┬───────────────────────┘  │
│                     │                          │
│  IterableDataset:                              │
│  ┌──────────────────────────────────────────┐  │
│  │    _IterableWrapper                      │  │
│  │  • DP × worker sharding only             │  │
│  │  • no stateful resume                    │  │
│  └──────────────────┬───────────────────────┘  │
│                     │ raw item                 │
│  ┌──────────────────▼───────────────────────┐  │
│  │    _DataPackerIterableDataset            │  │
│  │    (subclass of PackingIterableDataset)  │  │
│  │                                          │  │
│  │  • fill pool (pool_size samples)         │  │
│  │  • greedy bin-pack within max_tokens     │  │
│  │  • cap at max_batch_size                 │  │
│  │                                          │  │
│  │  → DataPacker.sft_process_sample  ← you  │  │
│  │  → DataPacker.compute_num_tokens  ← you  │  │
│  │  → DataPacker.sft_collate_fn      ← you  │  │
│  └──────────────────────────────────────────┘  │
└────────────────────────────────────────────────┘
        │ fully-collated batch dict
        ▼
     Trainer / model.forward()

Key point: all map-style datasets (whether shuffle=True or shuffle=False) are routed through _ShuffledMapIterableDataset, which attaches position metadata to every sample. This means stateful checkpoint/resume works regardless of whether shuffle is enabled.

---

DataPackerDataLoader

Step 1 — Prepare your data source

DataPackerDataLoader accepts either a map-style torch.utils.data.Dataset or an iterable-style torch.utils.data.IterableDataset. Plain lists and generators are rejected with a TypeError.

Type	Notes
torch.utils.data.Dataset (map-style)	Pass directly. Supports shuffle=True/False and stateful checkpoint/resume.
torch.utils.data.IterableDataset	Pass directly. No shuffle, no stateful resume — shuffle externally if needed.
HuggingFace Dataset	Is a torch.utils.data.Dataset subclass — pass directly, shuffle=True works.
HuggingFace IterableDataset (streaming)	Is a torch.utils.data.IterableDataset — pass directly, use .shuffle() externally.

Loading from HuggingFace (simplest)

from cosmos_framework.data.vfm.data_packer_dataloader import load_data_source

# HuggingFace Hub dataset (downloaded, map-style)
data_source = load_data_source("liuhaotian/LLaVA-Instruct-150K", split="train")

# Dataset saved with dataset.save_to_disk()
data_source = load_data_source("/path/to/my_saved_dataset", split="train")

# Then pass with shuffle for per-epoch shuffling + stateful resume
DataPackerDataLoader(data_source=data_source, ..., shuffle=True, seed=42)

Streaming from HuggingFace (no disk space)

from datasets import load_dataset

data_source = load_dataset(
    "lmms-lab/LLaVA-OneVision-Data", name="si", split="train", streaming=True
)
# shuffle before passing — IterableDataset does not support internal shuffle
data_source = data_source.shuffle(seed=42, buffer_size=10_000)

Custom map-style dataset

class MyMapDataset(torch.utils.data.Dataset):
    def __len__(self): return 10_000
    def __getitem__(self, idx): return {"video": ..., "text": ...}

# Pass directly — DataPackerDataLoader handles sharding and shuffle internally
DataPackerDataLoader(data_source=MyMapDataset(), ..., shuffle=True, seed=42)

---

Step 2 — Write your DataPacker

DataPacker is an abstract base class. Implement all three methods, then place the class in the same experiment config file that uses it.

from cosmos_framework.data.vfm.data_packer import DataPacker

class MyDataPacker(DataPacker):

    def sft_process_sample(self, item: dict) -> dict:
        """
        Convert one raw item from data_source into a training-ready sample.
        Called inside DataLoader workers — tokenization, decoding, transforms go here.
        """
        ...
        return {"input_ids": ..., "labels": ..., ...}

    def compute_num_tokens(self, sample: dict) -> int:
        """
        Return the token cost of one processed sample.
        Used by the packing engine to decide how many samples fit in a batch.
        """
        return int(sample["input_ids"].shape[0])

    def sft_collate_fn(self, samples: list[dict], max_len: int,
                       ignore_label_id: int = -100) -> dict:
        """
        Collate a list of processed samples into one batch dict.
        max_len is the longest token sequence in this batch (for padding).
        """
        ...
        return {"input_ids": ..., "labels": ..., ...}

Note on extra batch keys: For map-style datasets, DataPackerDataLoader automatically appends sample_worker_id, sample_epoch, and sample_index to every batch dict. These are used by DataLoaderStateCallback for stateful checkpoint/resume and are transparent to the model as long as training_step accesses the batch by key (not **kwargs unpack).

Token counting for Generator models

import math

def compute_num_tokens(self, sample: dict) -> int:
    tokens = 1 + len(sample.get("text_token_ids", []))
    v = sample.get("video")   # shape [C, T, H, W]
    if v is not None:
        _, T, H, W = v.shape
        latent_h = math.ceil(H / (self.spatial_compression * self.patch_spatial))
        latent_w = math.ceil(W / (self.spatial_compression * self.patch_spatial))
        latent_t = 1 + (T - 1) // self.temporal_compression
        tokens += latent_h * latent_w * latent_t + 2
    return tokens

Typical values: spatial_compression=16, temporal_compression=4, patch_spatial=2.

---

Step 3 — Wire everything into an experiment config

from cosmos_framework.utils.lazy_config import LazyCall as L, LazyDict
from cosmos_framework.data.vfm.data_packer_dataloader import DataPackerDataLoader, load_data_source
from cosmos_framework.callbacks.dataloader_state import DataLoaderStateCallback
from hydra.core.config_store import ConfigStore

cs = ConfigStore.instance()

my_experiment = LazyDict(dict(
    defaults=[...],   # inherit model, optimizer, scheduler from a base

    trainer=dict(
        callbacks=dict(
            # Tracks per-worker (epoch, position) for checkpoint/resume.
            # Works with both shuffle=True and shuffle=False for map-style datasets.
            dataloader_state=L(DataLoaderStateCallback)(distributor_type="data_packer"),
        ),
    ),

    dataloader_train=L(DataPackerDataLoader)(
        data_source=L(load_data_source)(name="my-org/my-dataset", split="train"),
        data_packer=L(MyDataPacker)(...),
        max_tokens=16000,
        pool_size=16,
        max_batch_size=1,
        shuffle=True,            # per-epoch randperm, different order every epoch
        seed=42,                 # epoch e uses seed+e → reproducible permutations
        num_workers=4,
        prefetch_factor=4,
        persistent_workers=True,
        pin_memory=True,
    ),
    dataloader_val=None,
), flags={"allow_objects": True})

cs.store(group="experiment", package="_global_", name="my_experiment", node=my_experiment)

Launch:

torchrun --nproc_per_node=8 -m cosmos_framework.scripts.train \
    --config=cosmos_framework/configs/base/config.py -- \
    experiment=my_experiment \
    trainer.max_iter=1000

---

Key parameters

Parameter	What it controls
max_tokens	Token budget per batch. Packing stops when adding one more sample would exceed this. For Generator, counts video latent tokens; for Reasoner, counts input_ids length.
pool_size	Samples to buffer before bin-packing. Larger pool → better packing efficiency, more memory. Default: 16.
max_batch_size	Hard cap on samples per batch regardless of token budget. Use 1 for Reasoner (one image per step), 128–256 for action policy training.
shuffle	True → per-epoch randperm shuffle for map-style datasets (no effect on IterableDataset, a warning is logged). False → sequential, still resumable.
seed	Base seed for the shuffle permutation. Epoch e uses seed + e → reproducible, different ordering every epoch. Default: 0.
name	Optional string that namespaces resume env vars. Required when multiple DataPackerDataLoader instances share the same process (i.e., inside JointDataPackerDataLoader). Each inner loader must have a unique name matching its key in the dataloaders dict. Leave empty (default) for single-loader setups.
long_threshold	Samples with token count ≥ this are emitted as singleton batches, bypassing packing. Default: 6400.
batching_strategy	"prefer_closest" (default) picks candidates nearest in token length. "prefer_first" picks the first that fits.
num_workers	DataLoader workers for sft_process_sample. Use 0 for debugging.
persistent_workers	Automatically promoted to True for all map-style datasets when num_workers > 0 (required for correct resume behaviour).

---

Shuffle and stateful checkpoint/resume

For map-style datasets, DataPackerDataLoader tracks each worker's position and resumes training from exactly where it left off after a checkpoint. This works for both shuffle=True and shuffle=False.

How it works

1. Each epoch, a permutation is generated with torch.randperm(n, generator=torch.Generator().manual_seed(seed + epoch)) (or list(range(n)) when shuffle=False).
2. Each (dp_rank, worker_id) pair sees a disjoint stride: perm[stream_id :: total_streams] where stream_id = dp_rank * num_workers + worker_id.
3. After each training step, DataLoaderStateCallback reads sample_epoch and sample_index from the batch and tracks the high-water mark per worker.
4. At checkpoint, the DCP checkpointer saves the state to iter_XXXXXXXXX/dataloader/rank_{rank}.pkl.
5. On resume, load_state_dict sets DP_STATE_WORKER_{worker_id}_EPOCH/INDEX env vars before workers start, and workers fast-forward past already-seen samples.

At most pool_size (default 16) samples are re-processed at each resume (they pass through sft_process_sample again but are trained on only once).

Required wiring

from cosmos_framework.callbacks.dataloader_state import DataLoaderStateCallback

exp["trainer"]["callbacks"]["dataloader_state"] = L(DataLoaderStateCallback)(
    distributor_type="data_packer"
)

Use ckpt_type=dcp (the default) — not ckpt_type=dummy which disables all checkpointing.

Limitations

• Map-style datasets only. Stateful resume is not supported for IterableDataset sources.
• fork start method required (the default for Linux/CUDA). spawn is not supported.
• persistent_workers=True required when num_workers > 0 (auto-enforced for all map-style datasets).

---

Data-parallel sharding

DataPackerDataLoader automatically shards data_source across ranks and DataLoader workers. Each (dp_rank, worker_id) pair receives a disjoint subset — a strided slice of the (shuffled) permutation.

If your dataset already shards internally (like SFTDataset), disable its sharding before passing it to DataPackerDataLoader:

def get_my_dataset_no_dp(**kwargs):
    dataset = MyDataset(**kwargs)
    dataset.shard_world_size = 1   # disable internal sharding
    dataset.shard_rank = 0
    return dataset

For FSDP + TP/PP: pass parallel_dims so the correct DP rank is used (global rank ≠ DP rank in these setups):

DataPackerDataLoader(..., parallel_dims=parallel_dims)

---

JointDataPackerDataLoader

When to use it

JointDataPackerDataLoader wraps multiple DataPackerDataLoader instances with ratio-based seeded selection. Use it when training on multiple datasets with different modalities or formats — for example, video + action data at a 3:1 ratio.

Semantics mirror IterativeJointDataLoader:

• One batch = one dataset — samples from different datasets never share a packed batch.
• Ratios control how frequently each dataset is visited (per batch, not per sample).
• Selection is deterministic: step i always picks the same dataset given the same seed.
• Stateful checkpoint/resume: both the outer step counter (global_id) and each inner loader's per-worker position are saved and restored.

How to wire it up

Each inner DataPackerDataLoader must be given a unique name that matches its key in the dataloaders dict. The name namespaces the resume env vars to prevent conflicts between concurrent loaders.

from cosmos_framework.data.vfm.data_packer_dataloader import DataPackerDataLoader, JointDataPackerDataLoader
from cosmos_framework.callbacks.dataloader_state import JointDataLoaderStateCallback
from cosmos_framework.utils.lazy_config import LazyCall as L

# Build the joint loader
joint_loader = JointDataPackerDataLoader(
    dataloaders={
        "video": {
            "dataloader": DataPackerDataLoader(
                data_source=MyVideoDataset(...),
                data_packer=MyVideoDataPacker(...),
                max_tokens=45056,
                shuffle=True,
                seed=0,
                name="video",          # must match the key above
                num_workers=4,
                persistent_workers=True,
                pin_memory=True,
            ),
            "ratio": 3,                # video 3×, action 1×
        },
        "action": {
            "dataloader": DataPackerDataLoader(
                data_source=MyActionDataset(...),
                data_packer=MyActionDataPacker(...),
                max_tokens=999_999,
                max_batch_size=128,
                shuffle=True,
                seed=0,
                name="action",         # must match the key above
                num_workers=4,
                persistent_workers=True,
                pin_memory=True,
            ),
            "ratio": 1,
        },
    },
    seed=42,   # controls outer dataset selection sequence
)

# Wire into the experiment config
exp["dataloader_train"] = joint_loader
exp["trainer"]["callbacks"]["dataloader_state"] = JointDataLoaderStateCallback(
    outer_loader=joint_loader,
    distributor_type="data_packer",
)

Reserved name: "global_id" cannot be used as a dataset name — it is reserved by the checkpoint state format.

JointDataPackerDataLoader parameters

Parameter	What it controls
dataloaders	Dict mapping dataset name → {"dataloader": DataPackerDataLoader, "ratio": int}. Entries with ratio <= 0 are skipped.
seed	Base seed for outer dataset selection. Step i uses np.random.RandomState(seed + i) → same sequence on every rank.

JointDataLoaderStateCallback

This single callback replaces the per-inner-loader DataLoaderStateCallback instances. It saves:

• global_id — the outer step counter, which determines which dataset fires at each step on resume.
• Per-dataset, per-worker (epoch, index) — each inner loader's position.

All state is written to a single DCP checkpoint entry (checkpoint_component="dataloader").

Stateful checkpoint/resume

At checkpoint step N:

• global_id = N is saved.
• Each inner loader saves its per-worker (epoch, index) under its name key.

On resume:

1. JointDataLoaderStateCallback.load_state_dict calls set_start_iteration(N) on the outer loader → selection sequence resumes from step N.
2. Each inner DataLoaderStateCallback.load_state_dict sets namespaced env vars (DP_STATE_{name}_WORKER_{id}_EPOCH/INDEX) → workers fast-forward to the saved position.

Inner loader iterators are created lazily on the first __iter__ call (not at __init__ time), ensuring workers fork after env vars have been set.

---

Real-world examples

Reasoner — HuggingFace image-text dataset

File: cosmos_framework/configs/base/vlm/experiment/llava_ov_datapacker_experiment.py

data_source:  lmms-lab/LLaVA-OneVision-Data  (streaming IterableDataset)
DataPacker:   VLMDataPacker
  sft_process_sample:  ShareGPT → OpenAI messages → Qwen3-VL processor
  compute_num_tokens:  len(input_ids)
  sft_collate_fn:      unsqueeze batch dim, keep pixel_values flat
max_batch_size: 1
max_tokens:    ~16000
shuffle:       False  (streaming IterableDataset — use .shuffle() externally)

Action Policy — Robot learning (LIBERO)

File: cosmos_framework/configs/base/experiment/action/posttrain_config/libero_policy_datapacker_experiment.py

data_source:  LIBERODataset  (map-style Dataset, passed directly)
DataPacker:   ActionDataPacker
  sft_process_sample:  full ActionTransformPipeline (resize, tokenize, pad action)
  compute_num_tokens:  VAE video tokens + text tokens
  sft_collate_fn:      action/domain_id/sequence_plan fields + video + text
max_batch_size: 128   (token budget disabled — batch bounded by max_batch_size)
max_tokens:    999999
shuffle:       True, seed=0

LIBERODataset is a map-style Dataset passed directly. shuffle=True enables per-epoch shuffling and stateful checkpoint/resume. This pattern (high max_tokens

• bounded max_batch_size) is standard for action policy training where you want a fixed number of demonstrations per step.

---

Checklist for a new dataset

Single dataset (DataPackerDataLoader)

• Choose a data_source: map-style Dataset or IterableDataset (no plain lists/generators)
• For map-style: pass directly; use shuffle=True, seed=<N> for per-epoch shuffle
• For iterable: shuffle externally before passing (e.g. .shuffle(buffer_size=N))
• If dataset has internal DP sharding, disable it (shard_world_size=1)
• Subclass DataPacker and implement sft_process_sample, compute_num_tokens, sft_collate_fn
• Choose max_tokens and max_batch_size for your modality
• Add DataLoaderStateCallback(distributor_type="data_packer") to the experiment's callbacks (works for both shuffle=True and shuffle=False on map-style datasets)
• Use ckpt_type=dcp (not dummy) for real checkpoint/resume
• Register in Hydra ConfigStore with cs.store(group="experiment", ...)
• Smoke-test with ckpt_type=dummy trainer.max_iter=10 before a full run

Multiple datasets (JointDataPackerDataLoader)

• Give each inner DataPackerDataLoader a unique name matching its key in dataloaders
• Set appropriate ratio for each dataset (controls visit frequency per batch)
• Use JointDataLoaderStateCallback(outer_loader=joint_loader) instead of DataLoaderStateCallback
• Do not also register standalone DataLoaderStateCallback for inner loaders — JointDataLoaderStateCallback handles all of them
• Avoid using "global_id" as a dataset name (reserved)
• Use ckpt_type=dcp for real checkpoint/resume