Diff Coverage

) -> DistributedDatasetConfig:
    """Derive distributed-packing sizes and retain optional runtime tuning."""
    try:
        seq_len = int(data_config.get("seq_length", max_seq_len))
    except (KeyError, TypeError, ValueError) as exc:
        raise ValueError("Distributed loading requires data_config.seq_length or transform.max_seq_len") from exc
    if max_seq_len is not None and seq_len != max_seq_len:
        raise ValueError("data_config.seq_length must match transform.max_seq_len for native batch balancing")
    options = dict(dataloader_target.load_balance)
    trainer_managed_options = {
        "seq_len",
        "local_batch_size",

def _trainer_model_batch_spec(target: Any, collate_fn: Any, data_config: Mapping[str, Any]) -> Any:
    """Resolve declared model fields without reading samples on model peers."""
    configured = getattr(target, "model_batch_spec", None)
    if configured is not None:
        return configured.build() if callable(getattr(configured, "build", None)) else configured
    declared = getattr(collate_fn, "model_batch_spec", None)
    if declared is not None:
        return declared
    if isinstance(collate_fn, TextPackingCollator):
        return {
            "input_ids": TensorFieldSpec((1, None), torch.int64),
            "labels": TensorFieldSpec((1, None), torch.int64),
            "cu_seq_lens": TensorFieldSpec((None,), torch.int32),
        }

    # Their layouts must be configured explicitly rather than guessed from a sample.
    if collate_fn is default_collate and not data_config.get("is_instruction_dataset", False):
        spec = {name: TensorFieldSpec((None, None), torch.int64) for name in ("input_ids", "labels")}
        if data_config.get("create_ltor_fields_in_dataloader", False):
            spec.update({
                "loss_mask": TensorFieldSpec((None, None), torch.float32),
                "position_ids": TensorFieldSpec((None, None), torch.int64),
                "attention_mask": TensorFieldSpec((None, 1, None, None), torch.bool),
            })

        balancing_algorithm: BalancingAlgorithm | None = None,
) -> Any:
    """Reuse native Dataset outputs and the existing collator in collective loading."""
    if batch_sampler is None:
        raise ValueError("Native BatchSampler balancing requires a mapping Dataset")
    if getattr(mesh_context, "device_mesh", None) is None:
        raise ValueError("Native BatchSampler balancing requires mesh_context.device_mesh")
    if int(getattr(mesh_context, "pp_size", 1)) != 1:
        raise ValueError("Native BatchSampler balancing does not yet support pipeline parallelism")
    # Target exposes its resolved callable only as _target_; inspecting it
    # prevents replacing a second-stage dynamic selector without warning.
    target_callable = getattr(dataloader_target, "_target_", None)
    if isinstance(target_callable, type) and issubclass(target_callable, _DynamicBatchLoader):

        dataset: Any, dataloader_target: Any, training_config: Any, mesh_context: Any, seed: int,
) -> Any:
    """Build rank-local sampling for mapping datasets; iterable sources own their order."""
    if _is_iterable_dataset(dataset):
        return None
    return build_dataset_batch_sampler(
        total_samples=len(dataset),
        micro_batch_size=training_config.micro_batch_size,
        global_batch_size=training_config.global_batch_size,

    if len(datasets) != 3:
        raise ValueError("datasets must contain train, validation, and test entries")

    if all(dataset is None for dataset in datasets):
        return (None, None, None), (None, None, None)

    if dataloader_target is None:
        raise ValueError("dataloader_target must define a build target")

        if source_type in {"online", "packed_text"}:
            return OnlineBoundaryResolver()
        if source_type == "indexed":
            return IndexedBoundaryResolver(getattr(tokenizer, "eod", None))
        raise ValueError(f"Unsupported batch source type: {source_type!r}")

    def _validate_attention_inputs(self, attention_mode: str, create_attention_mask: bool) -> None:
        """Require shifted labels and document isolation for prepared text records."""
        if self.source_type == "packed_text":

    def _validate_attention_inputs(self, attention_mode: str, create_attention_mask: bool) -> None:
        """Require shifted labels and document isolation for prepared text records."""
        if self.source_type == "packed_text":
            if not self.labels_are_shifted:
                raise ValueError("Packed text batches require shifted labels; labels_are_shifted must be True")
            if not create_attention_mask or (attention_mode == "compressed" and self.runtime_input_adapter is None):
                raise ValueError(
                    "Packed text batches require document-boundary attention: enable "
                    "create_attention_mask_in_dataloader and supply a runtime_input_adapter for compressed attention"
                )
        if self.is_sft_dataset:

        if num_epochs == 1:
            separate_final_epoch = False
            samples_without_final_epoch = 0
        else:
            samples_without_final_epoch = (
                (num_epochs - 1) * num_tokens_per_epoch - extra_token
            ) // self.config.sequence_length
            samples_from_final_epoch = self.num_samples - samples_without_final_epoch
            samples_per_epoch = (num_tokens_per_epoch - extra_token) // self.config.sequence_length

            # Keep a small final epoch separate from the main shuffled range.
            threshold = 0.80
            separate_final_epoch = samples_from_final_epoch < int(threshold * samples_per_epoch)
        return num_epochs, separate_final_epoch, samples_without_final_epoch

    def _build_document_sample_shuffle_indices(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Build or load the document, sample, and shuffle indices.

    pixels = sample.get("pixel_values")
    if input_ids is None or input_ids.ndim != 1 or input_ids.shape[0] == 0:
        raise ValueError("VLM metadata requires one non-empty 1-D input_ids tensor before collation")
    if grid is None or grid.ndim != 2 or grid.shape[1] != 3:
        raise ValueError("VLM metadata requires image_grid_thw with shape [num_images, 3]")
    if grid.device.type != "cpu":
        raise ValueError("VLM metadata must be extracted on CPU before device transfer")
    image_patches = 0
    for dimensions in grid.tolist():
        if any(not isinstance(dimension, int) or isinstance(dimension, bool) for dimension in dimensions):
            raise ValueError("VLM image_grid_thw must contain integer patch dimensions")
        if any(dimension <= 0 for dimension in dimensions):
            raise ValueError("VLM image_grid_thw dimensions must be positive")
        image_patches += dimensions[0] * dimensions[1] * dimensions[2]
    if pixels is None or pixels.ndim != 2 or pixels.shape[0] != image_patches:
        raise ValueError("VLM pixel_values rows must equal the raw patch count from image_grid_thw")
    seq_len = int(input_ids.shape[0])

        if source_rank is None or self.parallel_context.tp_world_size == 1:
            parallel_batch = self._move_mapping_to_device(cp_local_batch)
            return parallel_batch

        is_source = self.parallel_context.tp_rank == source_rank
        if is_source:
            parallel_batch = self._move_mapping_to_device(cp_local_batch)
            metadata = []
            for name, value in parallel_batch.items():

                parallel_batch[name] = value
                continue
            if not is_source:
                parallel_batch[name] = torch.empty(shape, dtype=dtype, device=self.device)
            dist.broadcast(parallel_batch[name], group=self.parallel_context.tp_group, group_src=source_rank)
        return parallel_batch

    def _move_mapping_to_device(self, batch: Mapping[str, Any]) -> dict[str, Any]:
        """Transfer Omni tensors without changing their native dtypes."""

        return dataloader_context

    tp_size = int(getattr(mesh_context, "tp_size", 1))
    if tp_source_rank is not None and not 0 <= tp_source_rank < tp_size:
        raise ValueError("tp_source_rank must be a rank within the TP group or None")
    build_on_rank = _always_build if tp_source_rank is None else partial(
        _is_dataloader_rank, mesh_context, tp_source_rank,
    )
    build_cache_on_rank = partial(_is_global_cache_builder_rank, shared_storage)

    def _validate_policy_and_flags(self) -> None:
        """Validate policy names, boolean flags and the relative-gain threshold."""
        if self.oversized_policy not in ("error", "single"):
            raise ValueError("oversized_policy must be 'error' or 'single'.")
        for name in (
                "dataset_already_sharded",
                "drop_last",
                "metadata_mode",

            or not 0.0 <= self.min_balance_gain < 1.0
        ):
            raise ValueError("min_balance_gain must be in [0, 1).")
        if self.communication_backend not in ("gloo", "hccl"):
            raise ValueError("communication_backend must be 'gloo' or 'hccl'.")
        if not self.drop_last:
            raise ValueError(
                "Dynamic distributed packing currently requires drop_last=True so every DP rank receives the "
                "same number of non-empty packing bins."
            )

                not isinstance(self.planner_rank, int)
                or isinstance(self.planner_rank, bool)
                or self.planner_rank < 0
        ):
            raise ValueError("planner_rank must be a non-negative integer or None.")

    @staticmethod
    def _validate_rank_tuple(value: tuple[int, ...] | None, name: str) -> None:
        if value is None:

    @staticmethod
    def _validate_rank_tuple(value: tuple[int, ...] | None, name: str) -> None:
        if value is None:
            return
        if not isinstance(value, tuple) or not value:
            raise ValueError(f"{name} must be a non-empty tuple or None.")
        if any(not isinstance(rank, int) or isinstance(rank, bool) or rank < 0 for rank in value):
            raise ValueError(f"{name} must contain non-negative integer ranks.")
        if len(value) != len(set(value)):
            raise ValueError(f"{name} must not contain duplicate ranks.")

    @staticmethod
    def _validate_name_tuple(value: tuple[str, ...] | None, name: str) -> None:
        if value is None:

    def _validate_name_tuple(value: tuple[str, ...] | None, name: str) -> None:
        if value is None:
            return
        if not isinstance(value, tuple) or not value or any(not isinstance(item, str) or not item for item in value):
            raise ValueError(f"{name} must be a non-empty tuple of strings or None.")
        if len(value) != len(set(value)):
            raise ValueError(f"{name} must not contain duplicate names.")


def _resolve_service_ranks(
        topology: DataTopology,

    """Resolve Reader and planner ownership within the root mesh."""
    dataset_reader_ranks = config.dataset_reader_ranks or topology.constructor_ranks
    unknown_ranks = set(dataset_reader_ranks) - set(topology.rank_list)
    if unknown_ranks:
        raise ValueError(f"dataset_reader_ranks are outside the root mesh: {sorted(unknown_ranks)}.")
    planner_rank = min(topology.constructor_ranks) if config.planner_rank is None else config.planner_rank
    data_plane_ranks = set(dataset_reader_ranks) | set(topology.constructor_ranks)
    if planner_rank not in data_plane_ranks:
        raise ValueError(f"planner_rank {planner_rank} must be a Dataset Reader or Data Constructor rank.")
    return dataset_reader_ranks, planner_rank


def _normalize_dataloader_kwargs(

    """Resolve validated DataLoader execution options."""
    if dataloader_kwargs is None:
        supplied_options = {}
    elif not isinstance(dataloader_kwargs, Mapping):
        raise ValueError("dataloader_kwargs must be a mapping or None.")
    else:
        try:
            supplied_options = dict(dataloader_kwargs)
        except Exception as exc:
            raise ValueError(f"dataloader_kwargs cannot be copied: {exc}") from exc
    _validate_dataloader_option_keys(supplied_options)
    options = dict(_DEFAULT_DATALOADER_KWARGS)
    options.update(supplied_options)
    _validate_worker_options(options)

def _validate_dataloader_option_keys(supplied_options: Mapping[str, Any]) -> None:
    """Reject unsupported keys and overrides of internally managed DataLoader options."""
    non_string_keys = [key for key in supplied_options if not isinstance(key, str)]
    if non_string_keys:
        raise ValueError(f"dataloader_kwargs keys must be strings, but got {non_string_keys!r}.")
    managed_keys = sorted(set(supplied_options) & _MANAGED_DATALOADER_KWARGS)
    if managed_keys:
        raise ValueError(
            f"dataloader_kwargs cannot override distributed sampling, batching, collation, or deterministic "

    num_workers = options["num_workers"]
    timeout = options.get("timeout", 0)
    try:
        normalized_timeout = float(timeout)
    except (TypeError, ValueError, OverflowError):
        normalized_timeout = math.nan
    if (
            not isinstance(timeout, (int, float))
            or isinstance(timeout, bool)
            or not math.isfinite(normalized_timeout)

            or isinstance(timeout, bool)
            or not math.isfinite(normalized_timeout)
            or normalized_timeout < 0
    ):
        raise ValueError(f"dataloader_kwargs timeout must be finite and non-negative, but got {timeout!r}.")
    if num_workers == 0 and timeout != 0:
        raise ValueError("dataloader_kwargs timeout must be zero when num_workers=0.")

        return
    if options["num_workers"] == 0:
        raise ValueError("dataloader_kwargs multiprocessing_context requires num_workers > 0.")
    if not isinstance(multiprocessing_context, (str, BaseContext)):
        raise ValueError("dataloader_kwargs multiprocessing_context is invalid.")
    start_method = multiprocessing_context
    if isinstance(multiprocessing_context, BaseContext):
        try:
            start_method = multiprocessing_context.get_start_method()
        except Exception as exc:
            raise ValueError(
                f"dataloader_kwargs multiprocessing_context cannot report its start method: {exc}"
            ) from exc
    if start_method not in torch.multiprocessing.get_all_start_methods():
        raise ValueError(f"dataloader_kwargs multiprocessing_context={start_method!r} is not supported.")


def _normalize_communication_device(communication_device: Any) -> torch.device | None:
    """Normalize a rank-local A2A device."""

def _normalize_communication_device(communication_device: Any) -> torch.device | None:
    """Normalize a rank-local A2A device."""
    if communication_device is None:
        return None
    try:
        return torch.device(communication_device)
    except Exception as exc:
        raise ValueError(f"communication_device is invalid: {communication_device!r}.") from exc

            split_fn: Callable[[Any], Sequence[Any]] | None,
    ) -> None:
        """Use native DP sampler owners as Readers without applying a second stride."""
        if self.dataset_reader_ranks != self.topology.constructor_ranks:
            raise ValueError("Native batch_sampler currently requires Dataset Readers on the Data Constructor ranks.")
        if config.dataset_already_sharded:
            raise ValueError(
                "Native batch_sampler owns DP slicing; disable dataset_already_sharded."
            )

            batch_sampler, data_rank=self.topology.data_rank,
            dp_size=self.topology.data_parallel_size, local_batch_size=config.local_batch_size,
        )
        if not self.is_reader:
            return
        if dataset is None:
            raise ValueError("Native batch_sampler requires a mapping Dataset on each Data Constructor.")
        sample_loader = PlannedSampleLoader(dataset, seed=config.seed, **loader_options)
        sample_loader.set_epoch(batch_sampler.epoch)
        if config.metadata_mode:
            if not callable(getattr(metadata, "__getitem__", None)) or len(metadata) != len(dataset):

        collate_fn: Callable[[Sequence[Any]], Any] | None,
) -> None:
    """Validate callback types and mutually exclusive metadata providers."""
    if metadata_fn is not None and not callable(metadata_fn):
        raise ValueError("metadata_fn must be callable or None.")
    if metadata_fn is not None and metadata is not None:
        raise ValueError("Provide either metadata or metadata_fn, but not both.")
    if pack_fn is not None and not callable(pack_fn):
        raise ValueError("pack_fn must be callable or None.")
    if collate_fn is not None and not callable(collate_fn):
        raise ValueError("collate_fn must be callable or None.")


def _resolve_repacking_callbacks(
        config: DistributedDatasetConfig,

        Collective iterator yielding one constructed local batch per micro-batch.
    """
    if external_step_source is not None:
        if model_batch_spec is not None:
            raise ValueError("model_batch_spec is only supported with a native batch_sampler.")
        if split_fn is not None:
            raise ValueError("split_fn is only supported with a native batch_sampler.")
        if config.metadata_mode:
            raise ValueError("external_step_source requires metadata_mode=False for online metadata.")
        if batch_sampler is not None or metadata is not None:
            raise ValueError("external_step_source cannot be combined with a sampler or metadata sequence.")
        if dataloader_kwargs:
            raise ValueError("Configure source worker options on external_step_source, not the balancing wrapper.")
        return build_local_balancing_dataloader(
            external_step_source, mesh, config, metadata_fn=metadata_fn,
            pack_fn=default_pack_fn if pack_fn is None else pack_fn,
            collate_fn=default_collate_fn if collate_fn is None else collate_fn,

        Args:
            objective: Maximum backbone load, or relative backbone-load variance.
        """
        if objective not in ("makespan", "balance"):
            raise ValueError("objective must be 'makespan' or 'balance'.")
        self.objective_name = objective
        self.padding_cost: Callable[[int], WorkloadCost] | None = None

    def assign(

    ) -> Sequence[Sequence[SampleKey]]:
        """Refine one-bin-per-slot assignments without token overflow or stage budgets."""
        lookup = {item.key: item for item in samples}
        if any(sum(lookup[key].pack_tokens for key in keys) > capacity for keys in candidate):
            return candidate
        if self.padding_cost is not None:
            candidate = min(
                (candidate, reference_bins),
                key=lambda bins: self._repair_objective(

        if not self._uses_scheduled_makespan() and (
                candidate is not reference_bins
                or any(sum(lookup[key].pack_tokens for key in keys) != capacity for keys in candidate)
        ):
            return candidate
        return self._repair_bin_pairs(lookup, candidate, capacity, group_size)

    def _uses_scheduled_makespan(self) -> bool:
        return self.objective_name == "makespan" and type(self).objective is LPTBalancingAlgorithm.objective

                    return self._pair_potential(costs, pair, group_size, total)

                pairs = sorted(pairs, key=lambda pair: (-_potential(pair), pair))
            else:
                pairs = sorted(pairs)
            for left, right in pairs:
                version = (left, right, versions[left], versions[right])
                if len(checked) == self.max_repair_pairs:
                    break

        order = sorted(range(len(costs)), key=lambda index: (-costs[index].llm, index))
        if self._uses_scheduled_makespan():
            leaders = order[::group_size]
        else:
            leaders = [
                max(range(start, start + group_size), key=lambda index: costs[index].llm)
                for start in range(0, len(costs), group_size)
            ]
        positions = {index: position for position, index in enumerate(order)}

                for side in range(2):
                    trial[pair[side]] += self._padding_llm_cost(capacity - trial_tokens[side])
                if scheduled_makespan:
                    score = math.fsum(sorted(trial, reverse=True)[::group_size])
                elif self.padding_cost is not None:
                    score = self._repair_objective([WorkloadCost(llm=value) for value in trial], group_size)
                else:
                    score = costs[rank]
                candidates.append((score, capacity - trial_tokens[rank], rank))
            if not candidates:
                return reference_bins
            rank = min(candidates)[-1]
            bins[rank].append(item.key)
            tokens[rank] += item.pack_tokens
            costs[rank] += item.cost.llm

        algorithm: Optional deterministic assignment and objective implementation.
    """
    resolved = FeasibleLPTBalancingAlgorithm() if algorithm is None else algorithm
    if not callable(getattr(resolved, "assign", None)) or not callable(getattr(resolved, "objective", None)):
        raise ValueError("balancing_algorithm must provide callable assign() and objective() methods.")
    return resolved


__all__ = [

        The sole source sample, unchanged.
    """
    del seq_len
    if len(samples) != 1:
        raise ValueError("Native BatchSampler mode must keep one Dataset output per sequence bin.")
    return samples[0]


@dataclass(frozen=True)

            split_fn: Callable[[Any], Sequence[Any]] | None = None,
    ) -> None:
        """Store the validated native sampler and metadata/payload sources."""
        if metadata is None and (sample_loader is None or not callable(metadata_fn)):
            raise ValueError("Online BatchSampler Reader requires payload loading and metadata_fn.")
        if split_fn is not None and (not callable(split_fn) or metadata is not None):
            raise ValueError("split_fn requires online metadata and must be callable.")
        self._sampler = sampler
        self._reader_rank = reader_rank
        self._metadata = metadata
        self._metadata_fn = metadata_fn

            metadata = [self._metadata[key.dataset_index] for key in keys]
        if self._split_fn is None:
            original_batches = tuple(tuple((key,) for key in batch) for batch in source_batches)
        if any(not isinstance(item, SampleMetadata) for item in metadata):
            raise ValueError("BatchSampler metadata must contain SampleMetadata entries.")
        return _ReaderWindow(
            metadata=tuple(replace(item, key=key) for key, item in zip(keys, metadata)),
            samples=samples,
            sampler_state=copy.deepcopy(self._sampler.state_dict()),

    def _sample_keys(self, indices: Sequence[int], batch_position: int) -> tuple[SampleKey, ...]:
        """Bind local Dataset indices to distinct global sampling occurrences."""
        if len(indices) != self._sampler.micro_batch_size:
            raise ValueError("Native BatchSampler must yield exactly local_batch_size indices on every DP rank.")
        # A native sampler exposes local indices, not its internal permutation
        # positions. Stable DP-major occurrence slots preserve duplicates.
        start = batch_position + self._sampler.dp_rank * self._sampler.micro_batch_size
        keys = []

        start = batch_position + self._sampler.dp_rank * self._sampler.micro_batch_size
        keys = []
        for offset, index in enumerate(indices):
            if isinstance(index, bool):
                raise ValueError("Native BatchSampler indices must be non-negative integers.")
            keys.append(SampleKey(self._reader_rank, operator.index(index), start + offset))
        return tuple(keys)

    def metadata(self) -> tuple[SampleMetadata, ...]:

            ValueError: If the checkpoint epoch or payload reader is incompatible.
        """
        sampler_state = state["sampler"]
        if sampler_state.get("epoch") != state.get("epoch"):
            raise ValueError("Native BatchSampler checkpoint epoch is inconsistent.")
        if (state.get("sample_loader") is None) != (self._sample_loader is None):
            raise ValueError("Native BatchSampler checkpoint payload reader ownership changed.")
        if "window" in state or sampler_state["consumed_samples"] % self._sampler.global_batch_size:
            raise ValueError("Native BatchSampler checkpoint must align to a complete optimizer step.")
        self.close()
        self._sampler.load_state_dict(sampler_state)

def _configuration_dict(config: Any) -> dict[str, Any]:
    """Accept an ordinary dictionary or a configuration with ``to_dict``."""
    if isinstance(config, Mapping):
        return dict(config)
    if callable(getattr(config, "to_dict", None)):
        value = config.to_dict()
        if isinstance(value, Mapping):
            return dict(value)
    raise ValueError("model_config must be a mapping or a configuration exposing to_dict().")


def _integer(value: Any, name: str, minimum: int = 1) -> int:
    if not isinstance(value, int) or isinstance(value, bool) or value < minimum:

            _integer(getattr(self, name), name)
        for name in ("qk_nope_head_dim", "qk_rope_head_dim", "n_shared_experts"):
            _integer(getattr(self, name), name, 0)
        if self.q_lora_rank is not None:
            _integer(self.q_lora_rank, "q_lora_rank")
        if self.qk_nope_head_dim + self.qk_rope_head_dim == 0:
            raise ValueError("The combined query/key head dimension must be positive.")
        if (
                not isinstance(self.mlp_layer_types, tuple)
                or len(self.mlp_layer_types) != self.num_hidden_layers
                or any(kind not in ("dense", "sparse") for kind in self.mlp_layer_types)

                not isinstance(self.mlp_layer_types, tuple)
                or len(self.mlp_layer_types) != self.num_hidden_layers
                or any(kind not in ("dense", "sparse") for kind in self.mlp_layer_types)
        ):
            raise ValueError("mlp_layer_types must contain one 'dense' or 'sparse' entry per backbone layer.")
        for name in ("moe_intermediate_size", "n_routed_experts", "num_experts_per_tok"):
            _integer(getattr(self, name), name, int("sparse" in self.mlp_layer_types))
        if self.num_experts_per_tok > self.n_routed_experts:
            raise ValueError("num_experts_per_tok cannot exceed n_routed_experts.")
        if self.attn_block_mode not in ("tail", "blockmask"):
            raise ValueError("attn_block_mode must be 'tail' or 'blockmask'.")
        if self.diff_time_embed_mode not in ("add", "prepend", "adaln"):
            raise ValueError("diff_time_embed_mode must be 'add', 'prepend', or 'adaln'.")

    @classmethod
    def from_model_config(cls, model_config: Any) -> "BackboneFlopsConfig":
        """Read the final model configuration after architecture overrides.

        Args:
            model_config: Configuration with explicit decoder dimensions and layer types.
        """
        config = _configuration_dict(model_config)
        text = _configuration_dict(config.get("text_config", config))
        required = (
            "hidden_size", "num_hidden_layers", "num_attention_heads", "intermediate_size",
            "kv_lora_rank", "q_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim",
        )
        missing = [name for name in required if name not in text]
        if missing:
            raise ValueError(f"Backbone cost model is missing configuration fields: {missing}.")
        pattern = text.get("mlp_layer_types")
        if not isinstance(pattern, (list, tuple)):
            raise ValueError("model_config must provide mlp_layer_types in actual decoder layer order.")
        qk_dim = (
            _integer(text["qk_nope_head_dim"], "qk_nope_head_dim", 0)
            + _integer(text["qk_rope_head_dim"], "qk_rope_head_dim", 0)
        )
        if text.get("qk_head_dim", qk_dim) != qk_dim:
            raise ValueError("qk_head_dim must equal qk_nope_head_dim + qk_rope_head_dim.")
        return cls(
            **{name: text[name] for name in required},
            mlp_layer_types=tuple(pattern),
            **{name: text.get(name, 0) for name in (
                "moe_intermediate_size", "n_routed_experts", "num_experts_per_tok", "n_shared_experts",

                or isinstance(training_multiplier, bool)
                or not math.isfinite(training_multiplier)
                or training_multiplier <= 0
        ):
            raise ValueError("training_multiplier must be finite and positive.")
        self.training_multiplier = float(training_multiplier)
        self.config = (
            model_config if isinstance(model_config, BackboneFlopsConfig)
            else BackboneFlopsConfig.from_model_config(model_config)

        """
        prefix = _integer(metadata.features.get("P"), "metadata.features['P']", 0)
        diffusion = _integer(metadata.features.get("D"), "metadata.features['D']", 0)
        if prefix + diffusion != metadata.pack_tokens:
            raise ValueError("metadata features P + D must equal pack_tokens before optional timestep-token insertion.")
        if diffusion and self.config.diff_time_embed_mode == "prepend":
            diffusion += 1
        pairs = 0.5 * prefix * prefix + diffusion * (prefix + diffusion)
        if self.config.attn_block_mode == "blockmask":
            runs = metadata.features.get("cond_image_token_lengths")
            if not isinstance(runs, (list, tuple)):
                raise ValueError("Blockwise attention requires cond_image_token_lengths (empty for no images).")
            lengths = [_integer(length, "conditional-image run length") for length in runs]
            if sum(lengths) > prefix:
                raise ValueError("Conditional-image run lengths cannot exceed the prefix token count P.")
            pairs += 0.5 * sum(length * length for length in lengths)
        return float(self.linear_flops_per_token * (prefix + diffusion) + self.attention_flops_per_pair * pairs)

    def __call__(self, metadata: SampleMetadata) -> WorkloadCost:
        """Return backbone-only workload without changing physical packing caps."""

    """
    if cost_model is None:
        return DefaultCostModel(model_config)
    if not callable(cost_model):
        raise ValueError("cost_model must be callable: SampleMetadata -> WorkloadCost.")
    return cost_model


__all__ = ["BackboneFlopsConfig", "CostModel", "DefaultCostModel"]

    Returns:
        Immutable raw samples in their deterministic planned order.
    """
    if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len < 1:
        raise ValueError(f"seq_len must be a positive integer, but got {seq_len!r}.")
    return tuple(samples)


def default_collate_fn(packed_sequences: Sequence[Any]) -> tuple[Any, ...]:

            collate_fn: Called once per local batch with packed sequences.
            seq_len: Maximum non-oversized sequence length.
        """
        if not callable(pack_fn) or not callable(collate_fn):
            raise ValueError("pack_fn and collate_fn must be callable.")
        if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len < 1:
            raise ValueError(f"seq_len must be a positive integer, but got {seq_len!r}.")
        self._pack_fn = pack_fn
        self._collate_fn = collate_fn
        self._seq_len = seq_len

        packed_sequences = []
        for pack_index, packing_bin in enumerate(plan):
            if packing_bin.pack_tokens > self._seq_len and not packing_bin.oversized:
                raise ValueError(
                    f"Packing bin {pack_index} has {packing_bin.pack_tokens} tokens, "
                    f"exceeding seq_len={self._seq_len}."
                )
            raw_samples = [payloads[key] for key in packing_bin.sample_keys]

) -> dict[str, Any]:
    """Merge additional DataLoader options without exposing owned semantics."""
    try:
        options = dict({} if dataloader_kwargs is None else dataloader_kwargs)
    except Exception as exc:
        raise ValueError(f"dataloader_kwargs cannot be copied: {exc}") from exc
    conflicting_options = sorted(set(options) & _INTERNAL_DATALOADER_OPTIONS)
    if conflicting_options:
        raise ValueError(f"dataloader_kwargs contains internally managed options {conflicting_options}.")
    options.update({
        "num_workers": num_workers,
        "pin_memory": pin_memory,
        "prefetch_factor": prefetch_factor,

    """Return the PyTorch device module that owns streams and events."""
    module = getattr(torch, device.type, None)
    required = ("Event", "Stream", "current_stream", "stream")
    if module is None or any(not callable(getattr(module, name, None)) for name in required):
        raise ValueError(
            f"Device {device} does not expose PyTorch accelerator Stream/Event APIs."
        )
    return module

        return {key: _move_to_device(item, device) for key, item in value.items()}
    if isinstance(value, tuple):
        moved = tuple(_move_to_device(item, device) for item in value)
        if hasattr(value, "_fields"):
            return type(value)(*moved)
        return moved
    if isinstance(value, list):
        return [_move_to_device(item, device) for item in value]
    return value

        return
    visited.add(value_id)

    if isinstance(value, torch.Tensor):
        if value.device.type != "cpu":
            value.record_stream(stream)
        return
    record_stream = getattr(value, "record_stream", None)
    if callable(record_stream):
        record_stream(stream)
        return

        return

    attributes = getattr(value, "__dict__", None)
    if isinstance(attributes, dict):
        for item in attributes.values():
            _record_stream(item, stream, visited)


def _pin_memory(value: Any) -> Any:
    """Pin final collated CPU tensors without changing metadata or dtype."""
    if isinstance(value, torch.Tensor):
        if value.device.type == "cpu" and value.layout == torch.strided and not value.is_pinned():
            return value.pin_memory()
        return value
    if isinstance(value, Mapping):
        return {key: _pin_memory(item) for key, item in value.items()}
    if isinstance(value, tuple):
        pinned = tuple(_pin_memory(item) for item in value)
        return type(value)(*pinned) if hasattr(value, "_fields") else pinned
    if isinstance(value, list):
        return [_pin_memory(item) for item in value]
    pin = getattr(value, "pin_memory", None)
    return pin() if callable(pin) else value


class DevicePrefetchedStep:
    """Hand the complete device batch and its copy dependency to training."""

            device: Target CUDA or NPU device, independent of payload transport.
        """
        self.device = torch.device(device)
        if self.device.type not in ("cuda", "npu"):
            raise ValueError("DeviceStepPrefetcher requires a CUDA or NPU device.")
        self._device_module = _get_device_module(self.device)
        self._copy_stream = None
        self._pending_staging: list[tuple[Any, Any]] = []

    if device is None:
        for device_type in ("npu", "cuda"):
            device_module = getattr(torch, device_type, None)
            if device_module is not None and device_module.is_available():
                device = torch.device(device_type, device_module.current_device())
                break
        else:
            device = "cpu"
    return torch.device(device)

        metadata_mode: bool,
) -> None:
    """Check Reader ownership and require mode-compatible loading components."""
    if not isinstance(metadata_mode, bool):
        raise ValueError("metadata_mode must be boolean.")
    if metadata_mode and dataset_reader is not None:
        raise ValueError("Metadata mode must not configure an online Dataset Reader.")
    if not metadata_mode and (metadata_reader is not None or direct_sample_loader is not None):
        raise ValueError("Online mode must not configure metadata loading components.")
    is_reader = topology.global_rank in dataset_reader_ranks
    planning_reader = metadata_reader if metadata_mode else dataset_reader
    if is_reader != (planning_reader is not None):
        raise ValueError("Dataset Reader ownership does not match dataset_reader_ranks.")
    if not metadata_mode and planning_reader is not None and not callable(
            getattr(planning_reader, "prepare_next_step", None)
    ):
        raise ValueError("Online Dataset Reader must provide prepare_next_step().")
    _validate_metadata_loader_owner(
        topology,
        metadata_mode=metadata_mode,
        direct_sample_loader=direct_sample_loader,

) -> None:
    if not metadata_mode:
        return
    if topology.is_constructor != (direct_sample_loader is not None):
        raise ValueError("Every metadata Data Constructor must own one plan-aware sample loader.")


@dataclass
class _PreparedWindow:

        self._pending_local_keys: set[SampleKey] = set()
        self._window_plans: deque[DistributedPackingPlan] = deque()
        for name, value in (("accumulation_steps", accumulation_steps), ("balance_window_size", balance_window_size)):
            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                raise ValueError(f"{name} must be a positive integer.")
        self._accumulation_steps = accumulation_steps
        self._balance_window_size = balance_window_size
        if (isinstance(device_prefetch_depth, bool) or not isinstance(device_prefetch_depth, int)
                or device_prefetch_depth < 0):
            raise ValueError("device_prefetch_depth must be a non-negative integer.")
        self._device_prefetch_depth = device_prefetch_depth
        self._device = torch.device(device)
        self._current: _PreparedWindow | None = None
        self._future: Future[_PreparedWindow] | None = None

            if self._topology.is_constructor:
                self._prefetch_device_batches(1, blocking=True)
                prepared = self._device_batches.popleft()
            else:
                prepared = self._next_host_batch()
                prefetch_wait_event = self._prefetch_slots.popleft()
                if prefetch_wait_event is not None:
                    prefetch_wait_event.synchronize()
        except BaseException as exc:  # pylint: disable=broad-exception-caught
            # All preparation failures must reach peers through the broadcast error header.
            prepared = _PreparedBatch(None, error=exc)
        module = getattr(torch, self._device.type)
        consumer_stream = module.current_stream(self._device)
        try:
            batch = prepared.batch

    def _prepare_batch(self) -> Any:
        """Construct one Host round; accelerator peers receive it at device delivery."""
        if self._stopped:
            raise StopIteration
        batch = self._prepare_next_batch()
        if batch is not None and self._device_prefetch is not None:
            batch = self._device_prefetch.prepare_host(batch)
        if self._device.type == "cpu":

    def _prepare_next_batch(self) -> Any:
        """Select, balance, route, and construct samples without model-group broadcast."""
        if self._data_plane.is_member:
            return self._produce_on_data_plane()
        return None

    def _broadcast_batch(self, constructed_batch: Any) -> Any:
        """Broadcast one batch on the caller stream, or on the Host producer for CPU."""
        return self._model_transport.broadcast(constructed_batch)

            ),
        }
        try:
            copied_state = copy.deepcopy(state)
        except Exception as exc:
            raise ValueError(f"Distributed DataLoader state is not checkpointable: {exc}") from exc
        return copied_state

    def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
        """Restore a fixed-topology rank-local checkpoint.

        """
        self._validate_restore_boundary()
        try:
            state = copy.deepcopy(dict(state_dict))
        except Exception as exc:
            raise ValueError(f"Distributed DataLoader state is not copyable: {exc}") from exc
        if state.get("global_rank") != self._topology.global_rank:
            raise ValueError(
                f"Distributed DataLoader checkpoint global_rank={state.get('global_rank')!r} "
                f"does not match {self._topology.global_rank!r}."