Data Loading

import datasets
import torch

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


def build_dataset(dist_context):
    # Only rank 0 downloads/processes; the others load from the cache it builds. Return it UNSHARDED —
    # AutoDataProvider shards across data-parallel ranks itself.
    with dist_context.main_process_first():
        return datasets.load_dataset("my/dataset", split="train")


def collate(samples: list[dict]) -> dict[str, torch.Tensor]:
    return {
        "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(
    data_provider=provider,
    ...,
).configure()

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,
    num_microbatches_for_global_batch,
    shard_dataset_data_parallel,
)
from d9d.loop.control import DataProvider, InitializeDataProviderContext


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[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
        return {...}

    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(str(self._config.tokenizer))

        # main_process_first ensures rank 0 builds the cache first; 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 (minimizes padding overhead).
        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,
        )
        return FixedCountMicrobatchPacker(loader, microbatches_per_step=microbatches_per_step, drop_last=True)