# How to Use Dataset Formats in Levanter

This guide walks you through configuring and using various dataset formats in Levanter.

---

## Step 1: Identify the Format for Each Dataset

Each dataset in your configuration can use a different format depending on its structure and intended use. Currently, the supported formats are:

| Format       | Best For                           | Example HF Dataset                                                                                   |
|--------------|------------------------------------|------------------------------------------------------------------------------------------------------|
| `text`       | Pretraining on plain text          | [`mlfoundations/dclm-baseline-1.0`](https://huggingface.co/datasets/mlfoundations/dclm-baseline-1.0) |
| `chat`       | Multi-turn chat models (SFT)       | [`allenai/tulu-3-sft-mixture`](https://huggingface.co/datasets/allenai/tulu-3-sft-mixture)           |
| `prebuilt`   | Pre-tokenized cache inputs         | (custom prebuilt data)                                                                               |

---

## Step 2: Prepare Your Data

All data should either be in [JSONL format](https://jsonlines.org/) — one JSON object per line –
or be a [Hugging Face dataset]](https://huggingface.co/docs/datasets/create_dataset) with appropriate field names.

### `text` format

```jsonl
{"text": "The quick brown fox jumps over the lazy dog."}
```

### `chat` format

```jsonl
{"messages": [ {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hi there!"} ]}
```
### `prebuilt` format

```jsonl
{"input_ids": [101, 2023, 2003, 1037, 7099], "loss_weights": [1, 1, 1, 1, 1]}
```

The `prebuilt` format is primarily intended for cases where you build the cache ahead of time using a custom
Zephyr pipeline or `SerialCacheWriter`. For the recommended approach and examples, see
[Direct Cache Construction](./Direct-Cache-Construction.md).

---

## Step 3: Write Your Config

> **URL vs Hugging Face dataset**
>
> - Use `train_urls:` and `validation_urls:` if your data is stored in GCS, S3, or local files.
> - Use `id:` if your dataset is hosted on Hugging Face (e.g. `allenai/tulu-3-sft-mixture`).
>   You can also specify `name:` if the dataset has named subsets.

Use either a single data source:

```yaml
data:
  train_urls: ["gs://bucket/train.jsonl.gz"]
  validation_urls: ["gs://bucket/val.jsonl.gz"]
  format:
    type: prebuilt
    input_ids_key: input_ids
    loss_weights_key: loss_weights
  tokenizer: marin-community/marin-tokenizer
  cache_dir: gs://bucket/cache
```

Or a mixture of multiple datasets:

```yaml
data:
  components:
    owt:
      source:
        type: url
        train_urls: ["gs://openwebtext/train.{1..128}.jsonl.gz"]
    alpaca:
      source:
        type: url
        train_urls: ["gs://bucket/alpaca_pretokenized.jsonl.gz"]
      format:
        type: prebuilt
        input_ids_key: input_ids
        loss_weights_key: loss_weights
    tulu:
      source:
        type: hf
        id: allenai/tulu-3-sft-mixture
      format:
        type: chat
        messages_key: messages
  train_weights:
    owt: 0.5
    alpaca: 0.3
    tulu: 0.2
  tokenizer: marin-community/marin-tokenizer
  cache_dir: gs://bucket/cache
```

!!! tip

     `train_weights` lets you mix datasets with different importance. Values are normalized.


!!! warning

    To use a chat format, your tokenizer must have a `chat_template`, or you must provide one in the config.
    This template must be formatted to work for training (which most are not, and it is not well documented in Hugging Face).
    The `marin-community/marin-tokenizer` has a default template that works. See our [chat template docs](../reference/Data-Formats.md#chat-templates) for more details.

https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L1530

### Shuffle Policies

`data.shuffle` controls how training examples are shuffled:

- `true`: full permutation shuffle (best mixing, highest random-access IO pressure)
- `false`: no shuffle
- object: hierarchical block shuffle (`BlockShuffleConfig`)
- omitted: hierarchical block shuffle with `io_block_size: 256`, `window_blocks: 512`, and `perm_type: feistel`

```yaml
data:
  # Hierarchical block shuffle (default)
  shuffle:
    io_block_size: 256
    window_blocks: 512
    perm_type: feistel
```

```yaml
data:
  # Full shuffle
  shuffle: true
```

Block shuffle is designed for TensorStore-backed training caches: it preserves more IO locality than full shuffle.
Internally it shuffles full blocks, then shuffles examples within a window of blocks.

Practical defaults:

- Set `io_block_size` to match one underlying cache storage chunk in sequence-space when possible.
- Start `window_blocks` around your train batch size in examples (e.g. batch size 512 -> `window_blocks: 512`).
- Increase `window_blocks` if you observe a persistent validation-loss gap vs full shuffle.

If the dataset length is not divisible by `io_block_size`, the final tiny block stays at the end of the dataset and is
locally permuted within itself.


### Training Mixture Schedules

Levanter supports training mixtures with different weights at different points in the training schedule.

To do so, you need to specify the transition points in the `train_weights` section:

```yaml
data:
  components:
    owt:
      source:
        type: url
        train_urls: ["gs://openwebtext/train.{1..128}.jsonl.gz"]
    alpaca:
      source:
        type: url
        train_urls: ["gs://bucket/alpaca_pretokenized.jsonl.gz"]
      format:
        type: prebuilt
        input_ids_key: input_ids
        loss_weights_key: loss_weights
    tulu:
      source:
        type: hf
        id: allenai/tulu-3-sft-mixture
      format:
        type: chat
        messages_key: messages
  train_weights:
    - [0, {"owt": 0.5, "alpaca": 0.3, "tulu": 0.2}]
    - [1000, {"owt": 0.2, "alpaca": 0.4, "tulu": 0.4}]
  tokenizer: marin-community/marin-tokenizer
```

(Again, the weights need not sum to 1.)

---

## Step 4: Launch Training

Run with:

```bash
python -m levanter.main.train_lm --config_path my_config.yaml
```

If the configured cache is missing, Levanter will first build it using the Zephyr-backed preprocessing pipeline. Training will use the cache after construction and consolidation complete.

---

## Tips

- Always validate that your JSONL is well-formed and contains all required fields.
- Each format must include a `format:` block unless you're using plain `text` with default settings.
- Use `train_weights: { mydataset: 0.0 }` to include a dataset for evaluation only.

For more details, see the [Dataset Formats Reference](../reference/Data-Formats.md).
