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):

                continue
            grad = param.grad
            # A tied end whose backward has not populated a grad has nothing to
            # contribute to the cross-stage sum; skip it.
            if grad is None:
                continue
            # Under FSDP the grad is a sharded DTensor; all-reduce the local
            # shard so the tied ends' matching dp shards sum to one gradient.
            if isinstance(grad, DTensor):
                grad = grad.to_local()
            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
            # 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)
            platform.all_reduce(grad, group_info)

    def get_last_stage_sens(self, last_stage_outputs):
        """Get last stage sens"""
        p_sens = None
        loss_scale = float(getattr(self, "loss_scale", 1.0))
        if isinstance(last_stage_outputs, (list, tuple)):
            p_sens = []
            for _, out_i in enumerate(last_stage_outputs):
                if isinstance(out_i, DTensor):

            p_sens = []
            for _, out_i in enumerate(last_stage_outputs):
                if isinstance(out_i, DTensor):
                    repeat_num = out_i.layout.repeat_num()
                    sens_i = platform.full_like(out_i.to_local(), loss_scale / repeat_num)
                else:
                    sens_i = platform.full_like(out_i, loss_scale)
                p_sens.append(sens_i)
        else:
            if isinstance(last_stage_outputs, DTensor):
                repeat_num = last_stage_outputs.layout.repeat_num()
                p_sens = platform.full_like(last_stage_outputs.to_local(), loss_scale / repeat_num)
            else:
                p_sens = platform.full_like(last_stage_outputs, loss_scale)

        return p_sens

    def _construct_forward_recv_info(self, micro_index, idx, global_rank, meta):

def _straight_through_scale(value, scale):
    """Scale the forward value while preserving the input gradient."""
    delta = value / scale - value
    return platform.stop_gradient(delta) + value


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

        # that ``AllReduce(SUM)`` of the Partial output reconstructs
        # ``x1 + true_x2`` while keeping the Replicate gradient unscaled.

        def _scaled_forward_only(x):
            return _straight_through_scale(x, scaling_factor)

        def _expand_impl1(x1, x2, *args, **kwargs):
            return func(_scaled_forward_only(x1), x2, *args, **kwargs)

        def _expand_impl2(x1, x2, *args, **kwargs):
            return func(x1, _scaled_forward_only(x2), *args, **kwargs)

        return _expand_impl1 if not x1_partial else _expand_impl2

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()

                    dt_out = out
                else:
                    dt_out = DTensor.from_local(out, device_mesh, (out_layout,))
                if self.reduce_dtype is not None 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
    model_cfg = getattr(getattr(base, "model", None), "config", None)
    vocab_size = getattr(model_cfg, "vocab_size", 32000) if model_cfg else 32000
    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. Megatron-LM uses int64 here for the same reason.
    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:
    """Memory-mapped view over the ``.idx`` companion file.

    is essentially free.
    """

    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.

        byte-compatible.
    """

    def __init__(self, path_prefix: str, *, mmap: bool = True, multimodal: bool = False) -> None:
        idx_path = get_idx_path(path_prefix)
        bin_path = get_bin_path(path_prefix)
        if not os.path.isfile(bin_path):
            raise FileNotFoundError(f"Megatron .bin not found: {bin_path}")
        self._index = _IndexReader(idx_path, multimodal=multimodal)
        if mmap:
            self._bin_buffer = np.memmap(bin_path, dtype=np.uint8, mode="r", order="C")
        else:
            with open(bin_path, "rb") as f:
                data = f.read()
            self._bin_buffer = np.frombuffer(data, dtype=np.uint8)
        self.path_prefix = path_prefix

    # ------------------------------------------------------------------
    # Sequence-level access
    # ------------------------------------------------------------------

    # Sequence-level access
    # ------------------------------------------------------------------

    def __len__(self) -> int:
        return int(self._index.sequence_count)

    @property
    def dtype(self) -> np.dtype:
        """Token dtype (numpy)."""
        return self._index.dtype

    @property
    def sequence_lengths(self) -> np.ndarray:
        """int32[N] token-counts for each sequence."""
        return self._index.sequence_lengths

    @property
    def document_indices(self) -> np.ndarray:
        """int64[D+1] sequence-index boundaries marking documents."""
        return self._index.document_indices

    def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> np.ndarray:
        """Read a slice of sequence ``idx``.

            A ``np.ndarray`` of the configured dtype. The returned array is
            a view into the mmap when ``mmap=True``; callers must not write
            to it.
        """
        seq_len = int(self._index.sequence_lengths[idx])
        if length is None:
            length = seq_len - offset
        if offset < 0 or length < 0 or offset + length > seq_len:
            raise ValueError(
                f"Slice [{offset}:{offset + length}] out of bounds for "
                f"sequence {idx} of length {seq_len}"
            )
        item_size = self._index.dtype.itemsize
        byte_ptr = int(self._index.sequence_pointers[idx]) + offset * item_size
        # ``sequence_pointers`` are accumulated ``nbytes`` of items, so
        # ``byte_ptr`` is always a multiple of ``item_size`` and a typed
        # view is alignment-safe; the slice is a zero-copy view into the
        # mmap (or a plain ndarray slice in the RAM path).
        return self._bin_buffer[byte_ptr: byte_ptr + length * item_size].view(self._index.dtype)

    def __getitem__(self, idx: int) -> np.ndarray:
        return self.get(idx)

    # ------------------------------------------------------------------
    # Document-level access
    # ------------------------------------------------------------------

        On disk ``document_count`` is ``D + 1`` because Megatron stores the
        full boundary array including the trailing ``N`` terminator.
        """
        return int(self._index.document_count) - 1

    def get_document(self, doc_idx: int) -> np.ndarray:
        """Concatenate all sequences belonging to ``doc_idx``."""
        if doc_idx < 0 or doc_idx >= self.num_documents:
            raise IndexError(
                f"document index {doc_idx} out of range [0, {self.num_documents})"
            )
        start = int(self._index.document_indices[doc_idx])
        end = int(self._index.document_indices[doc_idx + 1])
        parts = [self.get(i) for i in range(start, end)]
        if not parts:
            return np.array([], dtype=self._index.dtype)
        return np.concatenate(parts)


class IndexedDatasetBuilder:
    """Write a ``.bin``/``.idx`` pair sequentially.

        Python lists for the index metadata until :meth:`finalize`.
    """

    def __init__(self, path_prefix: str, *, dtype: np.dtype = np.int32) -> None:
        self.path_prefix = path_prefix
        self.dtype = np.dtype(dtype)
        # Make sure the dtype is supported before opening any file.
        dtype_code(self.dtype)
        self._bin_path = get_bin_path(path_prefix)
        self._idx_path = get_idx_path(path_prefix)
        os.makedirs(os.path.dirname(self._bin_path) or ".", exist_ok=True)
        self._bin_file = open(self._bin_path, "wb")  # pylint: disable=R1732
        self._sequence_lengths: List[int] = []
        self._sequence_pointers: List[int] = []
        self._document_indices: List[int] = [0]
        self._byte_offset = 0

    def add_item(self, tokens: np.ndarray) -> None:
        """Append one sequence of tokens to the dataset."""
        tokens = np.asarray(tokens, dtype=self.dtype)
        self._bin_file.write(tokens.tobytes(order="C"))
        self._sequence_lengths.append(int(tokens.size))
        self._sequence_pointers.append(self._byte_offset)
        self._byte_offset += int(tokens.nbytes)

    def end_document(self) -> None:
        """Mark the current sequence boundary as a document boundary."""
        self._document_indices.append(len(self._sequence_lengths))

    def finalize(self) -> Tuple[str, str]:
        """Close ``.bin`` and write the companion ``.idx``.