Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| """CosmosDataLoader — slim orchestrator that wires the four dataflow roles | |
| (DataDistributor -> RawItemProcessor -> SampleBatcher -> BatchCollator) inside | |
| each DataLoader worker. The canonical training dataloader. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| import torch.utils.data | |
| import numpy as np | |
| from cosmos_framework.utils import log | |
| from cosmos_framework.data.generator.dataflow.base import ( | |
| BatchCollator, | |
| DataDistributor, | |
| RawItemProcessor, | |
| SampleBatcher, | |
| ) | |
| from cosmos_framework.data.generator.dataflow.batchers import SimpleBatcher | |
| from cosmos_framework.data.generator.dataflow.collators import DefaultBatchCollator | |
| class _DataflowIterableDataset(torch.utils.data.IterableDataset): | |
| """Wires distributor -> processor -> batcher -> collator inside a worker.""" | |
| def __init__( | |
| self, | |
| distributor: DataDistributor, | |
| processor: RawItemProcessor, | |
| batcher: SampleBatcher, | |
| collator: BatchCollator, | |
| dp_rank: int, | |
| dp_world_size: int, | |
| ): | |
| super().__init__() | |
| self._distributor = distributor | |
| self._processor = processor | |
| self._batcher = batcher | |
| self._collator = collator | |
| self._dp_rank = dp_rank | |
| self._dp_world_size = dp_world_size | |
| def __iter__(self): | |
| info = torch.utils.data.get_worker_info() | |
| worker_id, num_workers = (info.id, info.num_workers) if info else (0, 1) | |
| raw = self._distributor.stream(self._dp_rank, self._dp_world_size, worker_id, num_workers) | |
| def _processed(): | |
| for item in raw: | |
| if isinstance(item, dict): | |
| meta = {k: item.pop(k) for k in list(item) if k.startswith("_dp_")} | |
| else: | |
| meta = {} | |
| s = self._processor.process(item) | |
| if meta and isinstance(s, dict): | |
| s.update(meta) | |
| yield s | |
| for group in self._batcher.batches(_processed()): | |
| has_meta = bool(group) and isinstance(group[0], dict) and "_dp_epoch" in group[0] | |
| if has_meta: | |
| epochs = [s["_dp_epoch"] for s in group] | |
| positions = [s["_dp_stream_pos"] for s in group] | |
| max_epoch = max(epochs) | |
| max_pos = max(positions) | |
| # Resume records (max_epoch, max_pos) and fast-forwards to max_pos+1 — | |
| # bit-for-bit with the legacy collate_batch. That is gap-free only when | |
| # this batch is a single sample (max_batch_size=1, all live recipes) or a | |
| # single-epoch contiguous run (sequential packing). A reordering batcher | |
| # (pool packing) at batch_size>1, or a batch spanning an epoch boundary, | |
| # would leave buffered lower positions unrecorded and skip them on resume. | |
| # Fail loudly rather than silently drop samples in that unsupported combo. | |
| if len(group) > 1: | |
| contiguous = min(epochs) == max_epoch and sorted(positions) == list( | |
| range(min(positions), max_pos + 1) | |
| ) | |
| if not contiguous: | |
| raise ValueError( | |
| "Map-style resume cannot safely stamp a multi-sample batch whose " | |
| "_dp_stream_pos values are non-contiguous or span epochs (reordering " | |
| "batcher + batch_size>1). Use max_batch_size=1 with pool packing, a " | |
| "sequential (order-preserving) batcher, or an iterable (non-resumable) " | |
| "source." | |
| ) | |
| clean = [{k: v for k, v in s.items() if not k.startswith("_dp_")} for s in group] | |
| batch = self._collator.collate(clean) | |
| batch["sample_worker_id"] = torch.tensor([worker_id] * len(group)) | |
| batch["sample_epoch"] = torch.tensor([max_epoch] * len(group)) | |
| batch["sample_index"] = torch.tensor([max_pos] * len(group)) | |
| else: | |
| batch = self._collator.collate(group) | |
| yield batch | |
| class CosmosDataLoader(torch.utils.data.DataLoader): | |
| """Public entry point: bring any dataset into training via four roles. | |
| Either pass an explicit ``batcher`` (and optional ``collator``), or pass a | |
| bare ``batch_size=N`` for stock fixed-size batching — the loader then builds | |
| ``SimpleBatcher(N)`` + ``DefaultBatchCollator()``. Passing both is an error. | |
| DP coordinates: ``parallel_dims.dp_coord`` > ``torch.distributed`` > (0, 1). | |
| """ | |
| def __init__( | |
| self, | |
| distributor: DataDistributor, | |
| processor: RawItemProcessor, | |
| batcher: SampleBatcher | None = None, | |
| collator: BatchCollator | None = None, | |
| batch_size: int | None = None, | |
| num_workers: int = 0, | |
| prefetch_factor: int | None = None, | |
| persistent_workers: bool = False, | |
| pin_memory: bool = False, | |
| parallel_dims=None, | |
| ): | |
| if batch_size is not None and batcher is not None: | |
| raise ValueError( | |
| "Pass either batch_size= (sugar) or an explicit batcher=, not both." | |
| ) | |
| if batch_size is None and batcher is None: | |
| raise ValueError("Provide either a batcher= or a batch_size=.") | |
| if batch_size is not None: | |
| batcher = SimpleBatcher(batch_size=batch_size) | |
| if collator is None: | |
| collator = DefaultBatchCollator() | |
| if parallel_dims is not None: | |
| dp_rank, dp_world_size = parallel_dims.dp_coord | |
| elif torch.distributed.is_initialized(): | |
| dp_rank = torch.distributed.get_rank() | |
| dp_world_size = torch.distributed.get_world_size() | |
| if dp_world_size > 1: | |
| log.info( | |
| "CosmosDataLoader: using global rank for DP sharding. " | |
| "For FSDP+TP/PP pass parallel_dims= for the correct DP rank.", | |
| rank0_only=True, | |
| ) | |
| else: | |
| dp_rank, dp_world_size = 0, 1 | |
| dataset = _DataflowIterableDataset( | |
| distributor=distributor, | |
| processor=processor, | |
| batcher=batcher, | |
| collator=collator, | |
| dp_rank=dp_rank, | |
| dp_world_size=dp_world_size, | |
| ) | |
| from cosmos_framework.data.generator.dataflow.distributors import MapDistributor | |
| if isinstance(distributor, MapDistributor) and num_workers > 0 and not persistent_workers: | |
| log.info( | |
| "CosmosDataLoader: MapDistributor requires persistent_workers=True for " | |
| "correct stateful resume; overriding to True.", | |
| rank0_only=True, | |
| ) | |
| persistent_workers = True | |
| if persistent_workers and num_workers == 0: | |
| log.info( | |
| "CosmosDataLoader: persistent_workers=True ignored because num_workers=0.", | |
| rank0_only=True, | |
| ) | |
| persistent_workers = False | |
| loader_kwargs: dict = dict( | |
| num_workers=num_workers, | |
| persistent_workers=persistent_workers, | |
| pin_memory=pin_memory, | |
| ) | |
| if num_workers > 0 and prefetch_factor is not None: | |
| loader_kwargs["prefetch_factor"] = prefetch_factor | |
| super().__init__(dataset, batch_size=None, **loader_kwargs) | |
| class JointCosmosDataLoader: | |
| """Wraps multiple ``CosmosDataLoader`` instances with ratio-based seeded selection. | |
| One output batch = one inner loader, selected deterministically by ratio at each | |
| step. Adds a ``"dataset_name"`` key to every yielded batch so downstream callbacks | |
| can route state updates to the correct inner loader. | |
| Parameters | |
| ---------- | |
| dataloaders: | |
| ``{name: {"dataloader": CosmosDataLoader, "ratio": int}}`` mapping. | |
| Entries with ``ratio <= 0`` are silently skipped. | |
| seed: | |
| Base seed for per-step dataset selection. Step ``i`` uses | |
| ``np.random.RandomState(seed + i)`` — fully reproducible on resume via | |
| ``set_start_iteration``. | |
| """ | |
| def __init__( | |
| self, | |
| dataloaders: dict, | |
| seed: int = 42, | |
| ) -> None: | |
| entries = [ | |
| (name, cfg["dataloader"], cfg["ratio"]) | |
| for name, cfg in dataloaders.items() | |
| if cfg.get("ratio", 0) > 0 | |
| ] | |
| if not entries: | |
| raise ValueError("JointCosmosDataLoader: no dataloaders with ratio > 0") | |
| self._names: list[str] = [e[0] for e in entries] | |
| if "global_id" in self._names: | |
| raise ValueError( | |
| "JointCosmosDataLoader: dataset name 'global_id' is reserved " | |
| "by the checkpoint state format; use a different name." | |
| ) | |
| self._loaders: list[CosmosDataLoader] = [e[1] for e in entries] | |
| ratios = np.array([e[2] for e in entries], dtype=float) | |
| self._probs: np.ndarray = ratios / ratios.sum() | |
| self._seed = seed | |
| self._global_id = 0 | |
| # Iterators are created lazily on the first __iter__ call so that | |
| # DataLoaderStateCallback.load_state_dict can install resume env vars | |
| # before workers are spawned (for num_workers > 0, iter(DataLoader) | |
| # forks workers immediately; env vars must be set in the parent first). | |
| self._iterators: list | None = None | |
| total = ratios.sum() | |
| lines = [f"JointCosmosDataLoader: {len(self._names)} streams"] | |
| for name, ratio in zip(self._names, ratios): | |
| lines.append(f" {name}: ratio={ratio:.4g} ({ratio / total:.1%})") | |
| log.info("\n".join(lines)) | |
| def set_start_iteration(self, iteration: int) -> None: | |
| """Restore deterministic selection sequence after checkpoint resume. | |
| Called by ``JointCosmosDataLoaderStateCallback.load_state_dict`` and by the | |
| trainer (if present) via ``hasattr`` guard. | |
| """ | |
| self._global_id = iteration | |
| def __iter__(self): | |
| # Lazy init: create iterators here (not in __init__) so that | |
| # load_state_dict can set resume env vars before workers fork. | |
| if self._iterators is None: | |
| self._iterators = [iter(loader) for loader in self._loaders] | |
| while True: | |
| rng = np.random.RandomState(self._seed + self._global_id) | |
| idx = int(rng.choice(len(self._loaders), p=self._probs)) | |
| try: | |
| batch = next(self._iterators[idx]) | |
| except StopIteration: | |
| # Inner CosmosDataLoaders are infinite; this guard handles | |
| # the unlikely case of a finite IterableDataset inner source. | |
| self._iterators[idx] = iter(self._loaders[idx]) | |
| batch = next(self._iterators[idx]) | |
| batch["dataset_name"] = self._names[idx] | |
| self._global_id += 1 | |
| yield batch | |