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](#overview)
2. [DataPackerDataLoader](#datapackerdataloader)
   - [Step 1 — Prepare your data source](#step-1--prepare-your-data-source)
   - [Step 2 — Write your DataPacker](#step-2--write-your-datapacker)
   - [Step 3 — Wire into an experiment config](#step-3--wire-everything-into-an-experiment-config)
   - [Key parameters](#key-parameters)
   - [Shuffle and stateful checkpoint/resume](#shuffle-and-stateful-checkpointresume)
   - [Data-parallel sharding](#data-parallel-sharding)
3. [JointDataPackerDataLoader](#jointdatapackerdataloader)
   - [When to use it](#when-to-use-it)
   - [How to wire it up](#how-to-wire-it-up)
   - [Stateful checkpoint/resume](#stateful-checkpointresume)
4. [Real-world examples](#real-world-examples)
5. [Checklist](#checklist-for-a-new-dataset)

---

## 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)

```python
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)

```python
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

```python
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.

```python
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

```python
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

```python
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:

```bash
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

```python
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`:

```python
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):

```python
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.

```python
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