syvain-training-data

Internal Syvain data utility. No secret sauce here, just a shared helper.

This is my dataloader. There are many like it, but this one is mine. My dataloader is my best friend. It is my life. I must master it as I must master my life. My dataloader, without me, is useless. Without my dataloader, I am useless.

Install

uv add 'syvain-training-data[torch]'

This installs the PyTorch data loader used in the primary example below. For manifest, storage, JSONL gzip, and MessagePack/Zstandard workflows that do not load training batches, use the lightweight base package:

uv add syvain-training-data

The base install does not install NumPy, PyTorch, PyArrow, or Safetensors. Parquet users can install syvain-training-data[parquet]. Consumers of the OLMo 3 representation publication can install syvain-training-data[representation]. The full extra installs every optional feature. Requesting an optional feature without its extra raises an import error that names the required extra.

Load data

from syvain_training_data import SyvainTrainingData

training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)


def collate(records):
    ...


loader = training_data.split_data_loader(
    "s3://my-training-bucket/path/to/data-manifest-v1.json",
    collate_fn=collate,
    dataloader_args={"batch_size": 32, "num_workers": 4, ...},
)

train_batches = loader.load("train")
valid_batches = loader.load("valid")
easy_batches = loader.load("train", curriculum_stage="easy")
early_curriculum_batches = loader.load("train", curriculum_stages=["easy", "medium"])
infinite_train_batches = loader.load("train", infinite_iter=True)

curriculum_stages selects the union of the named stages. It does not guarantee records are yielded in stage order, especially when num_workers is greater than zero.

Storage reads recover from transient S3/Tigris connection and body failures by opening a fresh client and resuming immutable shard streams at the last received byte. Point reads and writes retry the complete operation at the same URI. Missing objects, authentication failures, and invalid data still fail closed.

Load joined OLMo 3 representation examples

The representation loader verifies the pinned final manifest, the ETag-bound indexes, the index-bound artifact manifest, and each selected artifact. It then joins prepared tokens to the requested teacher projection rows.

examples = training_data.iter_representation_examples(
    manifest_uri,
    manifest_sha256,
    split="analysis",
    sources=["broad_replay", "dolmino"],
    per_source_count=64,
    teacher_layer=31,
    state_kind="post_mlp",
)

The loader processes source groups in caller order. Within each source, it selects and yields the first per_source_count examples by ascending block_id. Each RepresentationExample owns its returned tensors, so no Safetensors mapping or full artifact remains live after the loader copies the selected rows.

The package uses a 10-second connect timeout, a 60-second read-inactivity timeout, and a 10-minute overall request timeout. Obstore's internal retry window is deliberately short; the package owns the longer 15-minute no-progress recovery window so a failed connection pool can be discarded.

When using worker processes, leave PyTorch DataLoader(timeout=0) unless the training runtime has a specific worker watchdog. A positive DataLoader timeout must be longer than the storage recovery window plus normal shard processing; a value such as 120 seconds can terminate a healthy worker while it is retrying a transient object-store outage.

Derive data

Use the manifest's format to stream source shards when generating a derived dataset:

from syvain_training_data import iter_shard

for shard in manifest.splits["train"].shards:
    for record in iter_shard(
        manifest.data_format,
        shard,
        storage_config=storage_config,
    ):
        ...

Save data

from concurrent.futures import ProcessPoolExecutor

from syvain_training_data import SyvainTrainingData

def generate_data(split, curriculum_stage, shard_id):
    ...

def save_shard(job):
    saver, split, curriculum_stage, metadata, shard_id = job
    records = generate_data(split, curriculum_stage, shard_id)
    saver.save(
        split,
        curriculum_stage,
        records,
        curriculum_metadata=metadata,
        shard_id=str(shard_id),
    )


training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)

saver = training_data.dataset_saver(
    "s3://my-training-bucket/path/to/dataset/data-manifest-v1.json",
)

jobs = [
    (saver, "train", stage["name"], stage, shard_id)
    for stage in [
        {"name": "easy", "family": "arithmetic", "weight": 1.0},
        {"name": "medium", "family": "control", "weight": 2.0},
        {"name": "hard", "family": "composition", "weight": 3.0},
    ]
    for shard_id in range(32)
] + [
    (saver, "valid", None, None, shard_id) for shard_id in range(4)
] + [
    (saver, "test", None, None, shard_id) for shard_id in range(4)
]

with ProcessPoolExecutor(max_workers=8) as pool:
    list(pool.map(save_shard, jobs))

manifest = saver.commit_manifest()

Each deterministic shard writes a completion descriptor after its data object. A restarted saver verifies that descriptor and reuses the shard without consuming records. The data manifest is written last as the publication marker; committing the same completed publication again is idempotent. Call saver.recover(split, curriculum_stage, shard_id=...) to inspect a completed shard explicitly without constructing a records iterator.

Deterministic Parquet and framed MessagePack shards use an atomic conditional upload and therefore must each be smaller than 5 GiB. Use more logical shard IDs when generating a larger derived dataset.

Copy a manifest

from syvain_training_data import SyvainTrainingData

training_data = SyvainTrainingData(
    s3_base_url="https://t3.storage.dev",
    region="auto",
    access_key_id="...",
    secret_access_key="...",
)

manifest = training_data.load_manifest("s3://my-training-bucket/shared/data-manifest-v1.json")

# Do modifications if needed

training_data.save_manifest("s3://my-training-bucket/new-run/data-manifest-v1.json", manifest)