Zephon

Zephon is a data loader for foundation model training, built by DatologyAI. It tokenizes, packs and mixes data while training runs, and it gives the same batches when you change the number of GPUs (elastic determinism). Zephon reads datasets directly from the Hugging Face Hub with hf:// paths.

Installation

The hf extra adds support for Hub datasets and Parquet files. The examples below also need transformers for tokenization:

pip install "zephon[hf]" transformers

Load a dataset from the Hub

Pass an hf:// path to Dataset.from_path. The path names a dataset, a configuration and a split, for example the train split of the sample-10BT configuration of HuggingFaceFW/fineweb-edu. Zephon reads Hub datasets through its shard cache, so you must enable the cache. Put the cache on fast local disk and set a size limit:

from zephon import Pipeline
from zephon.io import CacheOptions, Dataset, StoreOptions
from zephon.work import MixtureSpec, StaticMixtureWorkSource

dataset = Dataset.from_path("fineweb_edu", "hf://HuggingFaceFW/fineweb-edu/sample-10BT/train")
work_source = StaticMixtureWorkSource(
    datasets=[dataset],
    mixture=MixtureSpec({"fineweb_edu": 1.0}),
    seed=42,
)
cache = CacheOptions(enabled=True, root="/tmp/zephon-cache", limit_bytes=50 << 30)  # 50 GiB

pipeline = (
    Pipeline(work_source)
    .options(io_options=StoreOptions(cache=cache))
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=2049, algorithm="wrap")
    .batch(microbatch_size=8)
)

for batch in pipeline:
    tensors = batch.to_training(return_labels=True)
    print(tensors["input_ids"].shape, tensors["labels"].shape)  # (8, 2048) (8, 2048)
    break

tokenize converts the text column to token IDs. pack_flat joins the tokenized documents into sequences of max_length tokens. Set max_length to the context length plus one: with return_labels=True, to_training shifts the labels by one position, so input_ids and labels both have the context length.

A Pipeline is a Python iterable, so you can use it in any training loop. Zephon downloads a complete shard before it reads from it, so the first batch can take some time.

To read private or gated datasets, log in with your Hugging Face account first, for example with hf auth login. Zephon also reads the HF_TOKEN environment variable.

Path format

Zephon’s hf:// paths for Hub datasets have this format:

hf://{org}/{name}[@[{revision}][~{source}]]/[{config}/]{split}

• revision: a branch, tag or commit. The default is main. Zephon pins the path to a commit when it loads the dataset, so all ranks read the same files.
• source: original reads the files uploaded to the repository, and parquet reads the Hub’s automatic Parquet conversion. Put the source after the @, for example @~parquet. By default, Zephon reads the uploaded files when it supports their format, and uses the Parquet conversion if it doesn’t.
• config: the dataset configuration. Include it when the dataset has more than one configuration.

Zephon paths name a dataset and a split, not a file. They don’t use the hf://datasets/... file paths that pandas, Polars and DuckDB use. Use hf://HuggingFaceFW/fineweb-edu/sample-10BT/train, not hf://datasets/HuggingFaceFW/fineweb-edu/sample-10BT/train.

Mix datasets

To train on several Hub datasets, add them to the work source with the proportion of each one. This example mixes FineWeb-Edu with English Wikipedia, targeting a 70/30 mixture by tokens:

from zephon import Pipeline
from zephon.io import CacheOptions, Dataset, StoreOptions
from zephon.work import MixtureSpec, StaticMixtureWorkSource, TokenEstimation

work_source = StaticMixtureWorkSource(
    datasets=[
        Dataset.from_path("fineweb_edu", "hf://HuggingFaceFW/fineweb-edu/sample-10BT/train"),
        Dataset.from_path("wikipedia", "hf://wikimedia/wikipedia/20231101.en/train"),
    ],
    mixture=MixtureSpec({"fineweb_edu": 0.7, "wikipedia": 0.3}),
    seed=42,
    token_estimation=TokenEstimation(primer="measure", calibration_samples=2048, calibration_shards_max=4),
)
cache = CacheOptions(enabled=True, root="/tmp/zephon-cache", limit_bytes=50 << 30)  # 50 GiB

pipeline = (
    Pipeline(work_source)
    .options(io_options=StoreOptions(cache=cache))
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=2049, algorithm="wrap")
    .batch(microbatch_size=8)
)

To mix by samples instead, leave out token_estimation.

Learn more

• Zephon documentation
• Launch blog post
• Zephon paper
• DatologyAI on the Hub