Custom Datasets for Generator and Reasoner Training¶

Bring your own dataset into Cosmos training with CosmosDataLoader — the OSS-facing data layer that works without any internal infrastructure (no WebDataset, no object-store credentials).

CosmosDataLoader turns any dataset into training batches by composing four small, swappable roles. Pick a built-in for each slot, or write your own — the loader wires them together in a fixed, safe order.

DataDistributor   →   RawItemProcessor   →   SampleBatcher   →   BatchCollator
(raw items, sharded     (one raw item        (a stream of         (a group of
 across DP ranks ×       → one sample         samples → groups     samples → one
 workers, shuffled,      dict)                that form a batch)   batch dict for
 resumable)                                                        model.forward)

DataDistributor owns the dataset and yields this rank/worker's disjoint slice of raw items (sharding, shuffle, checkpoint/resume).
RawItemProcessor turns one raw item into one training-ready sample dict (decode, tokenize, etc.).
SampleBatcher pulls from the sample stream and decides which samples go together in a batch (fixed size, token-budget packing, …).
BatchCollator turns a chosen group of samples into one batch dict.

Everything lives in cosmos_framework.data.generator.dataflow. The loader is a torch.utils.data.DataLoader subclass, so it drops into existing training loops.

1. Quickstart (60 seconds)¶

"I have a map-style dataset and just want normal, shuffled, resumable batches":

from cosmos_framework.data.generator.dataflow import (
    CosmosDataLoader, MapDistributor, IdentityProcessor,
)

loader = CosmosDataLoader(
    distributor=MapDistributor(my_dataset, shuffle=True, seed=0),  # any torch map Dataset
    processor=IdentityProcessor(),                                 # dataset already yields samples
    batch_size=32,                                                 # sugar → SimpleBatcher + DefaultBatchCollator
)

for batch in loader:
    out = model(**batch)

batch_size=N is convenience sugar: when you don't pass an explicit batcher, the loader builds a SimpleBatcher(N) + DefaultBatchCollator() (stock torch.utils.data stacking). Pass an explicit batcher/collator for anything fancier. (Passing both batch_size= and batcher= is an error.)

2. The four roles¶

Each role is a tiny ABC in cosmos_framework.data.generator.dataflow.base. Implement the one method (plus, for distributors, optional resume hooks).

class DataDistributor(ABC):
    def stream(self, dp_rank, dp_world_size, worker_id, num_workers) -> Iterator[Any]:
        """Yield this (rank, worker)'s disjoint slice of raw items, indefinitely."""
    def state_dict(self) -> dict: ...          # optional, for resume
    def load_state_dict(self, state) -> None: ...

class RawItemProcessor(ABC):
    def process(self, item) -> dict: ...        # one raw item → one sample dict

class SampleBatcher(ABC):
    def batches(self, samples: Iterator[dict]) -> Iterator[list[dict]]: ...
    def sample_size(self, sample) -> int: ...   # only packing batchers need this

class BatchCollator(ABC):
    def collate(self, samples: list[dict]) -> dict: ...   # group → batch dict

The loader passes the rank/worker coordinates into stream() — you never read get_worker_info() yourself. The fixed order (distribute → process → batch → collate) is enforced by the loader, so the stages can't be misordered.

Built-ins¶

Role	Built-in	Use it when
Distributor	`IterableDistributor(iterable)`	streaming / `IterableDataset` source (round-robin shard; not resumable)
	`MapDistributor(dataset, seed=0, shuffle=True, name="")`	map-style `Dataset` (per-epoch shuffle, slice shard, resumable)
	`RankPartitionedDistributor({name: {"dataset":…, "ratio":…}})`	assign whole DP ranks to different datasets by ratio
	`MixtureDistributor({name: (distributor, ratio)}, seed=0)`	mix several distributors into one stream at the sample level
Processor	`IdentityProcessor()`	the dataset already yields finished sample dicts
	(write your own)	decode/tokenize/transform a raw record
Batcher	`SimpleBatcher(batch_size, drop_last=False)`	fixed-size batches
	`PoolPackingBatcher(max_tokens, pool_size=16, max_batch_size=1, long_threshold=6400, batching_strategy="prefer_closest", apply_long_sample_halving=True, size_fn=None)`	token-budget bin-packing (reorders within a pool to minimize padding)
	`SequentialPackingBatcher(max_sequence_length, …, max_samples_per_batch=None)`	order-preserving pack-until-budget (no reordering)
Collator	`DefaultBatchCollator()`	stack with `torch.utils.data.default_collate`
	`VFMListCollator()`	VFM packed batches (media kept as per-sample lists)

Recipe-specific roles live next to their recipes, e.g. VLMProcessor / VLMCollator (in configs/base/vlm/experiment/dataflow_roles.py) and VideoPhy2Processor.

3. Recipes by use-case¶

Bring your own map-style dataset (shuffle + resume)¶

class MyImageCaptionDataset(torch.utils.data.Dataset):
    def __len__(self): return len(self.records)
    def __getitem__(self, i): return self.records[i]   # a plain dict

loader = CosmosDataLoader(
    distributor=MapDistributor(MyImageCaptionDataset(...), shuffle=True, seed=42),
    processor=MyProcessor(),     # turns the record into {"input_ids": …, "pixel_values": …}
    batch_size=8,
    num_workers=4,
)

Map-style sources are resumable (see §5).

Bring your own streaming / iterable dataset¶

hf_stream = load_dataset("some/dataset", split="train", streaming=True)
loader = CosmosDataLoader(
    distributor=IterableDistributor(hf_stream),   # round-robin shard across rank×worker
    processor=MyProcessor(),
    batch_size=8,
)

Iterable sources are not resumable (you can't random-access to fast-forward).

Token-budget packing for variable-length sequences¶

loader = CosmosDataLoader(
    distributor=IterableDistributor(stream),
    processor=MyProcessor(),                       # yields {"input_ids": Tensor[L], …}
    batcher=PoolPackingBatcher(max_tokens=16000, pool_size=16, max_batch_size=1),
    collator=MyCollator(),
)

PoolPackingBatcher.sample_size defaults to len(sample["input_ids"]); pass size_fn=lambda s: … (or subclass and override sample_size) for a custom cost.

Order-preserving sequence packing¶

batcher=SequentialPackingBatcher(
    max_sequence_length=45056,
    tokenizer_spatial_compression_factor=16,
    tokenizer_temporal_compression_factor=4,
    patch_spatial=2,
)

Packs samples in stream order until the token budget is hit (no reordering). Exactly one of max_sequence_length (token-budget mode) or max_samples_per_batch (count-only mode) must be set.

Mix multiple datasets by ratio (one pipeline)¶

distributor=MixtureDistributor(
    {"webvid": (IterableDistributor(webvid), 3.0),
     "internal": (IterableDistributor(internal), 1.0)},
    seed=0,
)

Ratio-weighted merge into a single stream — use when the datasets share one processor/batcher/collator (homogeneous join).

Interleave heterogeneous pipelines (different processors/collators)¶

from cosmos_framework.data.generator.dataflow import JointCosmosDataLoader

joint = JointCosmosDataLoader(
    dataloaders={
        "vlm": {"dataloader": vlm_loader, "ratio": 1},
        "vfm": {"dataloader": vfm_loader, "ratio": 3},
    },
    seed=42,
)

Each output batch comes from one selected inner CosmosDataLoader (ratio-weighted, seeded). Use when the joined datasets need different processing — each inner loader is a full four-role pipeline. Every yielded batch is tagged with "dataset_name". ("global_id" is reserved by the checkpoint state and cannot be used as a dataset name.)

4. Wiring into a training recipe (Hydra)¶

Recipes build the loader with LazyCall so CLI overrides work:

from cosmos_framework.utils.lazy_config import LazyCall as L

dataloader_train = L(CosmosDataLoader)(
    distributor=L(MapDistributor)(dataset=L(my_dataset_factory)(...), shuffle=True),
    processor=L(MyProcessor)(...),
    batcher=L(PoolPackingBatcher)(max_tokens="${data_setting.max_tokens}", max_batch_size=1),
    collator=L(MyCollator)(),
    num_workers=2,
)

Override from the CLI like any Hydra node, e.g. dataloader_train.batcher.max_tokens=8000. See the live recipes for full examples: pre_exp012_llava_ov (VLM), videophy2_sft_nano (videophy2), pre_exp012_llava_ov_mapstyle_dataloader (map-style resumable VLM), and vision_sft_nano_mapstyle_dataloader (VFM, alongside the legacy vision_sft_nano).

Structured-TOML launches. When you launch a VLM recipe via --sft-toml, the flat [dataloader_train] knobs max_samples_per_batch and max_sequence_length are routed onto the loader's nested batcher (dataloader_train.batcher.max_batch_size and …batcher.max_tokens) by PATH_REMAPS["vlm"] in configs/toml_config/toml_config_helper.py. This only works for experiments whose batcher actually has those fields (e.g. PoolPackingBatcher).

5. Checkpoint / resume¶

Resume is handled by CosmosDataLoaderStateCallback:

from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback
cb = CosmosDataLoaderStateCallback()

Use a MapDistributor source. On save, the callback records each worker's (epoch, index) from the per-batch sample_worker_id/sample_epoch/sample_index tensors the loader stamps. On load, it sets COSMOS_DL_STATE_* env vars before workers fork; MapDistributor.stream reads them and fast-forwards to the exact next sample — no duplicated or skipped samples.
Iterable sources are not resumable (no position to fast-forward to); the stream restarts from the beginning.
For multiple loaders sharing a process (e.g. inside JointCosmosDataLoader), give each a distinct name= so resume env vars are namespaced (COSMOS_DL_STATE_{name}_WORKER_{id}_{EPOCH,INDEX}), and use a single JointCosmosDataLoaderStateCallback(outer_loader=joint_loader) instead of one CosmosDataLoaderStateCallback per inner loader.
Use ckpt_type=dcp (the default) — not ckpt_type=dummy, which disables all checkpointing. The on-disk checkpoint format is unchanged.

Validated: a live save→stop→resume on pre_exp012_llava_ov_mapstyle_dataloader (8 dp ranks, save_iter=100) reproduces the original run's per-rank input_ids shapes exactly across the resume boundary — no duplicated or skipped samples on any rank.

6. Distributed & sharding¶

The loader resolves the data-parallel coordinates as parallel_dims.dp_coord > torch.distributed > (0, 1). For FSDP+TP/PP, pass parallel_dims= so sharding uses the correct DP rank (not the global rank).
IterableDistributor/MapDistributor give each (dp_rank, worker_id) pair a disjoint, complete slice: stream i is taken iff i % (dp_world_size × num_workers) == dp_rank × num_workers + worker_id.
RankPartitionedDistributor instead assigns whole ranks to datasets by ratio and sets shard_world_size/shard_rank on the chosen dataset (which then self-shards across the ranks sharing it). If your dataset already shards internally, disable that (dataset.shard_world_size = 1; dataset.shard_rank = 0) before handing it to a per-(rank,worker)-sharding distributor.
A MapDistributor source with num_workers > 0 is automatically promoted to persistent_workers=True (required for correct stateful resume). The fork start method (the Linux/CUDA default) is required; spawn is not supported.

7. Troubleshooting / FAQ¶

OOM on large packed batches. PoolPackingBatcher(apply_long_sample_halving=True) (default) halves the token budget for any batch whose largest sample ≥ 1000 tokens. Set False only after validating memory headroom at the literal budget.
ValueError: Provide either a batcher= or a batch_size=. You passed neither; give one. Passing both is also an error.
ValueError: Map-style resume cannot safely stamp a multi-sample batch …. A reordering batcher (pool packing) with batch_size > 1 on a resumable MapDistributor can't record a gap-free resume position. Use max_batch_size=1 with pool packing, a sequential (order-preserving) batcher, or an iterable (non-resumable) source.
'int' object is not iterable / wrong tensor shapes in the model. Your BatchCollator is producing a different batch structure than the model expects. Match the structure the model consumes (for VFM, that's VFMListCollator, which keeps media as per-sample lists).
Oversized samples silently dropped (sequential packing). A single sample larger than max_sequence_length is discarded with a logged error — increase the budget or filter upstream.
num_workers / persistent_workers. persistent_workers=True is ignored (with a log) when num_workers=0.

8. End-to-end worked example (custom dataset → training)¶

A local image-caption folder, fully custom processor, normal batching:

import torch
from cosmos_framework.data.generator.dataflow import (
    CosmosDataLoader, MapDistributor, RawItemProcessor, DefaultBatchCollator, SimpleBatcher,
)

class ImageCaptionFolder(torch.utils.data.Dataset):
    def __init__(self, records): self.records = records          # [{"image_path":…, "caption":…}, …]
    def __len__(self): return len(self.records)
    def __getitem__(self, i): return self.records[i]

class ImageCaptionProcessor(RawItemProcessor):
    def __init__(self, tokenizer, image_loader):
        self.tokenizer, self.image_loader = tokenizer, image_loader
    def process(self, item):
        return {
            "pixel_values": self.image_loader(item["image_path"]),   # Tensor[C,H,W]
            "input_ids": self.tokenizer(item["caption"]),            # Tensor[L]
        }

loader = CosmosDataLoader(
    distributor=MapDistributor(ImageCaptionFolder(records), shuffle=True, seed=0),
    processor=ImageCaptionProcessor(tokenizer, image_loader),
    batcher=SimpleBatcher(batch_size=16),
    collator=DefaultBatchCollator(),
    num_workers=4,
)

for batch in loader:                  # batch = {"pixel_values": [16,C,H,W], "input_ids": [16,L]}
    loss = model(**batch)
    loss.backward()

To pack variable-length captions by token budget instead of fixed size, swap the batcher for PoolPackingBatcher(max_tokens=…, max_batch_size=1) and provide a collator that pads/stacks accordingly — nothing else changes.

9. Real-world examples¶

Reasoner (VLM) — HuggingFace image-text dataset, streaming¶

File: cosmos_framework/configs/base/reasoner/experiment/llava_ov_vlm.py (pre_exp012_llava_ov)

distributor: IterableDistributor(get_llava_ov_streaming(...))   # lmms-lab/LLaVA-OneVision-Data
processor:   VLMProcessor   (ShareGPT → OpenAI messages → Qwen3-VL processor)
batcher:     PoolPackingBatcher(max_tokens≈16000, max_batch_size=1)
collator:    VLMCollator

Streaming source → not resumable. For a resumable variant of the same recipe, see llava_ov_mapstyle_dataloader_experiment.py (pre_exp012_llava_ov_mapstyle_dataloader): it loads the subset as a real map-style Dataset (load_dataset(..., streaming=False)) and wraps it in a MapDistributor, so checkpoint/resume works (see §5).

Reasoner (VLM) — local video dialog dataset¶

File: cosmos_framework/configs/base/reasoner/experiment/videophy2_sft_nano.py (videophy2_sft_nano)

distributor: IterableDistributor(build_videophy2_local_dataset(...))
processor:   VideoPhy2Processor
batcher:     PoolPackingBatcher(max_tokens≈16000, max_batch_size=1)
collator:    VLMCollator

Generator (VFM) — Cosmos video SFT¶

File: the vision_sft_nano_mapstyle_dataloader experiment (the new-loader VFM variant, alongside the legacy vision_sft_nano).

distributor: IterableDistributor over the Cosmos video dataset
processor:   recipe processor (decode + tokenize)
batcher:     SequentialPackingBatcher(max_sequence_length=…)   # order-preserving token packing
collator:    VFMListCollator                                   # media kept as per-sample lists

10. Checklist for a new dataset¶

Single dataset (`CosmosDataLoader`)¶

[ ] Pick a distributor: MapDistributor (map-style Dataset, shuffle + resume) or IterableDistributor (streaming, not resumable).
[ ] Write a RawItemProcessor (or use IdentityProcessor if your dataset already yields finished sample dicts).
[ ] Pick a batcher: batch_size=N sugar, SimpleBatcher, PoolPackingBatcher (token-budget, reorders), or SequentialPackingBatcher (token-budget, order-preserving).
[ ] Pick a collator: DefaultBatchCollator, VFMListCollator, or your own (must match the structure the model consumes).
[ ] For real resume: use a MapDistributor, add CosmosDataLoaderStateCallback(), and ckpt_type=dcp (not dummy).
[ ] For FSDP+TP/PP, pass parallel_dims= so the correct DP rank is used.
[ ] Register the experiment in the Hydra ConfigStore (cs.store(group="experiment", …)).
[ ] Smoke-test with --dryrun (config build) then trainer.max_iter=10 before a full run.

Multiple datasets (`JointCosmosDataLoader`)¶

[ ] Build each inner pipeline as its own CosmosDataLoader; give each a unique name= matching its key in dataloaders (namespaces resume env vars).
[ ] Set each dataset's ratio (controls how often it is visited, per batch).
[ ] Use a single JointCosmosDataLoaderStateCallback(outer_loader=joint_loader) — do not also register a standalone CosmosDataLoaderStateCallback per inner loader.
[ ] Avoid "global_id" as a dataset name (reserved by the checkpoint state).
[ ] Use ckpt_type=dcp for real checkpoint/resume.

Reference: where things live¶

ABCs + built-ins: cosmos_framework/data/generator/dataflow/ (base.py, distributors.py, batchers.py, collators.py, processors.py, loader.py).
Public symbols are re-exported from cosmos_framework.data.generator.dataflow.
Live recipes using the loader: pre_exp012_llava_ov, pre_exp012_llava_ov_mapstyle_dataloader, videophy2_sft_nano, and vision_sft_nano_mapstyle_dataloader.