# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 import math import queue import threading from collections import deque from collections.abc import Iterator, Mapping from dataclasses import dataclass from typing import Any, ClassVar, Dict, Union import numpy as np import torch import webdataset from torch.utils.data.dataloader import default_collate from cosmos_framework.utils.lazy_config import instantiate from cosmos_framework.utils import log from cosmos_framework.model.generator.tokenizers.uniae.frame_math import ( get_uniae_chunk_frames, get_uniae_latent_num_frames, normalize_uniae_chunk_frames, ) from cosmos_framework.utils.generator.data_utils import read_positive_int_metadata _TIMING_KEYS = {"_sample_time", "_aug_time", "_pre_aug_time", "_aug_step_times"} _BATCH_TIMING_KEYS = { "_worker_batch_time", "_worker_aug_time", "_worker_io_time", "_worker_aug_step_times", "_worker_id", } def custom_collate_fn(batch): """ Collate function that works like default_collate for all keys other than "text_token_ids", "images", and "video". For "text_token_ids", "images", and "video" it simply returns them in a list, instead of stacking them as a tensor. """ list_collate_keys = { "text_token_ids", "images", "video", "action", "action_raw", "domain_id", "sequence_plan", "sound", "raw_action_dim", "image_size", "action_processing_record", "residual_gate", "residual_release_event", } # Data keys where a per-sample value of ``None`` is a meaningful signal # (e.g. audio extraction failed for that sample → ``sound=None`` paired # with ``plan.has_sound=False``). These keys must be kept as a list with # ``None`` placeholders so the model can align per-sample data 1:1 with # per-sample plans. Dropping the entire key on any None would leave the # remaining sound tensors mis-aligned with the plans whose ``has_sound`` # flag was set BEFORE collation, causing ``sequence_packing`` to index # past the end of ``x0_tokens_sound``. sparse_data_keys = {"sound"} # Handle the case where the batch is already a dictionary (e.g. column-wise batching) if isinstance(batch, dict): return {key: (value if key in list_collate_keys else default_collate(value)) for key, value in batch.items()} # Handle standard list of samples elem = batch[0] if isinstance(elem, dict): # Some Action datasets add optional metadata keys (for example # ``additional_view_description`` for concat-view captions) only for a # subset of samples. PyTorch can batch such samples together when # DataLoader batch_size > 1; collating only elem's keys and indexing # every sample by that key turns the optional field into a fatal # KeyError. Use the union of keys and skip optional keys that are not # present in every sample. Required training keys still fail loudly via # downstream assertions if actually missing. result = {} keys = set().union(*(d.keys() for d in batch)) for key in keys: if key in _TIMING_KEYS: continue values = [d.get(key) for d in batch] if key == "action_processing_record": result[key] = values continue if any(value is None for value in values): # Sparse data keys keep their None placeholders to preserve # 1:1 alignment with sequence_plan. Other (optional metadata) # keys not present in every sample are dropped. if key in sparse_data_keys: result[key] = values continue if key in list_collate_keys: result[key] = values else: result[key] = default_collate(values) result.update(_aggregate_worker_timing(batch)) return result else: return default_collate(batch) def _aggregate_worker_timing(samples: list[dict]) -> dict: """Extract per-sample timing keys, aggregate into per-batch scalars.""" info: dict[str, float | int] = {} if "_sample_time" in samples[0]: info["_worker_batch_time"] = sum(s.get("_sample_time", 0.0) for s in samples) if "_aug_time" in samples[0]: aug_total = sum(s.get("_aug_time", 0.0) for s in samples) info["_worker_aug_time"] = aug_total if "_worker_batch_time" in info: info["_worker_io_time"] = info["_worker_batch_time"] - aug_total if "_aug_step_times" in samples[0]: agg: dict[str, float] = {} for s in samples: for step_name, t in s.get("_aug_step_times", {}).items(): agg[step_name] = agg.get(step_name, 0.0) + t info["_worker_aug_step_times"] = agg worker_info = torch.utils.data.get_worker_info() info["_worker_id"] = worker_info.id if worker_info is not None else 0 return info @dataclass class _PackingMetrics: """Per-batch packing statistics collected during the packing loop. Also serves as the single source of truth for packing-related metric names via ``STATS_SPEC``, which the dataloading monitor callback consumes to drive accumulation and logging. """ current_sequence_length: int = 0 num_samples: int = 0 dropped_count: int = 0 from_buffer: int = 0 from_workers: int = 0 STATS_SPEC: ClassVar[list[tuple[str, str, str]]] = [ # (batch_key, wandb_suffix, aggregation_type) ("_num_tokens", "token_fraction", "scalar"), ("_num_samples", "samples_per_batch", "list"), ("_from_buffer", "from_buffer", "list"), ("_from_workers", "from_workers", "list"), ("_buffer_size", "buffer_size", "list"), ("_dropped_count", "dropped", "scalar"), ] def attach_to(self, output_batch: dict, buffer_size: int) -> None: """Write packing statistics into the output batch dict.""" output_batch["_num_tokens"] = self.current_sequence_length output_batch["_num_samples"] = self.num_samples output_batch["_from_buffer"] = self.from_buffer output_batch["_from_workers"] = self.from_workers output_batch["_buffer_size"] = buffer_size output_batch["_dropped_count"] = self.dropped_count @dataclass class _AsyncBatchBuildResult: """The outcome of building one packed batch.""" sequence_id: int batch: dict[str, Any] | None = None error: BaseException | None = None end_of_stream: bool = False class _AsyncBatchBuilder: """Build one requested batch at a time on a persistent daemon thread. The caller submits one sequence ID, the daemon advances the existing packing iterator once, and the caller receives that exact batch object from the result queue. A second request is rejected until the first result is consumed, so at most one future batch is built at a time. """ def __init__(self, batch_iterator: Iterator[dict[str, Any]], timeout_s: float) -> None: self._batch_iterator: Iterator[dict[str, Any]] = batch_iterator self._timeout_s: float = timeout_s self._requests: queue.Queue[int | None] = queue.Queue() self._results: queue.Queue[_AsyncBatchBuildResult | None] = queue.Queue() self._outstanding_sequence_id: int | None = None self._closed: bool = False self._thread: threading.Thread = threading.Thread(target=self._run, name="async-batch-builder", daemon=True) self._thread.start() @property def is_alive(self) -> bool: return self._thread.is_alive() def request_batch(self, sequence_id: int) -> None: """Ask the daemon thread to build exactly one batch.""" if self._closed: raise RuntimeError("Cannot request a batch after the async batch builder is closed.") if self._outstanding_sequence_id is not None: raise RuntimeError( "Cannot request another batch while async batch building is already in progress " f"for sequence_id={self._outstanding_sequence_id}." ) self._outstanding_sequence_id = sequence_id self._requests.put(sequence_id) def get_batch(self) -> dict[str, Any]: """Wait for and return the exact batch object produced for the outstanding request.""" sequence_id = self._outstanding_sequence_id if sequence_id is None: raise RuntimeError("No async batch building request is outstanding.") if self._closed: raise RuntimeError("Async batch builder was closed while waiting for a batch.") try: result = self._results.get(timeout=self._timeout_s) except queue.Empty as error: producer_alive = self.is_alive # A timed-out request may still be running in the daemon thread, so it cannot be # safely abandoned or retried. Close the builder and make the timeout terminal. self.close() raise TimeoutError( f"Timed out waiting for async batch building sequence_id={sequence_id}; " f"producer_alive={producer_alive}." ) from error if self._closed or result is None: raise RuntimeError("Async batch builder was closed while waiting for a batch.") if result.sequence_id != sequence_id: raise RuntimeError( "Async batch builder returned an unexpected result: " f"expected sequence_id={sequence_id}, got sequence_id={result.sequence_id}." ) self._outstanding_sequence_id = None if result.error is not None: raise RuntimeError(f"Async batch building failed for sequence_id={sequence_id}.") from result.error if result.end_of_stream: raise StopIteration if result.batch is None: raise RuntimeError(f"Async batch builder returned no batch for sequence_id={sequence_id}.") return result.batch def _run(self) -> None: """Build one batch for each request and publish its result by reference.""" while True: sequence_id = self._requests.get() if sequence_id is None or self._closed: return result = _AsyncBatchBuildResult(sequence_id=sequence_id) try: result.batch = next(self._batch_iterator) except StopIteration: result.end_of_stream = True except BaseException as error: result.error = error if self._closed: return # queue.Queue stores object references within this process. The batch is neither # copied nor serialized here; the consumer receives the exact dictionary above. self._results.put(result) if result.end_of_stream or result.error is not None: return def close(self) -> None: """Permanently close the builder and perform a bounded join of its daemon thread.""" if not self._closed: self._closed = True self._requests.put(None) self._results.put(None) if self._thread is not threading.current_thread(): self._thread.join(timeout=min(1.0, self._timeout_s)) class JointDataLoader(webdataset.WebLoader): r""" A joint dataloader that supports loading both images and videos. """ _DEFAULT_LOOKAHEAD_LIMIT: ClassVar[int] = 10 def __init__( self, dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]], tokenizer_spatial_compression_factor: int, tokenizer_temporal_compression_factor: int, patch_spatial: int, max_sequence_length: int | None, max_samples_per_batch: int | None, sound_latent_fps: float = 0, audio_sample_rate: int = 48000, prewarm: bool = True, default_lookahead_limit: int = _DEFAULT_LOOKAHEAD_LIMIT, lookahead_limits: Dict[str, int] | None = None, uniae_chunk_frames: int | Mapping[str, int] | None = None, uniae_pad_frames: int | None = None, ): """ Initialize the JointDataLoader with multiple datasets. The effective mini-batch size can be controlled with either max_sequence_length or max_samples_per_batch. To use max_sequence_length, max_samples_per_batch needs to be None. Vice versa, to use max_samples_per_batch, max_sequence_length needs to be None. max_sequence_length and max_samples_per_batch cannot both be None simultaneously. Args: dataloaders: key - dataset_name; value - {"dataloader": dataloader, "ratio": data_ratio} tokenizer_spatial_compression_factor: The spatial compression factor of the tokenizer. tokenizer_temporal_compression_factor: The temporal compression factor of the tokenizer. patch_spatial: Spatial pathification factor. max_samples_per_batch: Max number of samples per packed batch (alternative to max_sequence_length). sound_latent_fps: Sound tokenizer latent rate in Hz (e.g. 25). If 0, sound tokens are not counted. audio_sample_rate: Audio sample rate in Hz (e.g. 48000). Used with sound_latent_fps to estimate sound token count. default_lookahead_limit: Packing-loop look-ahead fallback for dataloaders not in ``lookahead_limits``. lookahead_limits: Optional ``{dataset_name: int}`` per-dataloader override. uniae_chunk_frames: Optional UniAE full chunk size, or resolution-keyed chunk sizes. uniae_pad_frames: Optional UniAE boundary padding frames per chunk. Example: joint_loader = IterativeJointDataLoader( dataloaders{ "image_data": { "dataloader": webdataset.WebLoader(...), "ratio": 4, }, "video_data": { "dataloader": torch.utils.data.DataLoader(...), "ratio": 1, }, } ) """ self.dataloader_list, self.dataset_name_list, self.data_ratios = [], [], [] self.lookahead_limits: list[int] = [] self.tokenizer_spatial_compression_factor = tokenizer_spatial_compression_factor self.tokenizer_temporal_compression_factor = tokenizer_temporal_compression_factor self.patch_spatial = patch_spatial self.max_sequence_length = max_sequence_length self.max_samples_per_batch = max_samples_per_batch self.sound_latent_fps = sound_latent_fps self.audio_sample_rate = audio_sample_rate self.default_lookahead_limit = int(default_lookahead_limit) self.uniae_pad_frames = int(uniae_pad_frames) if uniae_pad_frames is not None else None self.uniae_chunk_frames = self._normalize_uniae_chunk_frames(uniae_chunk_frames) assert (self.max_sequence_length is None) != (self.max_samples_per_batch is None), ( "Exactly one of max_sequence_length or max_samples_per_batch must be None, but not both." ) _lookahead_overrides: Dict[str, int] = dict(lookahead_limits) if lookahead_limits else {} unknown = set(_lookahead_overrides) - set(dataloaders) assert not unknown, f"lookahead_limits references unknown dataloaders {unknown}; valid: {sorted(dataloaders)}" for dataset_name, dataloader_data in dataloaders.items(): if dataloader_data is None: continue assert set(dataloader_data.keys()) == {"dataloader", "ratio"}, f"Invalid config: {dataloader_data}" if dataloader_data["ratio"] <= 0: continue self.dataset_name_list.append(dataset_name) self.dataloader_list.append(instantiate(dataloader_data["dataloader"], collate_fn=custom_collate_fn)) self.data_ratios.append(dataloader_data["ratio"]) self.lookahead_limits.append(int(_lookahead_overrides.get(dataset_name, self.default_lookahead_limit))) self.global_id = 0 self.ratio_sum = sum(self.data_ratios) total = self.ratio_sum if self.ratio_sum > 0 else 1.0 lines = [f"JointDataLoader: {len(self.dataset_name_list)} streams"] for name, ratio in zip(self.dataset_name_list, self.data_ratios): lines.append(f" {name}: ratio={ratio:.4g} ({ratio / total:.1%})") log.info("\n".join(lines)) self.data_len = 0 self.dataloaders = [iter(dataloader) for dataloader in self.dataloader_list] self.buffers = [deque() for _ in range(len(self.dataloader_list))] for data in self.dataloader_list: self.data_len += len(data) # Pre-warm all dataloaders: force worker process spawning and first # batch loading so that slow dataset initialisation (e.g. action # datasets with spawn workers) happens here rather than mid-training # where it would cause NCCL collective timeouts. if prewarm: self._prewarm_dataloaders() else: log.info( "JointDataLoader: prewarm DISABLED (debug mode); first iteration may incur per-stream cold-load cost" ) def _normalize_uniae_chunk_frames( self, uniae_chunk_frames: int | Mapping[str, int] | None ) -> int | dict[str, int] | None: return normalize_uniae_chunk_frames( uniae_chunk_frames, pad_frames=self.uniae_pad_frames, temporal_compression_factor=self.tokenizer_temporal_compression_factor, temporal_divisibility_name="tokenizer_temporal_compression_factor", ) def _get_uniae_chunk_frames(self, spatial_shape: tuple[int, int]) -> int: assert self.uniae_chunk_frames is not None return get_uniae_chunk_frames(self.uniae_chunk_frames, spatial_shape=spatial_shape) def _compute_vision_latent_t_shape(self, T: int, H: int, W: int) -> int: if T < 1: raise ValueError(f"Vision media must contain at least one frame, got {T}.") if T == 1 or self.uniae_chunk_frames is None: return 1 + (T - 1) // self.tokenizer_temporal_compression_factor assert self.uniae_pad_frames is not None return get_uniae_latent_num_frames( T, self.uniae_chunk_frames, pad_frames=self.uniae_pad_frames, temporal_compression_factor=self.tokenizer_temporal_compression_factor, spatial_shape=(H, W), ) def _prewarm_dataloaders(self) -> None: """Force all dataloader iterators to spawn workers and produce one batch. The first ``next()`` call on an ``InfiniteDataLoader`` iterator triggers ``DataLoader.__iter__()`` which spawns worker processes. For action dataloaders using ``multiprocessing_context='spawn'``, each worker must fully initialise heavy datasets (BridgeOrigLeRobotDataset, embodiment_a, etc.) from scratch. If this happens lazily during training, the resulting delay (potentially minutes) causes NCCL collective timeouts when faster ranks enter the forward pass while slower ranks are still loading data. By pulling one batch from every dataloader here — before any training iteration — we ensure all workers are alive and warmed up. The fetched samples are pushed into the per-dataloader buffer so they are consumed normally by the first iteration that selects that dataloader. A ``dist.barrier()`` at the end synchronises all ranks so that training only begins once every rank has finished pre-warming. """ import time for i, (name, dl_iter) in enumerate(zip(self.dataset_name_list, self.dataloaders)): t0 = time.monotonic() try: batch = next(dl_iter) except StopIteration: log.warning(f"Pre-warm: dataloader {name!r} is empty, skipping") continue elapsed = time.monotonic() - t0 # Split the collated batch into individual samples and push them # into the buffer — identical to the splitting logic in # _get_next_sample — so the samples are not wasted. is_image_batch = "images" in batch input_images_or_videos = batch["images" if is_image_batch else "video"] batch_size = len(input_images_or_videos) for j in range(batch_size): sample = {} for k, v in batch.items(): if k in _BATCH_TIMING_KEYS: sample[k] = v elif isinstance(v, list) and k in self._MULTI_ITEM_KEYS: elem = v[j] if isinstance(elem, list): sample[k] = elem else: sample[k] = v[j : j + 1] elif isinstance(v, list): sample[k] = v[j] elif isinstance(v, torch.Tensor) and v.dim() > 0: sample[k] = v[j : j + 1] else: sample[k] = v[j : j + 1] self.buffers[i].append(sample) log.info( f"Pre-warm: dataloader {name!r} ready — {batch_size} samples buffered in {elapsed:.1f}s", rank0_only=False, ) # Synchronise so training only starts once every rank is warmed up. if torch.distributed.is_initialized(): log.info("Pre-warm: waiting at barrier for all ranks …") torch.distributed.barrier() log.info("Pre-warm: all ranks ready") def _compute_num_tokens_per_sample(self, data_batch: dict) -> int: """ This function computes the number of tokens per sample in the data batch. This includes text + vision generation tokens + action tokens. Args: data_batch (dict): The data batch containing the text tokens. Returns: int: The number of tokens per sample. """ # The token sequence we have is # [] # The spatial dimension of image tokens is compressed by # vae spatial downsampling factor + pathification # The temporal dimension of image tokens is compressed by # vae temporal downsampling factor # Action tokens have 1 token per time step (no spatial dimension) has_text_tokens = "text_token_ids" in data_batch num_tokens = 0 if has_text_tokens: text_token_ids = data_batch["text_token_ids"] if isinstance(text_token_ids, list): num_text_tokens = text_token_ids[0].shape[0] else: num_text_tokens = text_token_ids.shape[1] num_tokens += num_text_tokens + 1 # Vision part is_image_batch = "images" in data_batch input_images_or_videos = data_batch["images" if is_image_batch else "video"] if "enable_per_camera_vae_encoding" in data_batch: sample_n_views_values = read_positive_int_metadata(data_batch, "sample_n_views", expected_count=1) frames_per_view_values = read_positive_int_metadata( data_batch, "num_video_frames_per_view", expected_count=1, ) if sample_n_views_values is None or frames_per_view_values is None: raise ValueError( "sample_n_views and num_video_frames_per_view are required when per-camera VAE encoding is enabled." ) sample_n_views = sample_n_views_values[0] frames_per_view = frames_per_view_values[0] else: sample_n_views = None frames_per_view = None # iterate over all the media in the batch for media in input_images_or_videos if isinstance(input_images_or_videos, list) else [input_images_or_videos]: if is_image_batch: _, H, W = media.shape T = 1 else: _, T, H, W = media.shape latent_h_shape = H // self.tokenizer_spatial_compression_factor latent_w_shape = W // self.tokenizer_spatial_compression_factor patch_h_shape = math.ceil(latent_h_shape / self.patch_spatial) patch_w_shape = math.ceil(latent_w_shape / self.patch_spatial) if sample_n_views is not None and frames_per_view is not None and not is_image_batch: # The multiview dataset resizes every camera to the same H/W and # selects the same synchronized frame count before concatenating # the camera-major clips along T. expected_frames = sample_n_views * frames_per_view if T != expected_frames: raise ValueError( "Multiview vision length must equal sample_n_views * num_video_frames_per_view: " f"got T={T}, sample_n_views={sample_n_views}, num_video_frames_per_view={frames_per_view}." ) latent_t_shape = sample_n_views * self._compute_vision_latent_t_shape(frames_per_view, H, W) else: latent_t_shape = self._compute_vision_latent_t_shape(T, H, W) num_vision_tokens = patch_h_shape * patch_w_shape * latent_t_shape if has_text_tokens: num_vision_tokens += 2 num_tokens += num_vision_tokens # Action part: each action time step is 1 token. # Action tensor shape is (T_action, D) per sample; stored as a single-element list. if "action" in data_batch: list_of_actions = data_batch["action"] for action in list_of_actions: # skip None actions if action is None: continue num_action_tokens = action.shape[0] num_tokens += num_action_tokens # Sound part — estimate sound tokens from audio waveform length if self.sound_latent_fps > 0 and "sound" in data_batch: sound_data = data_batch["sound"] if isinstance(sound_data, list) and len(sound_data) > 0: first_sound = sound_data[0] # Unwrap nested list if needed if isinstance(first_sound, list): first_sound = first_sound[0] if first_sound is not None and isinstance(first_sound, torch.Tensor): num_audio_samples = first_sound.shape[-1] audio_duration = num_audio_samples / self.audio_sample_rate num_sound_tokens = int(audio_duration * self.sound_latent_fps) num_tokens += num_sound_tokens return num_tokens # Keys whose value per sample is a list of tensors to be flattened into one list in the batch _FLATTEN_LIST_KEYS = {"image_size"} def _update_output_batch(self, output_batch: dict, output: dict): for key, value in output.items(): if key in _BATCH_TIMING_KEYS: if key not in output_batch: output_batch[key] = value elif key in self._FLATTEN_LIST_KEYS and isinstance(value, list): if key not in output_batch: output_batch[key] = value else: output_batch[key].extend(value) elif key not in output_batch: output_batch[key] = [value] else: output_batch[key].append(value) def __len__(self) -> int: return self.data_len # Keys where each sample may hold multiple tensors (e.g. multiple video # clips in a packed sequence). Kept as single-element lists per sample # via v[i:i+1] so that _update_output_batch yields list[list[Tensor]]. _MULTI_ITEM_KEYS = {"text_token_ids", "images", "video", "action", "action_raw", "sound"} def _get_next_sample(self, index_id: int) -> dict: """Pop the next single-sample dict from the buffer for the given dataloader. If the buffer is empty, fetches the next collated batch from the inner dataloader and splits it into individual samples. Splitting rules: - Multi-item list values (keys in ``_MULTI_ITEM_KEYS``): sliced via ``v[i:i+1]`` to yield a single-element list ``[tensor]``. A packed sequence can contain multiple items per key. - Per-sequence metadata list values (all other list keys, e.g. ``sequence_plan``, ``domain_id``): direct-indexed via ``v[i]`` to yield the bare element. - Tensor values ``(B, ...)``: sliced to ``(1, ...)`` via ``v[i : i + 1]`` to preserve the batch dimension. After ``_update_output_batch`` accumulates samples, the packed output batch has the following shapes: - Multi-item keys (``text_token_ids``, ``video``, ``images``, ``action_raw``, ``action``): ``list[list[Tensor]]`` — each inner list has one element from one sub-sample. - Per-sequence metadata keys (``sequence_plan``, ``domain_id``, ``dataset_name``, etc.): ``list[element]`` — flat list. - Tensor-origin keys: ``list[Tensor(1, ...)]``. Args: index_id: Index of the dataloader to fetch from. Returns: A single-sample dictionary. """ buffer = self.buffers[index_id] if not buffer: try: batch = next(self.dataloaders[index_id]) except StopIteration: raise is_image_batch = "images" in batch input_images_or_videos = batch["images" if is_image_batch else "video"] batch_size = len(input_images_or_videos) for i in range(batch_size): sample = {} for k, v in batch.items(): if k in _BATCH_TIMING_KEYS: sample[k] = v elif isinstance(v, list) and k in self._MULTI_ITEM_KEYS: # For multi-item keys (images, video, etc.), the collated # value is a list with one element per sample. If the element # is itself a list (e.g. image editing: [src, tgt]), use v[i] # directly to avoid wrapping it in a redundant single-element # list. Otherwise keep the v[i:i+1] slice so that # _update_output_batch produces list[list[Tensor]]. elem = v[i] if isinstance(elem, list): sample[k] = elem else: sample[k] = v[i : i + 1] elif isinstance(v, list): sample[k] = v[i] else: sample[k] = v[i : i + 1] buffer.append(sample) return buffer.popleft() def set_start_iteration(self, iteration: int): self.global_id = iteration def __iter__(self): raise NotImplementedError("__iter__ function is not implemented yet") class IterativeJointDataLoader(JointDataLoader): r""" An iterative joint dataloader that supports loading multiple modalities. The behavior depends on the ``seed`` parameter: - **seed is not None** (Default): The modality is randomly selected at each iteration based on the probability distribution derived from the ratios. The random state is seeded with ``seed + global_id``, ensuring that all ranks select the same modality at the same iteration (assuming synchronized global_id). This prevents load imbalance due to mixed resolutions across ranks. - **seed is None**: The modality selection follows a deterministic round-robin pattern based on the ratios. For example, with 2 modalities (image and video) and ratio 2:1: - Iterations 0, 1: all ranks process images - Iteration 2: all ranks process videos - ... and so on. This also ensures all ranks process the same modality at the same iteration. ``enable_async_batch_building=True`` advances this same batch iterator on one daemon thread and builds exactly one future batch. The default keeps the original calling-thread behavior. """ def __init__( self, dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]], tokenizer_spatial_compression_factor: int, tokenizer_temporal_compression_factor: int, patch_spatial: int, max_sequence_length: int | None = None, max_samples_per_batch: int | None = None, sound_latent_fps: float = 0, audio_sample_rate: int = 48000, seed: int | None = 42, prewarm: bool = True, default_lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT, lookahead_limits: Dict[str, int] | None = None, uniae_chunk_frames: int | Mapping[str, int] | None = None, uniae_pad_frames: int | None = None, enable_async_batch_building: bool = False, async_batch_building_timeout_s: float = 1200.0, ) -> None: if async_batch_building_timeout_s <= 0: raise ValueError(f"async_batch_building_timeout_s must be positive, got {async_batch_building_timeout_s}.") # Keep PyTorch/WebDataset/Lance worker creation on the main thread. # The packer thread starts only after constructor prewarm has completed. if enable_async_batch_building and not prewarm: raise ValueError("enable_async_batch_building=True requires prewarm=True.") super().__init__( dataloaders, tokenizer_spatial_compression_factor, tokenizer_temporal_compression_factor, patch_spatial, max_sequence_length, max_samples_per_batch, sound_latent_fps=sound_latent_fps, audio_sample_rate=audio_sample_rate, prewarm=prewarm, default_lookahead_limit=default_lookahead_limit, lookahead_limits=lookahead_limits, uniae_chunk_frames=uniae_chunk_frames, uniae_pad_frames=uniae_pad_frames, ) self.seed = seed # Calculate probabilities for random sampling total_ratio = sum(self.data_ratios) self.data_probs = np.array([ratio / total_ratio for ratio in self.data_ratios]) # Async batch builder state self.enable_async_batch_building: bool = enable_async_batch_building self.async_batch_building_timeout_s: float = float(async_batch_building_timeout_s) self._async_batch_builder: _AsyncBatchBuilder | None = None def __iter__(self) -> Iterator[dict[str, Any]]: if not self.enable_async_batch_building: return self._iter_synchronous() return self._iter_asynchronous() def _iter_synchronous(self) -> Iterator[dict[str, Any]]: """Run the existing packing path on the calling thread.""" while True: if self.seed is not None: rng = np.random.RandomState(self.seed + self.global_id) index_id = rng.choice(len(self.dataloader_list), p=self.data_probs) else: data_id = self.global_id % self.ratio_sum index_id = self._get_dataloader_index(data_id) metrics = _PackingMetrics() output_batch = dict() skipped_samples = deque() lookahead_limit = self.lookahead_limits[index_id] lookahead_count = 0 while True: # Check max samples limit first if self.max_samples_per_batch is not None and metrics.num_samples >= self.max_samples_per_batch: break # If we have started packing and tried lookahead_limit times to find a fitting sample but failed, stop. if len(output_batch) > 0 and lookahead_count >= lookahead_limit: break had_buffer = len(self.buffers[index_id]) > 0 try: output = self._get_next_sample(index_id) except StopIteration: break # No more data in this dataloader if had_buffer: metrics.from_buffer += 1 else: metrics.from_workers += 1 num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output) if ( self.max_sequence_length is not None and metrics.current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length ): if len(output_batch) == 0: # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length # In this case, we should simply discard the current sample and get the next sample. log.info( f"Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}", rank0_only=False, ) metrics.dropped_count += 1 continue # current_sequence_length > 0 and selected sample is too large to fit in the remaining space. # Instead of stopping immediately (creating large padding), we buffer this large sample # and try to find a smaller one that fits in the remaining space. skipped_samples.append(output) lookahead_count += 1 continue metrics.current_sequence_length += num_tokens_in_current_sample metrics.num_samples += 1 output["dataset_name"] = self.dataset_name_list[index_id] self._update_output_batch(output_batch, output) # Add back skipped samples to the buffer for the next batch. # appendleft puts item at HEAD. So we insert S3, then S2, then S1. for sample in reversed(skipped_samples): self.buffers[index_id].appendleft(sample) if len(output_batch) == 0: return metrics.attach_to(output_batch, buffer_size=len(self.buffers[index_id])) self.global_id += 1 yield output_batch def _start_async_batch_builder(self) -> _AsyncBatchBuilder: """Create the async builder and request its first batch.""" if self._async_batch_builder is not None: raise RuntimeError("Async batch building has already started.") async_batch_builder = _AsyncBatchBuilder( self._iter_synchronous(), timeout_s=self.async_batch_building_timeout_s, ) self._async_batch_builder = async_batch_builder async_batch_builder.request_batch(self.global_id) return async_batch_builder def _iter_asynchronous(self) -> Iterator[dict[str, Any]]: """Yield each completed batch while the async builder builds its replacement.""" async_batch_builder = self._async_batch_builder or self._start_async_batch_builder() while True: try: batch = async_batch_builder.get_batch() except StopIteration: break # Request one replacement before yielding. The request queue contains only a # sequence ID, and the result queue returns this exact batch object by reference. async_batch_builder.request_batch(self.global_id) yield batch def set_start_iteration(self, iteration: int) -> None: """Set the batch sequence and eagerly request the first async batch.""" if self._async_batch_builder is not None: raise RuntimeError("set_start_iteration must be called before async batch building starts.") self.global_id = iteration if self.enable_async_batch_building: self._start_async_batch_builder() def close(self) -> None: """Permanently stop the optional async batch builder.""" if self._async_batch_builder is not None: # Keep the closed builder reference: its bounded join may return while the daemon # is still blocked in data loading, so starting a replacement could race with it. self._async_batch_builder.close() def _get_dataloader_index(self, data_id: float) -> int: """Maps global id to the corresponding dataloader index based on ratio.""" for i, r in enumerate(self.data_ratios): if data_id < r: return i data_id -= r raise ValueError("Invalid data_id") class RankPartitionedDataLoader: """Assigns each rank to exactly one dataset based on ratios. For N GPUs with datasets having ratios r_1:r_2:...:r_k, the first N * r_1 / sum(r) ranks are assigned dataset 1, the next N * r_2 / sum(r) ranks are assigned dataset 2, etc. Each rank instantiates a single PyTorch DataLoader for its assigned dataset. The sharding information (``shard_world_size`` and ``shard_rank``) is set on each dataset so that it shards data only across ranks that share the same dataset, rather than across the full world. Example: With 128 GPUs and datasets ``{"video": {"dataset": ..., "ratio": 3}, "image": {"dataset": ..., "ratio": 1}}``: - Ranks 0-95 -> video (shard_world_size=96, shard_rank=0..95) - Ranks 96-127 -> image (shard_world_size=32, shard_rank=0..31) """ def __init__( self, datasets: dict[str, dict[str, Any]], **dataloader_kwargs: Any, ): """ Args: datasets: Mapping of dataset name to config dict with keys: - ``"dataset"`` (required): a lazy config or dataset instance. - ``"ratio"`` (required): positive int weight. - ``"dataloader_kwargs"`` (optional): dict of keyword arguments that override the top-level ``**dataloader_kwargs`` for this dataset only (e.g. different ``num_workers`` or ``batch_size``). **dataloader_kwargs: Default kwargs forwarded to ``torch.utils.data.DataLoader``. ``collate_fn`` defaults to ``custom_collate_fn`` if not given. """ world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() log.info(f"RankPartitionedDataLoader: world_size: {world_size} and rank: {rank}", rank0_only=False) _VALID_KEYS = {"dataset", "ratio", "dataloader_kwargs"} names: list[str] = [] dataset_configs: list[Any] = [] ratios: list[int] = [] per_dataset_kwargs: list[dict[str, Any]] = [] for name, cfg in datasets.items(): extra = set(cfg.keys()) - _VALID_KEYS assert not extra, f"Dataset {name!r}: unexpected keys {extra}. Allowed: {_VALID_KEYS}" if cfg["ratio"] <= 0: log.warning( f"RankPartitionedDataLoader: Skipping dataset {name} with ratio {cfg['ratio']}", rank0_only=False ) continue names.append(name) dataset_configs.append(cfg["dataset"]) ratios.append(cfg["ratio"]) per_dataset_kwargs.append(cfg.get("dataloader_kwargs", {})) assert len(names) > 0, "No datasets with positive ratios provided." assert world_size >= len(names), ( f"world_size ({world_size}) must be >= number of datasets ({len(names)}) " f"so each dataset gets at least one rank." ) total_ratio = sum(ratios) ideal = [r / total_ratio * world_size for r in ratios] allocations = [max(1, int(q)) for q in ideal] remaining = world_size - sum(allocations) if remaining > 0: remainders = sorted(range(len(ratios)), key=lambda i: ideal[i] - allocations[i], reverse=True) for j in range(remaining): allocations[remainders[j]] += 1 elif remaining < 0: deficit = -remaining while deficit > 0: best = max( (i for i in range(len(allocations)) if allocations[i] > 1), key=lambda i: (allocations[i] - ideal[i], allocations[i]), ) allocations[best] -= 1 deficit -= 1 expected_ratios = [r / total_ratio for r in ratios] actual_ratios = [a / world_size for a in allocations] lines = [f"RankPartitionedDataLoader allocation ({world_size} GPUs):"] start = 0 for i, (name, alloc) in enumerate(zip(names, allocations)): end = start + alloc - 1 lines.append( f" {name} (ratio {ratios[i]}): ranks {start}-{end} ({alloc} GPUs) " f"| expected {expected_ratios[i]:.2%}, actual {actual_ratios[i]:.2%}" ) start += alloc log.info("\n".join(lines), rank0_only=False) cumulative = 0 my_dataset_idx = -1 for i, alloc in enumerate(allocations): if rank < cumulative + alloc: my_dataset_idx = i break cumulative += alloc assert my_dataset_idx >= 0 shard_rank = rank - cumulative shard_world_size = allocations[my_dataset_idx] dataset: Any = instantiate(dataset_configs[my_dataset_idx]) dataset.shard_world_size = shard_world_size dataset.shard_rank = shard_rank dataset.shard_id = my_dataset_idx merged_kwargs = {**dataloader_kwargs, **per_dataset_kwargs[my_dataset_idx]} merged_kwargs.setdefault("collate_fn", custom_collate_fn) self.dataloader = torch.utils.data.DataLoader(dataset, **merged_kwargs) self.dataset_name = names[my_dataset_idx] self.dataset = dataset def __iter__(self): return iter(self.dataloader) def __len__(self) -> int: return len(self.dataloader) class PackingDataLoader(JointDataLoader): """Packs multiple samples from a single dataloader into token-budget-constrained batches. Unlike the other ``JointDataLoader`` subclasses which manage multiple dataloaders with configurable ratios, this class wraps a single dataloader and greedily packs consecutive samples until the token budget (``max_sequence_length``) or sample count limit (``max_samples_per_batch``) is reached. """ def __init__( self, dataloader: torch.utils.data.DataLoader | webdataset.WebLoader, tokenizer_spatial_compression_factor: int, tokenizer_temporal_compression_factor: int, patch_spatial: int, max_sequence_length: int | None = None, max_samples_per_batch: int | None = None, sound_latent_fps: float = 0, audio_sample_rate: int = 48000, dataset_name: str = "default", lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT, uniae_chunk_frames: int | Mapping[str, int] | None = None, uniae_pad_frames: int | None = None, ): """ Args: dataloader: A single dataloader (or lazy config) to draw samples from. tokenizer_spatial_compression_factor: Spatial compression factor of the tokenizer. tokenizer_temporal_compression_factor: Temporal compression factor of the tokenizer. patch_spatial: Spatial patchification factor. max_sequence_length: Max total tokens per packed batch. Mutually exclusive with ``max_samples_per_batch``. max_samples_per_batch: Max number of samples per packed batch. Mutually exclusive with ``max_sequence_length``. sound_latent_fps: Sound tokenizer latent rate in Hz. If 0, sound tokens are not counted. audio_sample_rate: Audio sample rate in Hz. dataset_name: Name tag attached to every sample in the output batch. lookahead_limit: Packing-loop look-ahead for the wrapped dataloader. uniae_chunk_frames: Optional UniAE full chunk size, or resolution-keyed chunk sizes. uniae_pad_frames: Optional UniAE boundary padding frames per chunk. """ wrapped = {dataset_name: {"dataloader": dataloader, "ratio": 1}} super().__init__( dataloaders=wrapped, tokenizer_spatial_compression_factor=tokenizer_spatial_compression_factor, tokenizer_temporal_compression_factor=tokenizer_temporal_compression_factor, patch_spatial=patch_spatial, max_sequence_length=max_sequence_length, max_samples_per_batch=max_samples_per_batch, sound_latent_fps=sound_latent_fps, audio_sample_rate=audio_sample_rate, lookahead_limits={dataset_name: int(lookahead_limit)}, uniae_chunk_frames=uniae_chunk_frames, uniae_pad_frames=uniae_pad_frames, ) def __iter__(self): inner = self.dataloader_list[0] ds_name = getattr(inner, "dataset_name", self.dataset_name_list[0]) while True: current_sequence_length = 0 num_samples = 0 output_batch: dict = {} skipped_samples: deque = deque() # PackingDataLoader wraps a single dataloader, so lookahead_limits has one entry. lookahead_limit = self.lookahead_limits[0] lookahead_count = 0 while True: if self.max_samples_per_batch is not None and num_samples >= self.max_samples_per_batch: break if len(output_batch) > 0 and lookahead_count >= lookahead_limit: break try: output = self._get_next_sample(0) except StopIteration: break num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output) if ( self.max_sequence_length is not None and current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length ): if len(output_batch) == 0: # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length # In this case, we should simply discard the current sample and get the next sample. log.error( f"PackingDataLoader: Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}", rank0_only=False, ) continue skipped_samples.append(output) lookahead_count += 1 continue current_sequence_length += num_tokens_in_current_sample num_samples += 1 # Allows the dataset name to be overridden by the sample itself. output["dataset_name"] = output.get("dataset_name", ds_name) self._update_output_batch(output_batch, output) for sample in reversed(skipped_samples): self.buffers[0].appendleft(sample) if len(output_batch) == 0: return self.global_id += 1 yield output_batch class RandomJointDataLoader(JointDataLoader): r""" A random joint dataloader that supports loading multiple modalities with stochastic sampling. In this dataloader, the modality is randomly selected at each iteration based on the probability distribution derived from the ratios. Each rank independently samples a modality, so different ranks may process different modalities at the same iteration. For example, with 2 modalities (image and video) and ratio 2:1: - Each iteration has 66.7% probability of selecting images - Each iteration has 33.3% probability of selecting videos - The selection is independent across iterations and ranks Note: Unlike IterativeJointDataLoader, this does not guarantee synchronized modality selection across ranks. """ def __init__( self, dataloaders: Dict[str, Dict[str, Union[torch.utils.data.DataLoader, webdataset.WebLoader, int]]], tokenizer_spatial_compression_factor: int, tokenizer_temporal_compression_factor: int, patch_spatial: int, max_sequence_length: int | None = None, max_samples_per_batch: int | None = None, sound_latent_fps: float = 0, audio_sample_rate: int = 48000, default_lookahead_limit: int = JointDataLoader._DEFAULT_LOOKAHEAD_LIMIT, lookahead_limits: Dict[str, int] | None = None, uniae_chunk_frames: int | Mapping[str, int] | None = None, uniae_pad_frames: int | None = None, ): super().__init__( dataloaders, tokenizer_spatial_compression_factor, tokenizer_temporal_compression_factor, patch_spatial, max_sequence_length, max_samples_per_batch, sound_latent_fps=sound_latent_fps, audio_sample_rate=audio_sample_rate, default_lookahead_limit=default_lookahead_limit, lookahead_limits=lookahead_limits, uniae_chunk_frames=uniae_chunk_frames, uniae_pad_frames=uniae_pad_frames, ) # Convert data ratios to probabilities self.data_ratios = np.array([ratio / sum(self.data_ratios) for ratio in self.data_ratios]) def __iter__(self): while True: index_id = np.random.choice(len(self.dataloader_list), p=self.data_ratios) metrics = _PackingMetrics() output_batch = dict() skipped_samples = deque() lookahead_limit = self.lookahead_limits[index_id] lookahead_count = 0 while True: # Check max samples limit first if self.max_samples_per_batch is not None and metrics.num_samples >= self.max_samples_per_batch: break # If we have started packing and tried lookahead_limit times to find a fitting sample but failed, stop. if len(output_batch) > 0 and lookahead_count >= lookahead_limit: break had_buffer = len(self.buffers[index_id]) > 0 try: output = self._get_next_sample(index_id) except StopIteration: break # No more data in this dataloader if had_buffer: metrics.from_buffer += 1 else: metrics.from_workers += 1 num_tokens_in_current_sample = self._compute_num_tokens_per_sample(output) if ( self.max_sequence_length is not None and metrics.current_sequence_length + num_tokens_in_current_sample >= self.max_sequence_length ): if len(output_batch) == 0: # This case happens when current_sequence_length = 0 and num_tokens_in_current_sample > self.max_sequence_length # In this case, we should simply discard the current sample and get the next sample. log.info( f"Discarding oversized sample with {num_tokens_in_current_sample} tokens. Max sequence length: {self.max_sequence_length}", rank0_only=False, ) metrics.dropped_count += 1 continue # current_sequence_length > 0 and selected sample is too large to fit in the remaining space. # Instead of stopping immediately (creating large padding), we buffer this large sample # and try to find a smaller one that fits in the remaining space. skipped_samples.append(output) lookahead_count += 1 continue metrics.current_sequence_length += num_tokens_in_current_sample metrics.num_samples += 1 output["dataset_name"] = self.dataset_name_list[index_id] self._update_output_batch(output_batch, output) # Add back skipped samples to the buffer for the next batch. # appendleft puts item at HEAD. So we insert S3, then S2, then S1. for sample in reversed(skipped_samples): self.buffers[index_id].appendleft(sample) if len(output_batch) == 0: return metrics.attach_to(output_batch, buffer_size=len(self.buffers[index_id])) yield output_batch