Spaces:
Running on Zero
Running on Zero
File size: 2,115 Bytes
a1833c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""The four dataflow role ABCs.
A raw item flows through four independently-swappable roles in a fixed order
enforced by the loader:
DataDistributor -> RawItemProcessor -> SampleBatcher -> BatchCollator
See docs/superpowers/specs/2026-06-04-modular-dataflow-refactor-design.md.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any, Iterator
class DataDistributor(ABC):
"""Owns the raw dataset, shards it disjointly across DP ranks x workers,
shuffles, and (later) carries checkpoint/resume state."""
@abstractmethod
def stream(
self, dp_rank: int, dp_world_size: int, worker_id: int, num_workers: int
) -> Iterator[Any]:
"""Yield this (rank, worker)'s disjoint slice of raw items, indefinitely."""
def state_dict(self) -> dict:
"""Resume state. No-op default; resumable distributors override."""
return {}
def load_state_dict(self, state: dict) -> None:
"""Restore resume state. No-op default; resumable distributors override."""
return None
class RawItemProcessor(ABC):
"""Transforms one raw dataset item into one training-ready sample dict."""
@abstractmethod
def process(self, item: Any) -> dict:
...
class SampleBatcher(ABC):
"""Consumes a stream of samples and yields groups (the selection strategy)."""
@abstractmethod
def batches(self, samples: Iterator[dict]) -> Iterator[list[dict]]:
"""Pull from ``samples``; yield one ``list[dict]`` per batch."""
def sample_size(self, sample: dict) -> int:
"""Per-sample token cost for packing batchers. Non-packing batchers
never call this; packing batchers override it (or inject a size_fn)."""
raise NotImplementedError
class BatchCollator(ABC):
"""Collates one group of samples into one batch dict for ``model.forward()``."""
@abstractmethod
def collate(self, samples: list[dict]) -> dict:
...
|