Data Loading

import dataclasses

import datasets
import torch

from d9d.loop.auto import AutoDataConfig, AutoDataProvider
from d9d.loop.run import TrainingConfigurator


def build_dataset(dist_context):
    # Rank 0 downloads and processes the data first. The other ranks load from its cache.
    # Return the dataset unsharded: AutoDataProvider shards it across data-parallel ranks.
    with dist_context.main_process_first():
        return datasets.load_dataset("my/dataset", split="train")


@dataclasses.dataclass
class Batch:
    input_ids: torch.Tensor
    labels: torch.Tensor


# The rows of a Hugging Face dataset are dicts.
def collate(samples: list[dict]) -> Batch:
    return Batch(
        input_ids=torch.stack([s["input_ids"] for s in samples]),
        labels=torch.stack([s["labels"] for s in samples]),
    )


provider = AutoDataProvider(
    dataset_factory=build_dataset,
    collator=collate,
    config=AutoDataConfig(
        global_batch_size=256,  # Effective batch across all DP replicas and accumulation
        microbatch_size=4,  # Samples per microbatch on a single rank
        num_workers=4,
        shuffle=True,
    ),
)

# Hand it to the loop alongside the other providers.
trainer = TrainingConfigurator(
    mesh=...,
    parameters=...,
    model_provider=...,
    task_provider=...,
    data_provider=provider,
    optimizer_provider=...,
    lr_scheduler_provider=...,
).configure()

import dataclasses
from collections.abc import Sequence
from typing import Any

import torch
import datasets
from pydantic import BaseModel
from tokenizers import Tokenizer
from torch.utils.data import Dataset
from torchdata.stateful_dataloader import StatefulDataLoader

from d9d.core.protocol import MicrobatchPackStream
from d9d.core.types import TensorTree
from d9d.dataset import (
    BufferSortedDataset,
    DatasetImplementingSortKeyProtocol,
    FixedCountMicrobatchPacker,
    PinMemoryMicrobatchPackStream,
    num_microbatches_for_global_batch,
    shard_dataset_data_parallel,
)
from d9d.loop.control import DataProvider, InitializeDataProviderContext


@dataclasses.dataclass
class ProjectBatch:
    input_ids: torch.Tensor
    labels: torch.Tensor


class ProjectDataset(Dataset, DatasetImplementingSortKeyProtocol):
    def __init__(self, dataset: datasets.Dataset, tokenizer: Tokenizer):
        self._dataset = dataset
        self._tokenizer = tokenizer

    def sort_key(self, index: int) -> Any:
        # Used by BufferSortedDataset to group examples of similar length together.
        return self._dataset[index]["token_counts"]

    def __getitem__(self, index: int) -> TensorTree:
        return {...}

    @classmethod
    def collate(cls, batch: Sequence[TensorTree]) -> ProjectBatch:
        return ProjectBatch(...)

    def __len__(self) -> int:
        return len(self._dataset)


class DataConfig(BaseModel):
    dataset: str
    split: str
    tokenizer: str
    presort_buffer_size: int
    global_batch_size: int
    microbatch_size: int
    num_workers: int


class ProjectDataProvider(DataProvider):
    def __init__(self, config: DataConfig):
        self._config = config

    def __call__(self, context: InitializeDataProviderContext) -> MicrobatchPackStream:
        tokenizer = Tokenizer.from_file(self._config.tokenizer)

        # Rank 0 builds the cache first. The other ranks then load from it.
        with context.dist_context.main_process_first():
            data = datasets.load_dataset(self._config.dataset, split=self._config.split)

        dataset = ProjectDataset(data, tokenizer)

        # Length-bucketing buffer, which minimizes padding.
        dataset_buf = BufferSortedDataset(
            dataset,
            buffer_size=self._config.presort_buffer_size,
            pack_size=self._config.global_batch_size,
        )

        # Shard across data-parallel ranks (a custom DataProvider owns sharding).
        dataset_shard = shard_dataset_data_parallel(dataset_buf, context.dist_context)

        # Loader (layer 1): yields one collated microbatch at a time.
        loader = StatefulDataLoader(
            dataset_shard,
            batch_size=self._config.microbatch_size,
            collate_fn=ProjectDataset.collate,
            num_workers=self._config.num_workers,
            drop_last=True,
        )

        # Packer (layer 2): groups microbatches into one pack per step.
        microbatches_per_step = num_microbatches_for_global_batch(
            context.dist_context,
            global_batch_size=self._config.global_batch_size,
            microbatch_size=self._config.microbatch_size,
        )
        packer = FixedCountMicrobatchPacker(loader, microbatches_per_step=microbatches_per_step, drop_last=True)

        # Pin the packs, so the loop copies them to the device asynchronously.
        return PinMemoryMicrobatchPackStream(packer)