Diff Coverage

            raise TypeError(f"Argument 'shared_stage' must be list or tuple, \
                             but got type {type(shared_stage)}.")
        self._shared_stage = shared_stage
        self._parameter = parameter
        self._owner_module = owner_module
        self._param_name = param_name
        self.group = None

    @property
    def parameter(self):

    @property
    def parameter(self):
        """Return the shared parameter object."""
        if self._owner_module is not None and self._param_name is not None:
            return getattr(self._owner_module, self._param_name)
        return self._parameter

    @property
    def shared_stage(self):

            param = shared_param_info.parameter
            # On the meta-init -> FSDP-wrap -> load path the storage is still on
            # meta here; the trainer re-runs this after materialization, so skip
            # until it is real (broadcasting a meta tensor is invalid).
            if getattr(param, "is_meta", False):
                continue
            shared_stage = shared_param_info.shared_stage
            group, group_ranks = self._init_shared_parameter_group(shared_stage)
            shared_param_info.group = group
            # DTensor dispatch whitelists DistCommBroadcast and unwraps DTensor

        return platform.get_global_rank(self.pp_group, real_stage_index)

    def _init_shared_parameter_group(self, shared_stage):
        """init group of shared parameter."""
        group_ranks = [self._global_rank(stage) for stage in shared_stage]
        # When the shared stages span the whole PP sub-mesh (the tied-embedding
        # pp=2 case), reuse the PP group: it is already created per dp line, so
        # this avoids a world-collective ``new_group`` that would deadlock under
        # FSDP where each dp line needs its own {stage0, stage1} group.
        if self.mesh is not None and len(group_ranks) == self.mesh.size():
            return self.mesh.get_group(), group_ranks
        group = platform.create_group(group_ranks)
        return group, group_ranks

    def sync_shared_parameters_grad(self) -> None:

            group = shared_param_info.group
            # ``group`` is assigned by ``_sync_shared_parameters`` once the
            # parameter is materialized off meta; if that handshake has not run
            # there is no peer to reduce with yet, so skip.
            if group is None:
                continue

            grad = param.grad
            has_local_grad = grad is not None
            if has_local_grad:
                # Per-stage FSDP stores the reduced boundary-parameter grad as
                # a sharded DTensor. Reduce its local shard so matching DP
                # shards at the tied PP endpoints sum to one gradient.
                local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
            else:
                # Every member of a shared-parameter group must enter the same
                # collective even when this stage did not use the parameter in
                # backward. Derive the zero contribution from the live local

                # Every member of a shared-parameter group must enter the same
                # collective even when this stage did not use the parameter in
                # backward. Derive the zero contribution from the live local
                # parameter shard so shape, dtype, and device match TP/FSDP.
                local_param = param.to_local() if isinstance(param, DTensor) else param
                local_grad = platform.full_like(local_param, 0)

            # platform.all_reduce expects group_info (with .group for Torch, or str for MindSpore)
            group_info = group if isinstance(group, str) else SimpleNamespace(group=group)
            reduced_grad, handle = platform.all_reduce(local_grad, group_info, async_op=True)
            if handle is not None:
                handle.wait()

            if has_local_grad:
                # Torch may return a contiguous communication tensor for a
                # non-contiguous input. Preserve the original grad object (and
                # DTensor layout) while copying the completed reduction back.
                if reduced_grad is not local_grad:
                    platform.load_into_param(local_grad, reduced_grad)
            elif param.requires_grad:
                if isinstance(param, DTensor):
                    param.grad = DTensor.from_local(reduced_grad, param.device_mesh, param.placements)
                else:
                    param.grad = reduced_grad

    def _check_src_stage(self, src_stage):
        """check type for src_stage."""
        if src_stage is None:

        return [("isend", cur_out, info.global_rank) for cur_out, info in zip(out, rg_infos)]

    def execute_reduce_grad(self) -> None:
        """Reduce nested FSDP gradients and finalize every HSDP root in the stage."""
        fsdp_schedulers = []
        seen_scheduler_ids = set()
        for _, submod in platform.get_cells_and_names(self.submodule):
            if not isinstance(submod, HSDPModule):
                continue
            scheduler = submod.hsdp_scheduler
            scheduler_id = id(scheduler)
            if scheduler_id in seen_scheduler_ids:
                continue
            seen_scheduler_ids.add(scheduler_id)
            fsdp_schedulers.append(scheduler)

        if not fsdp_schedulers:
            return

        for scheduler in fsdp_schedulers:
            scheduler.scheduler_ctx.is_last_backward = True
            scheduler.set_reshard_after_backward(True)
            scheduler.set_requires_grad_sync(True)

        for scheduler in fsdp_schedulers:
            scheduler.hsdp_state.post_backward()
            scheduler.hsdp_state.reduce_params()

        # A plain pipeline-stage root may contain one or more independently
        # wrapped HSDP child roots. Finalize every runtime root so each context
        # clears its backward state; the platform hook also drains the global

        # wrapped HSDP child roots. Finalize every runtime root so each context
        # clears its backward state; the platform hook also drains the global
        # fused reduction queues. A pre-forward fallback retains the historical
        # single-root behavior for callers that invoke this method directly.
        root_schedulers = [
            scheduler for scheduler in fsdp_schedulers
            if scheduler._is_root  # pylint: disable=protected-access
        ]
        if not root_schedulers:
            root_schedulers = fsdp_schedulers[:1]
        for scheduler in root_schedulers:
            # No public API exposes root backward finalization. force_reduce=True
            # ensures a differentiable PP input cannot defer the final drain.
            scheduler._root_backward_hook(force_reduce=True)  # pylint: disable=protected-access

    def _build_padded_sens(self, micro_index):
        """Build an N-length sens list aligned with the forward output structure.

    def _validate_inplace_partial_inputs(op_name: str, args: tuple, kwargs: dict) -> None:
        """Reject local in-place add/sub when Partial contributions need gating."""
        if op_name not in {"InplaceAddExt", "InplaceSubExt"} or not args:
            return
        first = args[0]
        if len(args) >= 2:
            second = args[1]
        elif "other" in kwargs:
            second = kwargs["other"]
        else:
            return
        if not isinstance(first, DTensor):
            return
        mesh_ndim = len(first.layout.partial)
        first_partial = tuple(first.layout.partial)
        if isinstance(second, DTensor):
            second_partial = tuple(second.layout.partial)
            if len(second_partial) != mesh_ndim:
                raise ValueError(
                    f"For {op_name}, in-place input mesh dimensions must match, "
                    f"but got {mesh_ndim} and {len(second_partial)}."
                )
        else:
            second_partial = (None,) * mesh_ndim
        if first_partial != second_partial:
            raise ValueError(
                f"For {op_name}, input Partial placements must be identical for "
                f"local in-place execution, but got {first_partial} and {second_partial}."
            )

    @staticmethod
    def _restore_inplace_dtensor_result(op_name: str, args: tuple, output: Any) -> Any:
        """Return the original DTensor wrapper after a local in-place operation."""
        if op_name in {"add_", "sub_"} and args and isinstance(args[0], DTensor):
            return args[0]
        return output

    @staticmethod
    def _dispatch_new(func, distribute_op, packed_call, result) -> Tensor:

def _partial_signature(layout, mesh_ndim: int) -> tuple:
    """Return one Partial entry per output mesh dimension."""
    if layout is None:
        return (None,) * mesh_ndim
    partial = tuple(layout.partial)
    if len(partial) != mesh_ndim:
        raise ValueError(
            f"Input and output mesh dimensions must match, but got "
            f"input={len(partial)} and output={mesh_ndim}."
        )
    return partial

            output_partial_type is not None
            and input_partial[mesh_dim] is None
            and output_layout.mesh.get_local_rank(mesh_dim) != 0
        ):
            return False
    return True


def _zero_contribution(value):

def _zero_contribution(value):
    """Create a strict zero while retaining a floating tensor's grad edge."""
    is_complex = getattr(value, "is_complex", None)
    if callable(is_complex) and is_complex():
        return value * 0
    if hasattr(value, "clamp"):
        # ``value * 0`` turns +/-Inf into NaN. Clamp first so the forward value
        # is finite, then multiply by zero to make the derivative zero even at
        # the clamp boundary.
        return value.clamp(0, 0) * 0
    if isinstance(value, (bool, int, float, complex)):
        return type(value)(0)
    return value * 0


def _unwrap_local_value(value):
    """Convert DTensor-like values to local tensors while preserving containers."""

    @staticmethod
    def _canonicalize_binary_operands(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
        """Move semantic ``input`` and ``other`` operands before option values."""
        local_kwargs = dict(kwargs)
        if args:
            if "input" in local_kwargs:
                raise ValueError("Binary elementwise input was provided both positionally and by keyword.")
            lhs = args[0]
            trailing_args = args[1:]
        elif "input" in local_kwargs:
            lhs = local_kwargs.pop("input")
            trailing_args = ()
        else:
            return args, local_kwargs

        if trailing_args:
            if "other" in local_kwargs:
                raise ValueError("Binary elementwise other was provided both positionally and by keyword.")
            rhs = trailing_args[0]
            trailing_args = trailing_args[1:]
        elif "other" in local_kwargs:
            rhs = local_kwargs.pop("other")
        else:
            return args, local_kwargs
        return (lhs, rhs, *trailing_args), local_kwargs

    def preprocess(self, args: tuple, kwargs: dict) -> tuple:
        """Canonicalize binary operands before local unwrapping and layout caching.

        Returns:
            Local arguments, local keyword arguments, and layout cache values
            with the semantic left and right operands in the first two slots.
        """
        canonical_args, canonical_kwargs = self._canonicalize_binary_operands(args, kwargs)
        return super().preprocess(canonical_args, canonical_kwargs)

    def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:  # pylint: disable=W0221
        """
        Infer output layout for Add operator.

            ValueError: If propagated Partial status is not "sum" or None.
        """
        layouts = tuple(cache_values[:-1])
        if self.op_name in _INPLACE_ELEMENTWISE_OPS and len(layouts) >= 2:
            mesh_ndim = max(
                (len(layout.partial) for layout in layouts[:2] if layout is not None),
                default=0,
            )
            partial_signatures = tuple(
                _partial_signature(layout, mesh_ndim)
                for layout in layouts[:2]
            )
            if partial_signatures[0] != partial_signatures[1]:
                raise ValueError(
                    f"For {self.op_name}, input Partial placements should be identical "
                    f"for in-place execution, but got {partial_signatures}."
                )

        # ``*args``/``**kwargs`` forward add-specific extras such as ``alpha``.
        # A strict zero retains a valid zero-gradient autograd edge.
        def _expand_impl(x1, x2, *args, **kwargs):
            local_x1 = x1 if x1_contributes else _zero_contribution(x1)
            local_x2 = x2 if x2_contributes else _zero_contribution(x2)
            return func(local_x1, local_x2, *args, **kwargs)

        return _expand_impl

def _src_data_rank_for_tensor(tensor: Any, src_data_rank: Optional[int]) -> Optional[int]:
    """Disable source-rank communication while sharding meta tensors."""
    if getattr(tensor, "is_meta", False):
        return None
    return src_data_rank


__all__ = [

                    "RowwiseParallel expects a DTensor from Linear outputs; "
                    f"got {type(outputs)}. If this is an unsupported module, extend I/O hooks."
                )
        if reduce_dtype is not None and tuple(outputs.placements) != tuple(output_layouts):
            local_output = platform.cast_fp_tensor(reduce_dtype, outputs.to_local())
            outputs = DTensor.from_local(local_output, outputs.device_mesh, outputs.placements)
        if tuple(outputs.placements) != tuple(output_layouts):
            outputs = outputs.redistribute(device_mesh, output_layouts)
        if use_local_output:
            return outputs.to_local()

                    self.reduce_dtype is not None
                    and has_partial_output
                    and out_layout != desired_out_layout
                ):
                    local_out = platform.cast_fp_tensor(self.reduce_dtype, dt_out.to_local())
                    dt_out = DTensor.from_local(local_out, dt_out.device_mesh, dt_out.placements)
                if out_layout != desired_out_layout:
                    dt_out = dt_out.redistribute(device_mesh, (desired_out_layout,))
                prepared_outputs.append(
                    dt_out.to_local() if self.use_local_output else dt_out

        base_seed: Seed offset; each sample uses ``base_seed + idx``.
    """

    def __init__(self, num_samples: int, seq_length: int, vocab: int, base_seed: int = 42) -> None:
        self.num_samples = int(num_samples)
        self.seq_length = int(seq_length)
        self.vocab = int(vocab)
        self.base_seed = int(base_seed)

    def __len__(self) -> int:
        return self.num_samples

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        g = torch.Generator().manual_seed(self.base_seed + idx)
        input_ids = torch.randint(0, self.vocab, (self.seq_length,), generator=g)
        return {"input_ids": input_ids, "labels": input_ids.clone()}


@DATASET_REGISTRY.register("dummy")
def build_dummy(*, base: Any, args: Any, **_: Any) -> DummyDataset:

    Picks ``vocab`` from ``model.config.vocab_size`` when available, else
    defaults to 32000 (matches the pre-refactor behaviour in BaseTrainer).
    """
    seq_len = args.data.max_seq_len
    vocab_size = base.model.config.vocab_size
    base_seed = args.train.seed
    total_samples = base.state.max_steps * args.train.global_batch_size
    ds = DummyDataset(total_samples, seq_len, vocab_size, base_seed=base_seed)
    logger.info("Dummy dataset created: %d samples, seq_len=%d", total_samples, seq_len)
    return ds

    actually pulls a sample.
    """

    def __init__(self, hf_ds: Any) -> None:
        self.data = hf_ds

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, idx: int):
        item = self.data[idx]
        return {
            "input_ids": torch.tensor(item["input_ids"], dtype=torch.long),
            "labels": torch.tensor(item["labels"], dtype=torch.long),
        }

def _load_raw(args: Any, data_type: str) -> Any:
    """Run the appropriate ``load_dataset`` call for ``data_type``."""
    from datasets import load_dataset  # pylint: disable=C0415  # optional dep

    train_path = args.data.train_path
    if not train_path:
        raise ValueError(f"data.train_path is required when data.type='{data_type}'")

    if data_type == "json_file":
        return load_dataset("json", data_files=train_path, split="train")
    subset = args.data.subset
    if subset:
        return load_dataset(train_path, subset, split="train")
    return load_dataset(train_path, split="train")


def _maybe_truncate(ds: Any, args: Any) -> Any:
    """Apply ``data.train_size`` if it shrinks the dataset."""
    train_size = args.data.train_size
    if train_size and train_size < len(ds):
        ds = ds.select(range(train_size))
        logger.info("Dataset truncated to %d samples", train_size)
    return ds


def _build_hf(*, base: Any, args: Any, data_transform: Any, data_type: str, **_: Any) -> TokenizedDataset:
    """Shared loader for ``hf_datasets`` and ``json_file``.

    sequences and wraps in :class:`TokenizedDataset`. Updates
    ``base.state.max_steps`` to ``min(cfg.max_steps, len/global_bs)`` so
    epoch boundaries stay consistent with the data on hand.
    """
    logger.info(
        "Loading dataset: type=%s, path=%s", data_type, args.data.train_path,
    )
    ds = _maybe_truncate(_load_raw(args, data_type), args)

    if data_transform is not None:
        ds = ds.map(
            data_transform,
            batched=True,
            remove_columns=ds.column_names,
            desc="Tokenizing",

        )
    # Drop empty rows only when the dataset is already tokenized — without a
    # transform a raw text dataset has no ``input_ids`` column and the filter
    # predicate would raise ``KeyError`` instead of loading.
    if "input_ids" in ds.column_names:
        ds = ds.filter(lambda x: len(x["input_ids"]) > 0)

    wrapped = TokenizedDataset(ds)
    # max_steps clipping must happen pre-dataloader so the train loop's
    # epoch count matches the data on hand; the trainer reads
    # ``base.state.max_steps`` further down the build chain. The cap is
    # the TOTAL step budget across all epochs — clipping to a single

    # epoch count matches the data on hand; the trainer reads
    # ``base.state.max_steps`` further down the build chain. The cap is
    # the TOTAL step budget across all epochs — clipping to a single
    # epoch would silently truncate multi-epoch training.
    num_epochs = max(int(args.train.num_train_epochs or 1), 1)
    base.state.max_steps = min(
        args.train.max_steps,
        num_epochs * (len(wrapped) // max(args.train.global_batch_size, 1)),
    )
    logger.info(
        "Dataset ready: %d samples, max_steps=%d",
        len(wrapped), base.state.max_steps,
    )
    return wrapped


@DATASET_REGISTRY.register("hf_datasets")
def build_hf_datasets(**kwargs: Any) -> TokenizedDataset:

@DATASET_REGISTRY.register("hf_datasets")
def build_hf_datasets(**kwargs: Any) -> TokenizedDataset:
    """Build an HF hub / arrow dataset."""
    return _build_hf(data_type="hf_datasets", **kwargs)


@DATASET_REGISTRY.register("json_file")
def build_json_file(**kwargs: Any) -> TokenizedDataset:

@DATASET_REGISTRY.register("json_file")
def build_json_file(**kwargs: Any) -> TokenizedDataset:
    """Build a local Alpaca-style ``.json`` / ``.jsonl`` dataset."""
    return _build_hf(data_type="json_file", **kwargs)

    ``seed`` is accepted for API parity but unused; the algorithm is
    deterministic in ``(weights, num_samples)``.
    """
    if len(dataset_sizes) != len(weights):
        raise ValueError(
            f"dataset_sizes ({len(dataset_sizes)}) and weights ({len(weights)}) "
            f"must have the same length"
        )
    if any(w < 0 for w in weights):
        raise ValueError(f"weights must be non-negative, got {list(weights)}")
    total_w = float(sum(weights))
    if total_w <= 0:
        raise ValueError("weights must contain at least one non-zero value")
    norm_w = np.asarray([w / total_w for w in weights], dtype=np.float64)
    sizes = np.asarray([int(s) for s in dataset_sizes], dtype=np.int64)
    # A zero-length source with a non-zero weight would be selected by the
    # greatest-error rule and then hit ``counters[d] % sizes[d]`` → divide by
    # zero. Reject up-front (a source contributing samples must have samples).
    empty_weighted = [i for i, (s, w) in enumerate(zip(sizes, weights)) if s == 0 and w > 0]
    if empty_weighted:
        raise ValueError(
            f"BlendableDataset sources {empty_weighted} have a non-zero weight "
            f"but zero length; an empty source cannot contribute samples"
        )

    dataset_index = np.zeros(num_samples, dtype=np.int32)
    dataset_sample_index = np.zeros(num_samples, dtype=np.int64)
    realised = np.zeros(len(weights), dtype=np.float64)
    counters = np.zeros(len(weights), dtype=np.int64)
    for g in range(num_samples):
        # error[d] = how far behind source d's realised fraction is —
        # picking the argmax keeps every source within 1 sample of its
        # target share at every prefix (Megatron's invariant).
        error = norm_w * (g + 1) - realised
        d = int(np.argmax(error))
        dataset_index[g] = d
        # Modulo wrap so a small sub-dataset reused inside a long blend
        # cycles through its samples in order instead of indexing OOB.
        dataset_sample_index[g] = counters[d] % sizes[d]
        realised[d] += 1.0
        counters[d] += 1
    return dataset_index, dataset_sample_index


class BlendableDataset(Dataset):
    """Weighted blend over a list of datasets.

        weights: Sequence[float],
        num_samples: int,
        seed: int = 1234,
    ) -> None:
        if not datasets:
            raise ValueError("BlendableDataset requires at least one dataset")
        self.datasets = datasets
        sizes = [len(d) for d in datasets]
        self.num_samples = int(num_samples)
        self.dataset_index, self.dataset_sample_index = _build_blend_indices(
            sizes, weights, self.num_samples, seed,
        )
        logger.info(
            "BlendableDataset built: %d sub-datasets, weights=%s, total samples=%d",
            len(datasets), list(weights), self.num_samples,
        )

            len(datasets), list(weights), self.num_samples,
        )

    def __len__(self) -> int:
        return self.num_samples

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        d = int(self.dataset_index[idx])
        s = int(self.dataset_sample_index[idx])
        return self.datasets[d][s]

    - ``str`` of whitespace-separated tokens: ``"0.3 /data/a 0.7 /data/b"``
    - ``list[float | str]`` in the same alternating order
    - ``list[list]`` of ``[weight, prefix]`` pairs
    """
    if isinstance(raw, str):
        toks = raw.split()
    elif isinstance(raw, (list, tuple)):
        # Already in pair form?
        if raw and isinstance(raw[0], (list, tuple)):
            return [(float(w), str(p)) for w, p in raw]
        toks = list(raw)
    else:
        raise ValueError(f"Unsupported blend spec type: {type(raw)}")

    if len(toks) % 2 != 0:
        raise ValueError(
            f"Blend spec must have an even number of tokens "
            f"(alternating weight + prefix); got {len(toks)}: {toks}"
        )
    out: List[Tuple[float, str]] = []
    for i in range(0, len(toks), 2):
        out.append((float(toks[i]), str(toks[i + 1])))
    return out


def _looks_like_blend(train_path: Any) -> bool:
    """Return ``True`` when ``train_path`` is a multi-source blend spec.

    a list of pairs); a single corpus is a bare path-prefix. Keying off "the
    first token parses as a float" (rather than "contains whitespace") lets a
    single path that legitimately contains spaces still be read as one source.
    """
    if isinstance(train_path, (list, tuple)):
        return True
    if not isinstance(train_path, str):
        return False
    toks = train_path.split()
    if len(toks) < 2:
        return False
    try:
        float(toks[0])
    except ValueError:
        return False
    return True


@DATASET_REGISTRY.register("megatron")
def build_megatron(*, base: Any, args: Any, **_: Any) -> Any:

@DATASET_REGISTRY.register("megatron")
def build_megatron(*, base: Any, args: Any, **_: Any) -> Any:
    """Build a Megatron ``.bin``/``.idx`` dataset (single source or blend)."""
    del base
    data_cfg = args.data
    train_cfg = args.train
    train_path = data_cfg.train_path
    if not train_path:
        raise ValueError("data.train_path is required when data.type='megatron'")

    seq_length = int(data_cfg.max_seq_len)
    # ``megatron_seed`` may be present-but-None on a default DataConfig; fall
    # through to ``train.seed`` in that case so the typed config doesn't have
    # to invent a sentinel.
    megatron_seed = data_cfg.megatron_seed
    if megatron_seed is None:
        megatron_seed = train_cfg.seed
    seed = int(megatron_seed)
    pad_token_id = int(data_cfg.pad_token_id)
    eod_token_id = data_cfg.eod_token_id
    eod_mask_loss = bool(data_cfg.eod_mask_loss)
    mmap_bin = bool(data_cfg.mmap_bin_files)

    num_samples = int(train_cfg.max_steps * train_cfg.global_batch_size)
    if num_samples <= 0:
        raise ValueError(
            f"num_samples = max_steps * global_batch_size must be > 0; "
            f"got max_steps={train_cfg.max_steps}, global_bs={train_cfg.global_batch_size}"
        )

            f"got max_steps={train_cfg.max_steps}, global_bs={train_cfg.global_batch_size}"
        )

    # Single source: simplest path, no blend needed.
    if not _looks_like_blend(train_path):
        prefix = strip_suffix(train_path)
        idx_ds = IndexedDataset(prefix, mmap=mmap_bin)
        gpt_ds = GPTDataset(
            idx_ds, num_samples=num_samples, seq_length=seq_length,
            seed=seed, pad_token_id=pad_token_id,
            eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id,
        )
        logger.info(
            "Megatron dataset built: prefix=%s docs=%d sequences=%d samples=%d seq_len=%d",
            prefix, idx_ds.num_documents, len(idx_ds), len(gpt_ds), seq_length,
        )
        return gpt_ds

    # Blend: parse, build per-source GPTDataset, wrap in BlendableDataset.
    pairs = _parse_blend(train_path)
    weights = [w for w, _ in pairs]
    sub_datasets = []
    for w, p in pairs:
        prefix = strip_suffix(p)
        idx_ds = IndexedDataset(prefix, mmap=mmap_bin)
        # Each sub-dataset is sized to cover its share of the blend; we ask
        # for the full ``num_samples`` here because blending picks indices
        # modulo the sub-dataset length.
        sub_datasets.append(
            GPTDataset(
                idx_ds, num_samples=num_samples, seq_length=seq_length,
                seed=seed, pad_token_id=pad_token_id,
                eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id,

                seed=seed, pad_token_id=pad_token_id,
                eod_mask_loss=eod_mask_loss, eod_token_id=eod_token_id,
            )
        )
        logger.info(
            "  blend source w=%g prefix=%s docs=%d sequences=%d",
            w, prefix, idx_ds.num_documents, len(idx_ds),
        )
    return BlendableDataset(sub_datasets, weights, num_samples=num_samples, seed=seed)

    appears once — that randomness is part of the canonical training
    distribution and the realised epoch-boundary noise is what Megatron's
    schedulers expect.
    """
    doc_idx = np.tile(np.arange(num_documents, dtype=np.int32), num_epochs)
    rng.shuffle(doc_idx)
    return doc_idx


def _build_sample_index(
    sizes: np.ndarray,

    token so labels can be shifted by 1; the helper keeps the
    pre-Megatron-Cython slow-path semantics (no boundary correction —
    multi-document samples are reassembled by :meth:`GPTDataset._read_sample`).
    """
    tokens_to_consume = num_epochs * tokens_per_epoch - 1  # -1 to keep ``+1`` extra token in reach
    sample_stride = seq_length
    num_samples = tokens_to_consume // sample_stride
    # int64: column 0 indexes ``doc_idx`` whose length is
    # ``num_documents * num_epochs`` — int32 would wrap past ~2.1e9 slots on
    # large corpora trained for many epochs, silently selecting the wrong
    # document.
    sample_idx = np.zeros((num_samples + 1, 2), dtype=np.int64)

    cur_doc_index = 0
    cur_doc_offset = 0
    cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
    sample_idx[0, 0] = cur_doc_index
    sample_idx[0, 1] = cur_doc_offset
    for sample_id in range(1, num_samples + 1):
        remaining = sample_stride
        while remaining > 0:
            if cur_doc_remaining > remaining:
                cur_doc_offset += remaining
                cur_doc_remaining -= remaining
                remaining = 0
            else:
                remaining -= cur_doc_remaining
                cur_doc_index += 1
                if cur_doc_index >= doc_idx.shape[0]:
                    # Document stream exhausted — record where we landed
                    # and stop. Caller will pad the tail.
                    cur_doc_remaining = 0
                    break
                cur_doc_offset = 0
                cur_doc_remaining = int(sizes[doc_idx[cur_doc_index]])
        sample_idx[sample_id, 0] = cur_doc_index
        sample_idx[sample_id, 1] = cur_doc_offset
    return sample_idx


def _build_shuffle_index(num_samples: int, rng: np.random.RandomState) -> np.ndarray:
    """Random permutation of ``[0, num_samples)``."""
    shuffle = np.arange(num_samples, dtype=np.int64)
    rng.shuffle(shuffle)
    return shuffle


def _layout_fingerprint(doc_sizes: np.ndarray) -> str:
    """Hash the per-document token-count array.

    corpus regenerated at the same path with a different document layout — even
    one preserving document count and total token count — produces a different
    fingerprint and invalidates the stale cache.
    """
    return hashlib.sha1(np.ascontiguousarray(doc_sizes, dtype=np.int64).tobytes()).hexdigest()[:16]


def _cache_key(
    prefix: str,

    :func:`_layout_fingerprint`) so that regenerating the ``.bin``/``.idx``
    under the same path invalidates a stale cache instead of silently
    replaying indices computed against the old token stream.
    """
    h = hashlib.sha1(
        f"{prefix}|{seq_length}|{num_samples}|{num_epochs}|{seed}|{layout_fp}".encode("utf-8"),
    ).hexdigest()[:16]
    return h


class GPTDataset(Dataset):
    """Megatron-style GPT dataset on top of an :class:`IndexedDataset`.

        eod_mask_loss: bool = False,
        eod_token_id: Optional[int] = None,
        cache_dir: Optional[str] = None,
    ) -> None:
        if num_samples <= 0:
            raise ValueError(f"num_samples must be > 0, got {num_samples}")
        if seq_length <= 0:
            raise ValueError(f"seq_length must be > 0, got {seq_length}")

        self.indexed_dataset = indexed_dataset
        self.num_samples = int(num_samples)
        self.seq_length = int(seq_length)
        self.seed = int(seed)
        self.pad_token_id = int(pad_token_id)
        self.eod_mask_loss = bool(eod_mask_loss)
        self.eod_token_id = eod_token_id

        self.document_index, self.sample_index, self.shuffle_index = self._build_indices(cache_dir)
        # Adjacent samples often touch the same document — cache the concat
        # result for the last few requested doc ids so straddling reads
        # don't re-walk the IndexedDataset.
        self._cached_doc = functools.lru_cache(maxsize=128)(self._read_doc)

    def _build_indices(self, cache_dir: Optional[str]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Build (or load from cache) the three Megatron-style indices."""
        sizes = self.indexed_dataset.sequence_lengths
        num_docs = self.indexed_dataset.num_documents
        if num_docs == 0:
            raise ValueError("IndexedDataset has no documents to sample from")
        # Build a per-document length array using the existing
        # ``document_indices`` boundaries.
        doc_boundaries = self.indexed_dataset.document_indices
        doc_sizes = np.zeros(num_docs, dtype=np.int64)
        for d in range(num_docs):
            start = int(doc_boundaries[d])
            end = int(doc_boundaries[d + 1])
            doc_sizes[d] = int(sizes[start:end].sum())

        tokens_per_epoch = int(doc_sizes.sum())
        if tokens_per_epoch == 0:
            raise ValueError("IndexedDataset has zero total tokens")
        # ``+1`` because each sample reads ``seq_length + 1`` tokens (input
        # plus the next-token label). Ceil so we always cover ``num_samples``.
        tokens_required = self.num_samples * self.seq_length + 1
        num_epochs = max(1, int(np.ceil(tokens_required / tokens_per_epoch)))

        cache_dir = self._resolve_cache_dir(cache_dir)
        key = _cache_key(
            self.indexed_dataset.path_prefix, self.seq_length,
            self.num_samples, num_epochs, self.seed,
            _layout_fingerprint(doc_sizes),
        )
        doc_path = os.path.join(cache_dir, f"doc_idx_{key}.npy")
        sample_path = os.path.join(cache_dir, f"sample_idx_{key}.npy")
        shuffle_path = os.path.join(cache_dir, f"shuffle_idx_{key}.npy")

        if os.path.isfile(doc_path) and os.path.isfile(sample_path) and os.path.isfile(shuffle_path):
            return np.load(doc_path), np.load(sample_path), np.load(shuffle_path)

        rng = np.random.RandomState(self.seed)
        doc_idx = _build_document_index(num_docs, num_epochs, rng)
        sample_idx = _build_sample_index(
            doc_sizes, doc_idx, self.seq_length, tokens_per_epoch, num_epochs,
        )
        # ``sample_idx`` has at most ``num_samples + 1`` rows; clip in case
        # we built slightly more (rounded up by tokens_per_epoch / seq_len).
        usable_samples = max(sample_idx.shape[0] - 1, 1)
        shuffle_idx = _build_shuffle_index(min(self.num_samples, usable_samples), rng)
        np.save(doc_path, doc_idx)
        np.save(sample_path, sample_idx)
        np.save(shuffle_path, shuffle_idx)
        return doc_idx, sample_idx, shuffle_idx

    def _resolve_cache_dir(self, user_cache_dir: Optional[str]) -> str:
        """Return a writable cache directory, falling back to ``$TMPDIR``.

        Defaulting to ``<prefix>_cache`` next to the corpus is convenient,
        but production datasets often sit on read-only shared mounts where
        ``os.makedirs`` would raise ``EROFS``. Detect and silently relocate.
        """
        candidate = user_cache_dir or (self.indexed_dataset.path_prefix + "_cache")
        try:
            os.makedirs(candidate, exist_ok=True)
            # Probe writability — ``os.makedirs`` succeeds on directories
            # that exist read-only.
            probe = os.path.join(candidate, ".hp_cache_probe")
            with open(probe, "w", encoding="utf-8") as f:
                f.write("")
            os.remove(probe)
            return candidate
        except OSError as exc:
            fallback = os.path.join(tempfile.gettempdir(), "hp_megatron_cache")
            os.makedirs(fallback, exist_ok=True)
            logger.warning(
                "GPTDataset cache dir %s not writable (%s); using fallback %s",
                candidate, exc, fallback,
            )
            return fallback

    def __len__(self) -> int:
        return min(self.num_samples, int(self.shuffle_index.size))

    def _read_sample(self, sample_pos: int) -> np.ndarray:
        """Concatenate the ``seq_length + 1`` tokens for sample ``sample_pos``."""
        start_doc_idx, start_offset = self.sample_index[sample_pos]
        end_doc_idx, end_offset = self.sample_index[sample_pos + 1]
        parts = []
        if start_doc_idx == end_doc_idx:
            doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
            parts.append(doc_tokens[start_offset:end_offset + 1])
        else:
            doc_tokens = self._cached_doc(int(self.document_index[start_doc_idx]))
            parts.append(doc_tokens[start_offset:])
            for doc_id in range(start_doc_idx + 1, end_doc_idx):
                parts.append(self._cached_doc(int(self.document_index[doc_id])))
            tail_doc_tokens = self._cached_doc(int(self.document_index[end_doc_idx]))
            parts.append(tail_doc_tokens[:end_offset + 1])
        out = np.concatenate(parts) if len(parts) > 1 else parts[0]
        needed = self.seq_length + 1
        if out.size >= needed:
            return out[:needed]
        # Pad short tail (only happens for the very last sample).
        pad = np.full(needed - out.size, self.pad_token_id, dtype=out.dtype)
        return np.concatenate([out, pad])

    def _read_doc(self, doc_id: int) -> np.ndarray:
        """All tokens of ``doc_id`` as a single 1-D array (memoised)."""
        return self.indexed_dataset.get_document(doc_id)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        actual = int(self.shuffle_index[idx])
        tokens = self._read_sample(actual)
        input_ids = torch.from_numpy(tokens[: self.seq_length].astype(np.int64))
        labels = input_ids.clone()
        if self.eod_mask_loss and self.eod_token_id is not None:
            labels[input_ids == int(self.eod_token_id)] = -100
        return {"input_ids": input_ids, "labels": labels}

    The single source of truth for path-prefix normalisation — :func:`get_bin_path`,
    :func:`get_idx_path` and the dataset builder all route through here so the
    suffix rule never diverges across call sites.
    """
    return prefix[:-4] if prefix.endswith((".bin", ".idx")) else prefix


def get_bin_path(prefix: str) -> str:
    """Return ``<prefix>.bin`` (strips an already-present ``.bin``/``.idx``)."""
    return strip_suffix(prefix) + ".bin"


def get_idx_path(prefix: str) -> str:
    """Return ``<prefix>.idx``."""
    return strip_suffix(prefix) + ".idx"


def dtype_code(dtype: np.dtype) -> int:
    """Return the on-disk dtype code for ``dtype``.

    Raises:
        ValueError: When ``dtype`` is not part of the Megatron dtype table.
    """
    dtype = np.dtype(dtype)
    if dtype not in _DTYPE_CODES:
        raise ValueError(
            f"Unsupported Megatron indexed-dataset dtype: {dtype}. "
            f"Allowed: {sorted(d.name for d in _DTYPE_CODES)}"
        )
    return _DTYPE_CODES[dtype]


def code_dtype(code: int) -> np.dtype:
    """Inverse of :func:`dtype_code`."""
    try:
        return _DTYPES[code]
    except KeyError as exc:
        raise ValueError(f"Unknown Megatron dtype code: {code}") from exc


class _IndexReader:
    """In-memory metadata parsed from the ``.idx`` companion file.

    :class:`IndexedDataset`.
    """

    def __init__(self, idx_path: str, *, multimodal: bool = False) -> None:
        if not os.path.isfile(idx_path):
            raise FileNotFoundError(f"Megatron index file not found: {idx_path}")
        with open(idx_path, "rb") as f:
            magic = f.read(len(_MAGIC))
            if magic != _MAGIC:
                raise ValueError(
                    f"Not a Megatron indexed dataset (magic mismatch): "
                    f"{idx_path}. Expected {_MAGIC!r}, got {magic!r}"
                )
            version = struct.unpack("<Q", f.read(8))[0]
            if version != _VERSION:
                raise ValueError(
                    f"Unsupported indexed-dataset version: {version} "
                    f"(only v{_VERSION} is recognised)"
                )
            code = struct.unpack("<B", f.read(1))[0]
            self.dtype = code_dtype(code)
            self.sequence_count = struct.unpack("<Q", f.read(8))[0]
            self.document_count = struct.unpack("<Q", f.read(8))[0]

            # The index payload is small (4-12 bytes per sequence) — read it
            # into RAM rather than memmap. Avoids the fragility of slicing a
            # memmap with offsetted ``np.frombuffer`` while still being fast.
            self.sequence_lengths = np.frombuffer(
                f.read(self.sequence_count * 4), dtype=np.int32,
            )
            self.sequence_pointers = np.frombuffer(
                f.read(self.sequence_count * 8), dtype=np.int64,
            )
            self.document_indices = np.frombuffer(
                f.read(self.document_count * 8), dtype=np.int64,
            )
            if multimodal:
                self.sequence_modes = np.frombuffer(
                    f.read(self.sequence_count), dtype=np.int8,
                )
            else:
                self.sequence_modes = None


class IndexedDataset:
    """Read-only Megatron ``.bin``/``.idx`` dataset.