# Copyright 2025 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import os import random import traceback from functools import partial from typing import Any, Callable, Dict, List, Literal, Optional, Sequence import numpy as np import torch from datasets import IterableDataset as HFIterableDataset from datasets import interleave_datasets, load_dataset from datasets.distributed import split_dataset_by_node from huggingface_hub import hf_hub_download from torch.utils.data import Dataset, IterableDataset, get_worker_info from ..utils.registry import Registry try: from hdfs_io import isdir, listdir except ImportError: from ..utils.hdfs_io import isdir, listdir from ..distributed.parallel_state import get_parallel_state from ..utils import logging from ..utils.dist_utils import main_process_first from ..utils.multisource_utils import parse_multisource_config logger = logging.get_logger(__name__) DATASET_REGISTRY = Registry("Dataset") def build_dataset(dataset_name: str, **kwargs) -> "Dataset": return DATASET_REGISTRY[dataset_name](**kwargs) class MappingDataset(Dataset): def __init__(self, data: "Dataset", transform: Optional[Callable] = None): self._data = data self._transform = transform self.indices = list(range(len(self._data))) self.data_len = len(self.indices) def __len__(self) -> int: return self.data_len def __getitem__(self, index: int) -> List[Dict[str, "torch.Tensor"]]: if index >= len(self.indices): random.shuffle(self.indices) index = index % len(self.indices) mapped_idx = self.indices[index] if self._transform is not None: return self._transform(self._data[mapped_idx]) else: return self._data[mapped_idx] class IterativeDataset(IterableDataset): def __init__(self, data: "HFIterableDataset", transform: Optional[Callable] = None): self._data = data self._transform = transform def __iter__(self): for sample in self._data: if self._transform is not None: yield self._transform(sample) else: yield sample def load_state_dict(self, state_dict): self._data.load_state_dict(state_dict["dataset"]) def state_dict(self): return {"dataset": self._data.state_dict()} def set_epoch(self, epoch: int): self._data.set_epoch(epoch) class InterleavedIterableDataset(IterativeDataset): def __init__(self, data: "HFIterableDataset", transform: Optional[Callable] = None): self._data = data self._transform = transform def __iter__(self): for sample in self._data: if self._transform is not None: ds_idx = sample["ds_idx"] source_name = sample.get("source_name", None) transformed_sample = self._transform(sample, source_name=source_name) if isinstance(transformed_sample, List): for idx in range(len(transformed_sample)): transformed_sample[idx]["ds_idx"] = ds_idx yield transformed_sample else: transformed_sample["ds_idx"] = ds_idx yield transformed_sample else: yield sample class InterleavedMappingDataset(MappingDataset): def __init__(self, data: "Dataset", transform: Optional[Callable] = None): super().__init__(data, transform) def __getitem__(self, index: int) -> List[Dict[str, "torch.Tensor"]]: if index >= len(self.indices): random.shuffle(self.indices) index = index % len(self.indices) mapped_idx = self.indices[index] if self._transform is not None: sample = self._data[mapped_idx] ds_idx = sample["ds_idx"] source_name = sample.get("source_name", None) transformed_sample = self._transform(sample, source_name=source_name) if isinstance(transformed_sample, List): for idx in range(len(transformed_sample)): transformed_sample[idx]["ds_idx"] = ds_idx else: transformed_sample["ds_idx"] = ds_idx return transformed_sample else: return self._data[mapped_idx] class EnergonDataset(IterativeDataset): """ A specialized wrapper for Megatron-Energon datasets that provides: - Automatic WorkerConfig management - TextSample to dict conversion - Native state management using save_state/restore_state - Epoch-based state reset Args: data (Dataset): underlying Megatron-Energon dataset transform (Optional[Callable]): transform function """ def __init__(self, data: "Dataset", transform: Optional[Callable] = None): self._data = data self._transform = transform def __len__(self): """Get the length of the dataset.""" if hasattr(self._data, "__len__"): return len(self._data) def __iter__(self): """Iterate over the dataset with WorkerConfig management and TextSample conversion.""" # For Megatron-Energon datasets, we need to set up the WorkerConfig properly if hasattr(self._data, "worker_config"): try: from megatron.energon import WorkerConfig # Ensure active_worker_config is None before activation WorkerConfig.active_worker_config = None # Activate the worker config self._data.worker_config.worker_activate(sample_index=0) logger.debug("Activated WorkerConfig for Megatron-Energon dataset") except Exception as e: logger.warning(f"Failed to activate WorkerConfig: {e}") try: for sample in self._data: # Convert Megatron-Energon TextSample to dict for compatibility if hasattr(sample, "__dict__") and not isinstance(sample, dict): # Convert TextSample or similar objects to dict sample_dict = {} for key, value in sample.__dict__.items(): if not key.startswith("_"): # Skip private attributes sample_dict[key] = value # Handle special case for TextSample if hasattr(sample, "text"): sample_dict["text"] = sample.text sample = sample_dict if self._transform is not None: yield self._transform(sample) else: yield sample finally: # Clean up WorkerConfig if hasattr(self._data, "worker_config"): try: self._data.worker_config.worker_deactivate() logger.debug("Deactivated WorkerConfig for Megatron-Energon dataset") except Exception as e: logger.warning(f"Failed to deactivate WorkerConfig: {e}") def load_state_dict(self, state_dict): """Load the state of the dataset from checkpointing.""" if hasattr(self._data, "restore_state"): # Use Megatron-Energon's native restore_state method try: self._data.restore_state(state_dict["dataset"]) except Exception as e: logger.warning(f"Failed to restore state using restore_state: {e}") elif hasattr(self._data, "load_state_dict"): # Fallback to load_state_dict if available self._data.load_state_dict(state_dict["dataset"]) else: logger.warning(f"Dataset {type(self._data).__name__} does not support state restoration") def state_dict(self): """Get the state of the dataset for checkpointing.""" if hasattr(self._data, "save_state"): # Use Megatron-Energon's native save_state method try: state = self._data.save_state() return {"dataset": state} except Exception as e: logger.warning(f"Failed to save state using save_state: {e}") return {"dataset": {}} elif hasattr(self._data, "state_dict"): # Fallback to state_dict if available return {"dataset": self._data.state_dict()} else: # Return empty state dict for datasets that don't support state management return {"dataset": {}} def set_epoch(self, epoch: int): """Set the epoch for the dataset.""" if hasattr(self._data, "set_epoch"): self._data.set_epoch(epoch) elif hasattr(self._data, "reset_state_deep"): # For Megatron-Energon datasets, reset state when epoch changes try: self._data.reset_state_deep() logger.debug(f"Reset state for epoch {epoch}") except Exception as e: logger.warning(f"Failed to reset state for epoch {epoch}: {e}") else: logger.debug(f"Dataset {type(self._data).__name__} does not support set_epoch or state reset") def get_length_by_attention_mask_fn(sample): """Return token length from ``attention_mask``. Defined as a top-level helper instead of an inline lambda because ``spawn`` worker mode requires the callable to be pickleable. """ return int(sample["attention_mask"].sum()) def _supports_output_index_for_resume(dataset: Any) -> bool: return callable(getattr(dataset, "get_item", None)) and hasattr(dataset, "output_index_for_resume") class WeightedMultiSourceDataset(IterableDataset): """Multi-source dataset with weighted sampling. This dataset samples from multiple upstream iterable datasets according to a (possibly token-adjusted) weight distribution. It supports: - Per-epoch deterministic randomness (seeded by epoch, dp rank, and worker id). - Optional distributed sharding behavior controlled by ``upstream_sharded``. - Stopping strategies for how to behave when an upstream source is exhausted. - Optional resume-index passthrough for checkpointing buffers by index. """ def __init__( self, datasets: Sequence[IterableDataset], weights: Sequence[float], seed: int = 42, level: Literal["sample", "token"] = "sample", sample_token_len_fn: Optional[Callable[[Any], float]] = None, source_names: Optional[Sequence[str]] = None, source_ids: Optional[Sequence[str]] = None, upstream_sharded: bool = False, stopping_strategy: Literal["first_exhausted", "all_exhausted", "never_exhausted"] = "first_exhausted", output_index_for_resume: bool = False, ) -> None: """Initialize a WeightedMultiSourceDataset. Args: datasets: Upstream iterable datasets (one per source). weights: Sampling weights aligned with ``datasets``. seed: Base random seed. level: Sampling level. ``sample`` uses ``weights`` directly; ``token`` reweights by the inverse of the running average token length per source. sample_token_len_fn: Function that returns the token length of a sample. If not provided, a default heuristic is used. source_names: Optional display names for each source (for meta fields). source_ids: Optional stable IDs for each source (used in checkpoint state). upstream_sharded: If False, performs deterministic modulo-based sharding by dp rank on the produced samples. If True, assumes upstream datasets already handle sharding/splitting. stopping_strategy: - ``first_exhausted``: Stop the whole dataset once any source is exhausted. - ``all_exhausted``: Restart an exhausted source until all sources are exhausted. - ``never_exhausted``: Always restart exhausted sources and never terminate. output_index_for_resume: If True, yields ``(sample, (source_id, output_index))`` so downstream components can checkpoint buffers by output indices and reconstruct them later. Raises: ValueError: If input arguments are invalid. """ self._datasets = list(datasets) self._weights = np.asarray(weights, dtype=np.float64) self._seed = seed self._level = level self._sample_token_len_fn = sample_token_len_fn or self._default_sample_token_len self._source_names = list(source_names) if source_names is not None else None self._source_ids = list(source_ids) if source_ids is not None else [] self._upstream_sharded = upstream_sharded self._stopping_strategy = stopping_strategy self._ds_num = len(self._datasets) if not self._source_names: self._source_names = [] for i, dataset in enumerate(self._datasets): if callable(getattr(dataset, "get_name", None)): self._source_names.append(dataset.get_name()) else: self._source_names.append(f"source_{i}") if not self._source_ids: self._source_ids = copy.deepcopy(self._source_names) self._id2dataset = { source_id: (dataset, ds_idx) for ds_idx, (source_id, dataset) in enumerate(zip(self._source_ids, self._datasets)) } self._avg_len_sum = [0.0 for _ in range(self._ds_num)] self._avg_len_count = [0 for _ in range(self._ds_num)] self._global_sample_idx = 0 self._random_state = np.random.RandomState(seed=self._seed) self._iters: List[Any] = [] self._epoch = 0 self._exhausted = [False for _ in range(self._ds_num)] if self._weights.shape[0] != self._ds_num: raise ValueError("weights length must match datasets length") if self._source_names is not None and len(self._source_names) != self._ds_num: raise ValueError("source_names length must match datasets length") if len(self._source_ids) != self._ds_num: raise ValueError("source_ids length must match datasets length") if len(set(self._source_ids)) != self._ds_num: raise ValueError("source_ids must be unique") if self._level not in ("sample", "token"): raise ValueError("level must be 'sample' or 'token'") if self._stopping_strategy not in ("first_exhausted", "all_exhausted", "never_exhausted"): raise ValueError("stopping_strategy must be 'first_exhausted', 'all_exhausted', or 'never_exhausted'") parallel_state = get_parallel_state() self.dp_rank = max(0, int(getattr(parallel_state, "dp_rank", 0))) self.dp_size = max(1, int(getattr(parallel_state, "dp_size", 1))) self.output_index_for_resume = output_index_for_resume self._just_resumed = False @property def output_index_for_resume(self) -> bool: """Whether to yield output indices alongside samples for resume.""" return self._output_index_for_resume @output_index_for_resume.setter def output_index_for_resume(self, value: bool) -> None: """Enable or disable output-index emission for resume. When enabled, each upstream dataset must provide: - ``get_item(idx)`` to fetch a sample by index - ``output_index_for_resume`` attribute to switch yielding ``(sample, idx)`` Args: value: True to emit output indices for resume, False to disable. Raises: ValueError: If any upstream dataset cannot emit output indices for resume. """ if value: for source_id, (dataset, _ds_idx) in self._id2dataset.items(): if not _supports_output_index_for_resume(dataset): raise ValueError( f"output_index_for_resume is True, but dataset '{source_id}' does not have " f"get_item method or output_index_for_resume attribute to resume samples " f"in buffers based on idx" ) self._output_index_for_resume = value for dataset in self._datasets: if hasattr(dataset, "output_index_for_resume"): dataset.output_index_for_resume = value def get_item(self, resume_index): """Fetch a single sample by its source ID and index within that source. This is used by downstream checkpoint/resume logic that stores buffer contents as ``(source_id, idx)`` pairs instead of full samples. Args: resume_index: A ``(source_id, idx)`` tuple. ``source_id`` identifies the sub-dataset, and ``idx`` is the 0-based index within that sub-dataset. Returns: The sample returned by the underlying sub-dataset. Raises: AttributeError: If the underlying sub-dataset does not provide an index-based fetch API. """ source_id, idx = resume_index dataset, ds_idx = self._id2dataset[source_id] get_item_fn = getattr(dataset, "get_item", None) if callable(get_item_fn): sample = get_item_fn(idx) sample = self._attach_meta(sample, ds_idx) return sample raise AttributeError(f"dataset '{source_id}' does not implement get_item") def set_epoch(self, epoch: int) -> None: """Set the epoch for deterministic sampling. Args: epoch: Current epoch number. """ self._epoch = epoch for dataset in self._datasets: set_epoch_fn = getattr(dataset, "set_epoch", None) if callable(set_epoch_fn): set_epoch_fn(epoch) def __iter__(self): """Iterate and yield samples from multiple sources. Yields: If ``output_index_for_resume`` is False, yields a sample. If ``output_index_for_resume`` is True, yields ``(sample, (source_id, output_index))``. """ worker_info = get_worker_info() worker_id = worker_info.id if worker_info is not None else 0 if not self._just_resumed: seed_seq = np.random.SeedSequence([self._seed, self._epoch, self.dp_rank, worker_id]) current_seed = int(seed_seq.generate_state(1, dtype=np.uint32)[0]) self._random_state = np.random.RandomState(current_seed) self._exhausted = [False for _ in range(self._ds_num)] self._avg_len_sum = [0.0 for _ in range(self._ds_num)] self._avg_len_count = [0 for _ in range(self._ds_num)] self._global_sample_idx = 0 else: self._just_resumed = False self._iters = [iter(ds) for ds in self._datasets] while True: ds_idx = self._random_state.choice(self._ds_num, p=self._runtime_weights()) try: sample = self._next_sample(ds_idx) except StopIteration: return if sample is None: continue if self._output_index_for_resume: sample, output_index = sample[0], sample[1] token_len = self._sample_token_len_fn(sample) if token_len <= 0: continue if self._level == "token": self._avg_len_sum[ds_idx] += token_len self._avg_len_count[ds_idx] += 1 self._global_sample_idx += 1 if not self._upstream_sharded and self._global_sample_idx % self.dp_size != self.dp_rank: continue sample = self._attach_meta(sample, ds_idx) if self._output_index_for_resume: yield sample, (self._source_ids[ds_idx], output_index) else: yield sample def _runtime_weights(self) -> np.ndarray: """Compute the per-source sampling probabilities for the current runtime state. Returns: A probability vector of shape ``(num_sources,)`` that sums to 1. Raises: ValueError: If the weight sum is non-positive. """ if self._level == "sample": weights = self._weights else: avg_lens = [] for idx in range(self._ds_num): if self._avg_len_count[idx] > 0: avg_lens.append(self._avg_len_sum[idx] / self._avg_len_count[idx]) else: avg_lens.append(1.0) weights = self._weights / np.asarray(avg_lens, dtype=np.float64) total = float(np.sum(weights)) if total <= 0: raise ValueError("sum of weights must be positive") return weights / total def _next_sample(self, ds_idx: int) -> Any: """Fetch the next sample from a specific sub-dataset index. Args: ds_idx: Index of the sub-dataset to fetch from. Returns: The next sample from the chosen sub-dataset. Raises: StopIteration: When the dataset terminates under the configured stopping strategy. """ while True: try: return next(self._iters[ds_idx]) except StopIteration: if self._stopping_strategy == "first_exhausted": raise if self._stopping_strategy == "all_exhausted": self._exhausted[ds_idx] = True if all(self._exhausted): raise elif self._stopping_strategy == "never_exhausted": self._exhausted[ds_idx] = True if all(self._exhausted): self._exhausted = [False for _ in range(self._ds_num)] logger.warning( f"Data source #{ds_idx} (source_name: {self._source_names[ds_idx]}) is exhausted, reset and continue" ) self._iters[ds_idx] = iter(self._datasets[ds_idx]) try: return next(self._iters[ds_idx]) except StopIteration as e: raise RuntimeError( f"Data source #{ds_idx} (source_name: {self._source_names[ds_idx]}) remains exhausted " "immediately after reset" ) from e def _attach_meta(self, sample: Any, ds_idx: int) -> Any: """Attach per-source metadata fields onto a sample. Adds: - ``ds_idx``: the integer source index - ``source_name``: optional display name if provided Args: sample: A sample or list of samples. ds_idx: Source index for this sample. Returns: The updated sample (mutated in place when possible). """ source_name = self._source_names[ds_idx] if self._source_names is not None else None if isinstance(sample, list): for item in sample: if isinstance(item, dict): item["ds_idx"] = ds_idx if source_name is not None: item["source_name"] = source_name return sample if isinstance(sample, dict): sample["ds_idx"] = ds_idx if source_name is not None: sample["source_name"] = source_name return sample def _default_sample_token_len(self, sample: Any) -> float: """Default heuristic to estimate token length of a sample. Args: sample: A single sample or a list of samples. Returns: Estimated token length as a float. """ if sample is None: return 0 if isinstance(sample, list): return float(sum(self._default_sample_token_len(item) for item in sample)) if not isinstance(sample, dict): return 1.0 if "attention_mask" in sample: attention_mask = sample["attention_mask"] if isinstance(attention_mask, torch.Tensor): return float(attention_mask.sum().item()) if isinstance(attention_mask, list): return float(sum(attention_mask)) if "input_ids" in sample: input_ids = sample["input_ids"] if isinstance(input_ids, torch.Tensor): return float(input_ids.numel()) if isinstance(input_ids, list): return float(len(input_ids)) return 1.0 def state_dict(self) -> dict: """Return a checkpointable state dict for this dataset.""" dataset_states_by_id = {} for dataset, source_id in zip(self._datasets, self._source_ids): state_fn = getattr(dataset, "state_dict", None) getstate_fn = getattr(dataset, "__getstate__", None) if callable(state_fn): ds_state = state_fn() elif callable(getstate_fn): ds_state = getstate_fn() else: ds_state = None dataset_states_by_id[source_id] = ds_state avg_len_sum_by_id = {source_id: self._avg_len_sum[idx] for idx, source_id in enumerate(self._source_ids)} avg_len_count_by_id = {source_id: self._avg_len_count[idx] for idx, source_id in enumerate(self._source_ids)} # save _exhausted state exhausted_by_id = {source_id: self._exhausted[idx] for idx, source_id in enumerate(self._source_ids)} return { "version": 0, "topology": { "source_ids": list(self._source_ids), "source_names": list(self._source_names) if self._source_names is not None else None, "weights": self._weights.tolist(), "level": self._level, "stopping_strategy": self._stopping_strategy, }, "runtime": { "random_state": self._random_state.get_state(), "avg_len_sum": avg_len_sum_by_id, "avg_len_count": avg_len_count_by_id, "exhausted": exhausted_by_id, "global_sample_idx": self._global_sample_idx, "dataset_states": dataset_states_by_id, }, } def load_state_dict( self, state: dict, reconcile_policy: Literal["strict", "allow_add", "allow_add_remove", "warn_only"] = "allow_add_remove", ) -> None: """Restore state from a previous ``state_dict()``. Args: state: State dict previously produced by ``state_dict()``. reconcile_policy: Policy for handling source-id changes: - ``strict``: error on any added/removed source. - ``allow_add``: allow new sources but error on removed ones. - ``allow_add_remove``: allow both add and remove. - ``warn_only``: allow changes and log a warning. Raises: ValueError: If required state fields are missing or incompatible. """ if "topology" not in state or "runtime" not in state: raise ValueError("state_dict missing required keys: topology/runtime") runtime = state["runtime"] topology = state["topology"] if "source_ids" not in topology: raise ValueError("state_dict missing topology.source_ids") saved_source_ids = topology["source_ids"] added = [] removed = [] if saved_source_ids is not None: saved_set = set(saved_source_ids) added = [source_id for source_id in self._source_ids if source_id not in saved_set] removed = [source_id for source_id in saved_source_ids if source_id not in set(self._source_ids)] if added or removed: if reconcile_policy == "strict": raise ValueError( f"source_ids mismatch: added={added} removed={removed} with policy={reconcile_policy}" ) if reconcile_policy == "allow_add" and removed: raise ValueError( f"source_ids removed not allowed: removed={removed} with policy={reconcile_policy}" ) if reconcile_policy == "warn_only": logger.warning( f"source_ids changed: added={added} removed={removed} with policy={reconcile_policy}" ) random_state = runtime["random_state"] self._random_state.set_state(random_state) avg_len_sum = runtime["avg_len_sum"] avg_len_count = runtime["avg_len_count"] if not isinstance(avg_len_sum, dict) or not isinstance(avg_len_count, dict): raise ValueError("runtime.avg_len_sum and runtime.avg_len_count must be dicts keyed by source_id") self._avg_len_sum = [float(avg_len_sum.get(source_id, 0.0)) for source_id in self._source_ids] self._avg_len_count = [int(avg_len_count.get(source_id, 0)) for source_id in self._source_ids] self._global_sample_idx = runtime.get("global_sample_idx", 0) dataset_states = runtime["dataset_states"] if not isinstance(dataset_states, dict): raise ValueError("runtime.dataset_states must be a dict keyed by source_id") dataset_states_by_id = dataset_states for dataset, source_id in zip(self._datasets, self._source_ids): ds_state = dataset_states_by_id.get(source_id) if ds_state is None: continue load_state_fn = getattr(dataset, "load_state_dict", None) if callable(load_state_fn): load_state_fn(ds_state) # Ensure _exhausted is re-initialized for the current source count # This is important when sources are added/removed during checkpoint resume if "exhausted" in runtime and isinstance(runtime["exhausted"], dict): exhausted_dict = runtime["exhausted"] self._exhausted = [bool(exhausted_dict.get(source_id, False)) for source_id in self._source_ids] else: self._exhausted = [False for _ in range(self._ds_num)] self._just_resumed = True class DynamicBatchingSizeDataset(IterableDataset): """Dynamic batching dataset that yields micro batches based on token count. Unlike ``DynamicBatchSizeDataLoader``, which constructs micro batches in the main process after fetching from a plain DataLoader, ``DynamicBatchingSizeDataset`` performs batching inside each DataLoader worker process. It is also compatible with ``StatefulDataLoader``'s per-worker ``state_dict()`` / ``load_state_dict()`` mechanism, enabling exact checkpoint / resume for dynamic-batching workloads. Internally each worker maintains a sample buffer. A micro batch is emitted once the buffer holds at least ``ready_for_micro_batch_threshold`` samples **and** their combined token count reaches ``micro_batch_seq_length``. When the upstream dataset is exhausted, remaining buffer contents are drained and emitted as final batches regardless of the threshold. Attributes: dataset: The upstream iterable dataset to read samples from. ready_for_micro_batch_threshold: Minimum number of samples that must be in the buffer before a microbatch can be formed. micro_batch_seq_length: Target total token count per micro batch (soft upper bound; may be exceeded by a single overlong sample when ``force_generate_long_sequence`` is True). get_length_fn: Function that returns the token count of a single sample. save_by_idx: Whether to checkpoint the buffer as sample indices (smaller checkpoint size) rather than full sample tensors. force_generate_long_sequence: If True, a sample whose length alone exceeds ``micro_batch_seq_length`` is emitted as a single-sample batch instead of being silently discarded. This is not supported yet. """ def __init__( self, dataset: IterableDataset, micro_batch_seq_length: int, ready_for_micro_batch_threshold: int, dynamic_batching_collate_fn: Callable, save_by_idx: bool = True, get_length_fn: Optional[Callable] = get_length_by_attention_mask_fn, force_generate_long_sequence: bool = False, ) -> None: """Initialize the DynamicBatchingSizeDataset. Args: dataset: The underlying iterable dataset to batch from. micro_batch_seq_length: Target total token count per micro batch. ready_for_micro_batch_threshold: Minimum number of samples required in buffer before attempting to create a batch. save_by_idx: If True, saves sample indices for checkpoint resumption. Requires dataset to have get_item method and output_index_for_resume attribute. get_length_fn: Function to compute the length (token count) of a sample. Defaults to len. force_generate_long_sequence: If True, a sample whose length alone exceeds ``micro_batch_seq_length`` is emitted as a single-sample batch instead of being silently discarded. This is not supported yet. Resume flow when ``save_by_idx=True``:: Runtime path ------------ DynamicBatchingSizeDataset +- sets dataset.output_index_for_resume = True +- reads from dataset.__iter__() | `- WeightedMultiSourceDataset.__iter__() | +- reads from source dataset 'zh'.__iter__() | +- gets inner_index = 17 from source dataset 'zh' | +- sample comes from source_dataset['zh'].get_item(17) | | `- [ | | {'input_ids': [11, 22], 'attention_mask': [1, 1]}, | | {'input_ids': [33], 'attention_mask': [1]}, | | ] | `- yields ([...], ('zh', 17)) `- keeps +- runtime buffer: [ | ({'input_ids': [11, 22], 'attention_mask': [1, 1]}, 2), | ({'input_ids': [33], 'attention_mask': [1]}, 1), | ] `- checkpoint buffer: [(('zh', 17), 0), (('zh', 17), 1)] Resume path ----------- checkpoint['buffer'] `- [(('zh', 17), 0), (('zh', 17), 1)] `- load_state_dict() +- WeightedMultiSourceDataset.get_item(('zh', 17)) | +- 'zh' -> select source dataset 'zh' | `- 17 -> source_dataset['zh'].get_item(17) `- select sample_idx=1 from the returned list Raises: ValueError: If ``save_by_idx`` is True but ``dataset`` does not expose the ``get_item()`` method and ``output_index_for_resume`` attribute required to reconstruct the buffer from indices on resume. """ if not isinstance(dataset, IterableDataset): raise TypeError( f"DynamicBatchingSizeDataset does not support Mapping style datasets now, the dataset's type must be IterableDataset, got {type(dataset).__name__}" ) self.dataset = dataset self.dynamic_batching_collate_fn = dynamic_batching_collate_fn self.ready_for_micro_batch_threshold = ready_for_micro_batch_threshold self.micro_batch_seq_length = micro_batch_seq_length self.get_length_fn = get_length_fn self.save_by_idx = save_by_idx if force_generate_long_sequence: raise ValueError("force_generate_long_sequence is not supported yet.") self.force_generate_long_sequence = force_generate_long_sequence self._buffer = [] self._buffer_of_output_index = [] self._buffer_token_count = 0 self._just_resumed = False # Flag to indicate if the dataset has just been resumed from a checkpoint, used to skip buffer checks on the first iteration after resume. @property def save_by_idx(self) -> bool: return self._save_by_idx @save_by_idx.setter def save_by_idx(self, value: bool) -> None: if value and not _supports_output_index_for_resume(self.dataset): raise ValueError( "save_by_idx is True, but dataset does not have get_item method or output_index_for_resume attribute to resume samples in buffers based on idx" ) self._save_by_idx = value if hasattr(self.dataset, "output_index_for_resume"): self.dataset.output_index_for_resume = value def __iter__(self): """Iterate over the dataset and yield dynamically batched micro batches. Buffers samples from the underlying dataset and yields micro batches when the buffer contains enough samples and tokens. Each yielded batch is collated using the dynamic_batching_collate_fn. Yields: Collated micro batch when buffer conditions are met. Raises: Exception: Re-raises any exception other than StopIteration encountered during iteration. """ self._data_iter = iter(self.dataset) if not self._just_resumed: # Clear buffer state on new iteration unless we just resumed from a checkpoint, # in which case we want to keep the buffer contents. self._buffer = [] self._buffer_of_output_index = [] self._buffer_token_count = 0 else: self._just_resumed = False while True: try: if ( len(self._buffer) >= self.ready_for_micro_batch_threshold and self._buffer_token_count >= self.micro_batch_seq_length ): micro_batch = self._get_micro_batch() micro_batch = self.dynamic_batching_collate_fn(micro_batch) if micro_batch is not None: yield micro_batch else: logger.warning("dynamic_batching_collate_fn returned None, skip this micro_batch") item = next(self._data_iter) if self.save_by_idx: item, output_index = item else: output_index = None samples_to_add = item if isinstance(item, list) else [item] for sample_idx, sample in enumerate(samples_to_add): length = self.get_length_fn(sample) if length > self.micro_batch_seq_length and not self.force_generate_long_sequence: # TODO: record the count of discarded long examples for monitoring logger.warning( f"Sample length {length} exceeds micro batch seq length {self.micro_batch_seq_length}, skipping. If you want to force generate a micro batch with this sample, enable force_generate_long_sequence." ) continue self._buffer.append((sample, length)) if self.save_by_idx: # Save one output-index entry per buffered sample. # An upstream dataset may yield ``list[dict]`` in one # iteration, and ``sample_idx`` selects the buffered # sample within that list during resume. self._buffer_of_output_index.append((output_index, sample_idx)) self._buffer_token_count += length except Exception as e: if isinstance(e, StopIteration): while len(self._buffer) > 0: micro_batch = self._get_micro_batch() micro_batch = self.dynamic_batching_collate_fn(micro_batch) if micro_batch is not None: yield micro_batch else: logger.warning("dynamic_batching_collate_fn returned None, skip this micro_batch") return else: logger.error(f"DynamicBatchDataset iter data exception: {e} \n{traceback.format_exc()}") raise def _get_micro_batch(self): """Construct a micro batch from buffered samples using a greedy first-fit strategy. Iterates the buffer in order and greedily adds each sample whose length fits within the remaining token budget (``micro_batch_seq_length - seq_length``). Samples that do not fit are left in the buffer for subsequent batches. Special case: when the buffer's first sample alone exceeds ``micro_batch_seq_length`` and ``force_generate_long_sequence`` is True, that sample is taken unconditionally (``seq_length == 0`` guard) so that the dataset never stalls on an overlong sequence. Returns: list: Non-empty list of samples forming the micro batch. Raises: AssertionError: If no sample could be selected (should never happen under normal operation). """ micro_batch = [] seq_length = 0 indices_to_remove_from_buffer = [] for idx, item in enumerate(self._buffer): sample, length = item[0], item[1] if length + seq_length > self.micro_batch_seq_length: if seq_length > 0: continue elif not self.force_generate_long_sequence: # Usually it is impossible to reach this branch because too long samples would not be added to the buffer if force_generate_long_sequence is False. continue micro_batch.append(sample) seq_length += length self._buffer_token_count -= length indices_to_remove_from_buffer.append(idx) if seq_length >= self.micro_batch_seq_length: break # Remove selected items from buffer (iterate backwards to maintain indices) for idx in reversed(indices_to_remove_from_buffer): del self._buffer[idx] if self.save_by_idx: del self._buffer_of_output_index[idx] assert len(micro_batch) > 0 return micro_batch def state_dict(self): """Get the state dictionary for checkpointing. Saves the current buffer state and token count. If save_by_idx is True, only saves sample indices; otherwise saves the full buffer contents. Also saves the upstream dataset state if available. Returns: dict: State dictionary containing: - save_by_idx: Whether indices are saved instead of samples. - buffer_token_count: Total token count in the buffer. - buffer: Buffered samples or their indices. - dynamic_batch_upstream_dataset_state: Upstream dataset state (if available). """ state = { "save_by_idx": self.save_by_idx, # Make sure we store an integer instead of any tensor "buffer_token_count": int(self._buffer_token_count), } # the state_dict might be called frequently with StatefulDataloaders(see more details of snapshot_every_n_steps) # so we try to not include extra calculations here. if self.save_by_idx: state["buffer"] = copy.deepcopy(self._buffer_of_output_index) else: # deepcopy buffer so that it can be transfered through multiple processes state["buffer"] = copy.deepcopy(self._buffer) if hasattr(self.dataset, "state_dict"): state["dynamic_batch_upstream_dataset_state"] = self.dataset.state_dict() return state def load_state_dict(self, state_dict): """Load state from a checkpoint. Restores the buffer and token count from a saved state. Handles both index-based and full-sample buffer restoration based on the saved state. Also restores the upstream dataset state if available. Args: state_dict: State dictionary from a previous checkpoint, containing: - save_by_idx: Whether the saved buffer contains indices. - buffer: Saved buffer (samples or indices). - buffer_token_count: Saved token count. - dynamic_batch_upstream_dataset_state: Upstream dataset state (optional). Raises: AssertionError: If the restored ``buffer_token_count`` does not match the sum of token lengths recomputed from the reconstructed buffer. ValueError: If ``save_by_idx`` is True on the current instance but the checkpoint buffer holds some full samples instead of indices (incompatible checkpoint format). """ # prev_save_by_idx does not have to be equal to self.save_by_idx, however, we still need to resume the buffer according to it. prev_save_by_idx = state_dict["save_by_idx"] if prev_save_by_idx: self._buffer = [] self._buffer_of_output_index = [] cached_output_index = None cached_restored_samples = None for output_index_entry in state_dict["buffer"]: # Each checkpoint entry points to exactly one buffered sample: # ``output_index`` identifies the upstream item and ``sample_idx`` # selects one sample after flattening a possible ``list[dict]``. output_index, sample_idx = output_index_entry if output_index != cached_output_index: restored_item = self.dataset.get_item(output_index) cached_restored_samples = restored_item if isinstance(restored_item, list) else [restored_item] cached_output_index = output_index restored_sample = cached_restored_samples[sample_idx] length = self.get_length_fn(restored_sample) self._buffer.append((restored_sample, length)) if self.save_by_idx: self._buffer_of_output_index.append(output_index_entry) self._buffer_token_count += length else: self._buffer = state_dict["buffer"] if self.save_by_idx and len(self._buffer) > 0: raise ValueError("save_by_idx is True, but previous buffer contains valid samples instead of indices") self._buffer_of_output_index = [] self._buffer_token_count = state_dict["buffer_token_count"] # Verify buffer_token_count matches the sum of token lengths assert self._buffer_token_count == sum([item[1] for item in self._buffer]), ( "buffer_token_count does not match the sum of token lengths in buffer" ) assert self._buffer_token_count == sum(self.get_length_fn(item[0]) for item in self._buffer), ( "buffer_token_count does not match the sum of lengths computed from samples in buffer" ) del state_dict["buffer"] if "dynamic_batch_upstream_dataset_state" in state_dict: self.dataset.load_state_dict(state_dict["dynamic_batch_upstream_dataset_state"]) self._just_resumed = True def set_epoch(self, epoch: int): """Set the epoch for the upstream dataset. Passes the epoch to the upstream dataset if it supports set_epoch. Has no direct effect on dynamic batching itself. Args: epoch: The epoch number to set. """ if hasattr(self.dataset, "set_epoch"): self.dataset.set_epoch(epoch) def get_data_files(train_path): data_files = [] data_paths = train_path.split(",") for data_path in data_paths: if data_path.startswith("hdfs://"): if not isdir(data_path): raise FileNotFoundError(f"Dataset {data_path} not exists.") for filename in listdir(data_path): from ..utils.helper import get_cache_dir data_files.append(hf_hub_download(data_path, os.path.split(filename)[-1], cache_dir=get_cache_dir())) elif os.path.isdir(data_path): data_files.extend([os.path.join(data_path, fn) for fn in sorted(os.listdir(data_path))]) elif os.path.isfile(data_path): data_files.append(data_path) else: raise FileNotFoundError(f"Dataset {data_path} not exists.") file_extenstion = os.path.splitext(data_files[0])[-1][1:] if file_extenstion not in ["parquet", "jsonl", "json", "csv", "arrow"]: raise ValueError(f"{file_extenstion} files are not supported.") file_extenstion = "json" if file_extenstion == "jsonl" else file_extenstion return data_files, file_extenstion @DATASET_REGISTRY.register("mapping") def build_mapping_dataset( train_path: str, transform: Optional[Callable] = None, namespace: Literal["train", "test"] = "train", source_name: Optional[str] = None, **kwargs, ) -> "Dataset": """ Build mapping dataset. Args: train_path (str): data path transform (Optional[Callable]): transform function namespace (Literal["train", "test"]): dataset namespace source_name (Optional[str]): source name Returns: Dataset: mapping dataset """ logger.info_rank0("Start building mapping dataset") data_files, file_extenstion = get_data_files(train_path) with main_process_first(): dataset = load_dataset(file_extenstion, data_files=data_files, split=namespace) if transform: transform = partial(transform, source_name=source_name) return MappingDataset(data=dataset, transform=transform) @DATASET_REGISTRY.register("iterable") def build_iterable_dataset( train_path: str, transform: Optional[Callable] = None, namespace: Literal["train", "test"] = "train", seed: int = 42, source_name: Optional[str] = None, split_by_node: bool = True, shuffle: bool = True, **kwargs, ) -> "IterableDataset": """ Build iterative dataset. Args: train_path (str): data path transform (Optional[Callable]): transform function namespace (Literal["train", "test"]): dataset namespace seed (int): random seed source_name (Optional[str]): source name Returns: IterableDataset: iterative dataset """ logger.info_rank0("Start building iterative dataset") data_files, file_extenstion = get_data_files(train_path) dataset = load_dataset(file_extenstion, data_files=data_files, split=namespace, streaming=True) if shuffle: dataset = dataset.shuffle(seed=seed, buffer_size=10_000) if split_by_node: parallel_state = get_parallel_state() dataset = split_dataset_by_node(dataset, parallel_state.dp_rank, parallel_state.dp_size) if transform: transform = partial(transform, source_name=source_name) return IterativeDataset(dataset, transform=transform) @DATASET_REGISTRY.register("interleave") def build_interleave_dataset( train_path: str, datasets_type: str = "mapping", namespace: Literal["train", "test"] = "train", transform: Optional[Callable] = None, seed: int = 42, **kwargs, ): """ Build interleave dataset. Args: train_path (str): data path datasets_type (str): datasets type namespace (Literal["train", "test"]): dataset namespace transform (Optional[Callable]): transform function seed (int): random seed Returns: InterleavedIterableDataset: interleaved iterable dataset or InterleavedMappingDataset: interleaved mapping dataset """ logger.info_rank0("Start building interleave dataset") multisource_config = parse_multisource_config(train_path) logger.info_rank0(f"multisource_config: {multisource_config}") sources = multisource_config["sources"] schedule = multisource_config["schedule"] source_names = multisource_config["names"] if len(schedule) > 1 or schedule[0]["schedule_type"] != "const": logger.info_rank0("Interleaved dataset only supports const schedule type.") weights = schedule[0]["weights"] datasets = [] if datasets_type == "iterable": logger.info_rank0("Start building iterable multisource dataset") def add_ds_idx_to_iterable(dataset, ds_idx, source_name): def trans_example(example): return {**example, "ds_idx": ds_idx, "source_name": source_name} return dataset.map(trans_example) for idx, source in enumerate(sources): dataset = build_iterable_dataset(source, namespace=namespace, seed=seed, split_by_node=False) ds = dataset._data ds = add_ds_idx_to_iterable(ds, idx, source_names[idx]) datasets.append(ds) interleave_dataset = interleave_datasets(datasets=datasets, probabilities=weights, seed=seed) # split dataset by node parallel_state = get_parallel_state() interleave_dataset = split_dataset_by_node(interleave_dataset, parallel_state.dp_rank, parallel_state.dp_size) interleave_dataset = InterleavedIterableDataset( interleave_dataset, transform=transform, ) elif datasets_type == "mapping": logger.info_rank0("Start building mapping multisource dataset") for idx, source in enumerate(sources): dataset = build_mapping_dataset(source, namespace=namespace) ds = dataset._data ds = ds.add_column("ds_idx", [idx] * len(ds)) ds = ds.add_column("source_name", [source_names[idx]] * len(ds)) datasets.append(ds) interleave_dataset = InterleavedMappingDataset( interleave_datasets(datasets=datasets, probabilities=weights, seed=seed), transform=transform, ) else: raise ValueError(f"Unsupported datasets_type: {datasets_type}") return interleave_dataset @DATASET_REGISTRY.register("energon") def build_energon_dataset( train_path: str, transform: Optional[Callable] = None, namespace: Literal["train", "test"] = "train", max_samples_per_sequence: Optional[int] = None, virtual_epoch_length: Optional[int] = 0, shuffle_buffer_size: Optional[int] = None, num_workers: Optional[int] = None, **kwargs, ) -> "Dataset": """ Build Megatron-Energon native dataset using the official get_train_dataset function. This is the recommended way to use Megatron-Energon datasets as it provides: - Automatic length calculation based on virtual_epoch_length - Built-in field mapping (txt -> text) - Professional streaming dataset support - Built-in error handling and performance optimizations Args: train_path (str): Path to the energon dataset directory transform (Optional[Callable]): Transform function to apply to samples namespace (Literal["train", "test"]): Dataset namespace (not used for energon) max_samples_per_sequence (Optional[int]): Maximum samples per sequence virtual_epoch_length (Optional[int]): Virtual epoch length for length calculation shuffle_buffer_size (Optional[int]): Shuffle buffer size num_workers (Optional[int]): Number of workers (if None, will be auto-detected) Returns: Dataset: Megatron-Energon native dataset """ try: from megatron.energon import WorkerConfig, get_train_dataset except ImportError as e: raise ImportError( "Megatron-Energon is not installed. Please install it with: pip install megatron-energon" ) from e logger.info_rank0(f"Start building Megatron-Energon native dataset from {train_path}") # Get parallel state for distributed training parallel_state = get_parallel_state() # Auto-detect number of workers if not provided if num_workers is None: # Try to get from environment or use a reasonable default num_workers = int(os.environ.get("TORCH_DATA_WORKERS", "1")) # Create base WorkerConfig base_worker_config = WorkerConfig( rank=parallel_state.dp_rank, world_size=parallel_state.dp_size, num_workers=num_workers ) # Wrap it with our compatible version worker_config = base_worker_config logger.info(f"Created WorkerConfig: rank={parallel_state.dp_rank}, world_size={parallel_state.dp_size}") if virtual_epoch_length is None: # Estimate based on data path - look for .nv-meta/info.json try: meta_path = os.path.join(train_path, ".nv-meta", "info.json") if os.path.exists(meta_path): import json with open(meta_path) as f: info = json.load(f) if "splits" in info and "train" in info["splits"]: virtual_epoch_length = info["splits"]["train"].get("num_samples", 1000000) else: virtual_epoch_length = 0 except Exception as e: logger.warning(f"Could not determine virtual_epoch_length from metadata: {e}") if virtual_epoch_length is None: virtual_epoch_length = 0 # Fallback logger.info(f" - max_samples_per_sequence: {max_samples_per_sequence}") logger.info(f" - virtual_epoch_length: {virtual_epoch_length}") logger.info(f" - shuffle_buffer_size: {shuffle_buffer_size}") # Get the dataset using Megatron-Energon's official function dataset = get_train_dataset( path=train_path, split_part=namespace, worker_config=worker_config, batch_size=None, # No batching at dataset level shuffle_buffer_size=shuffle_buffer_size, max_samples_per_sequence=max_samples_per_sequence, virtual_epoch_length=virtual_epoch_length, repeat=True, # Always repeat for training ) logger.info(f"Dataset type: {type(dataset)} Dataset length: {len(dataset)}") # Wrap in our EnergonDataset for Megatron-Energon specific functionality return EnergonDataset(dataset, transform) @DATASET_REGISTRY.register("veomni_weighted_multisource") def build_weighted_multisource_dataset( train_path: str, transform: Optional[Callable] = None, seed: int = 42, shuffle: bool = True, **kwargs: Any, ) -> IterableDataset: multisource_config = parse_multisource_config(train_path) schedule = multisource_config["schedule"] if len(schedule) != 1 or schedule[0].get("schedule_type") != "const": raise ValueError("simple_multisource only supports a single const schedule now") weights = schedule[0]["weights"] sources = multisource_config["sources"] source_names = multisource_config.get("names") source_ids = multisource_config.get( "source_names", source_names ) # if source_ids is not provided, use source_names as source_ids level = multisource_config.get("level", "sample") stopping_strategy = multisource_config.get("stopping_strategy", "first_exhausted") split_by_node = multisource_config.get("upstream_sharded", True) datasets = [ build_iterable_dataset( train_path=source, seed=seed, transform=transform, split_by_node=split_by_node, shuffle=shuffle, ) for source in sources ] return WeightedMultiSourceDataset( datasets=datasets, weights=weights, seed=seed, level=level, source_names=source_names, source_ids=source_ids, upstream_sharded=split_by_node, stopping_strategy=stopping_strategy, )