Diff Coverage

        if current_step > lr_decay_steps:
            return min_lr_ratio

        progress = float(current_step - num_warmup_steps) / float(max(1, lr_decay_steps - num_warmup_steps))
        if not 0 <= progress <= 1:
            raise ValueError(f"cosine schedule progress must be in [0, 1], but got {progress}")
        factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
        factor = factor * (1 - min_lr_ratio) + min_lr_ratio
        return max(0, factor)

def _resolve_muon_advanced_options(options: Dict[str, Any]) -> _MuonAdvancedOptions:
    """Validate and fill low-frequency Muon keyword options."""
    unknown = options.keys() - _MUON_ADVANCED_DEFAULTS.keys()
    if unknown:
        raise TypeError(f"Muon got unexpected keyword arguments: {', '.join(sorted(unknown))}")
    resolved = dict(_MUON_ADVANCED_DEFAULTS)
    resolved.update(options)
    return _MuonAdvancedOptions(**resolved)

    ) -> Dict[torch.nn.Parameter, torch.Tensor]:
        """Compute first-order momentum and return bfloat16 NS inputs."""
        momentum = group["momentum"]
        if not isinstance(momentum, (list, tuple)):
            momentum = [momentum]
        if len(momentum) == 1:
            momentum = (momentum[0], momentum[0])
        momentum1, momentum2 = momentum
        # Pre-filter params with valid grads and ensure momentum buffers exist
        valid_params = []
        grads = []

        grads = [to_local_if_dtensor(grad) for grad in grads]
        bufs = [to_local_if_dtensor(buffer) for buffer in bufs]
        if self.momentum_update_fn is not None:
            custom_updates = list(self.momentum_update_fn(grads, bufs, momentum1, momentum2, group["nesterov"]))
            if len(custom_updates) != len(valid_params):
                raise ValueError(
                    "momentum_update_fn must return one update tensor per input gradient, "
                    f"got {len(custom_updates)} updates for {len(valid_params)} gradients"

                    "momentum_update_fn must return one update tensor per input gradient, "
                    f"got {len(custom_updates)} updates for {len(valid_params)} gradients"
                )
            custom_updates = [update.to(torch.bfloat16) for update in custom_updates]
            param_updates = dict(zip(valid_params, custom_updates))
            return param_updates

        # Match muon_update_core():
        # m_for_update = grad + momentum1 * m_old
        # m_new = grad + momentum2 * m_old

        torch._foreach_add_(bufs, grads)

        if group["nesterov"]:
            torch._foreach_mul_(local_us, momentum1)
            torch._foreach_add_(local_us, grads)

        if local_us[0].dtype != torch.bfloat16:
            local_us = [update.to(torch.bfloat16) for update in local_us]

                reshaped_inputs[0].shape,
            )

        for reshaped_input in reshaped_inputs:
            if reshaped_input.untyped_storage().data_ptr() != working_input.untyped_storage().data_ptr():
                raise ValueError("reshape_fn must return views that share storage with the working NS input tensor")

        return working_input, reshaped_inputs

    def _prepare_ns_transform(

            no_shard: bool = False,
    ) -> Dict[torch.nn.Parameter, torch.Tensor]:
        """Compute native reshape/view NS updates without transform bookkeeping."""
        if not p_list:
            return {}
        state = {"reshape_groups": defaultdict(list), "origin_shapes": {}, "working_inputs": {}, "restored": {}}
        for param in p_list:
            local_shape = getattr(param, "local_shape", None)
            if local_shape is None:

                scale = compute_muon_slice_scale(
                    update, group["matched_adamw_rms"], zero_rms_scale_mode=group["zero_rms_scale_mode"]
                )
                if group["apply_lr_in_update"]:
                    scale *= -group["lr"]
                update.mul_(scale)
                reshaped_input.copy_(update.contiguous().view_as(reshaped_input))

        for param in p_list:

        for param in p_list:
            update = ns_inputs[param].view(state["origin_shapes"][param])
            working_input = state["working_inputs"][param]
            if working_input.untyped_storage().data_ptr() != update.untyped_storage().data_ptr():
                update.copy_(working_input)
            state["restored"][param] = update
        return state["restored"]

    def _compute_batched_ns_updates_with_transform(

        """

        if not p_list:
            return {}

        state: _TransformNSState = {
            "reshape_groups": defaultdict(list),
            "origin_shapes": {},

            reshaped_update = ns_inputs[p].view(state["origin_shapes"][p])

            if (state["working_inputs"][p].untyped_storage().data_ptr()
                    != reshaped_update.untyped_storage().data_ptr()):
                reshaped_update.copy_(state["working_inputs"][p])
            state["restored"][p] = reshaped_update

        return state["restored"]

        param: torch.Tensor,
        stage_metas: List[ParamShardStageMeta],
) -> ParamShardMeta:
    """Construct shard metadata from real local-shape plans."""
    global_shape = tuple(param.shape)
    local_shape = tuple(to_local_if_dtensor(param.data).shape)
    return ParamShardMeta(
        global_shape=global_shape,
        global_numel=math.prod(global_shape),
        local_shape=local_shape,
        local_numel=math.prod(local_shape),

def _get_mesh_shape(device_mesh) -> Tuple[int, ...]:
    """Get mesh shape."""
    mesh = getattr(device_mesh, "mesh", None)
    if mesh is not None:
        return tuple(mesh.shape)
    mesh_shape = getattr(device_mesh, "mesh_shape", None)
    if mesh_shape is not None:
        return tuple(mesh_shape)
    return ()


def _get_local_shape_kinds(global_shape: Tuple[int, ...], shard_meta: ParamShardMeta) -> List[Tuple[int, ...]]:
    """Collect all distinct local shapes implied by the shard metadata."""
    dim_size_options: Dict[int, List[int]] = {}
    for stage_meta in shard_meta.stage_metas:
        tensor_dim = stage_meta.tensor_dim
        seen_sizes = []
        for split_size in stage_meta.split_sizes:
            if split_size not in seen_sizes:
                seen_sizes.append(split_size)
        dim_size_options[tensor_dim] = seen_sizes

    shape_kinds = [list(global_shape)]
    for tensor_dim, size_options in sorted(dim_size_options.items()):
        next_shape_kinds = []
        for shape_kind in shape_kinds:
            for size in size_options:
                next_shape = list(shape_kind)
                next_shape[tensor_dim] = size
                next_shape_kinds.append(next_shape)
        shape_kinds = next_shape_kinds

    unique_shape_kinds = []
    seen = set()
    for shape_kind in shape_kinds:
        shape_kind = tuple(shape_kind)
        if shape_kind in seen:
            continue
        seen.add(shape_kind)
        unique_shape_kinds.append(shape_kind)
    return unique_shape_kinds


def _debug_param_shard_metadata(hsdp_group, param_to_meta: Dict[torch.nn.Parameter, ParamShardMeta]) -> None:
    """Log shard metadata for the group on rank 0 when distributed is ready."""
    if not dist.is_available() or not dist.is_initialized():
        return

    rank = dist.get_rank()
    if rank != 0:
        return

    for record in hsdp_group.records:
        param = record.param
        shard_meta = param_to_meta.get(param)
        if shard_meta is None:
            continue
        local_shape_kinds = _get_local_shape_kinds(shard_meta.global_shape, shard_meta)
        logger.debug_rank0(
            "name=%s, placements=%s, mesh_shape=%s, global_shape=%s, local_shapes=%s",
            getattr(param, "model_name", f"param_{record.index}"),
            [str(placement) for placement in param.placements],
            _get_mesh_shape(param.device_mesh),

        shard_size: int,
        device: torch.device,
) -> List[List[Tuple[int, ...]]]:
    """All-gather per-parameter shapes within one shard process group."""
    if shard_pg is None or shard_size <= 1:
        return [[tuple(shape) for shape in local_shapes]]

    shape_tensor = torch.tensor(local_shapes, dtype=torch.int64, device=device)
    gathered = torch.empty((shard_size, *shape_tensor.shape), dtype=torch.int64, device=device)
    dist.all_gather_into_tensor(gathered.view(-1), shape_tensor.contiguous().view(-1), group=shard_pg)
    gathered_cpu = gathered.cpu().tolist()
    return [
        [tuple(shape) for shape in rank_shapes]
        for rank_shapes in gathered_cpu
    ]

        hsdp_group: HSDPCommGroup,
) -> Dict[torch.Tensor, ParamShardMeta]:
    """Build per-parameter shard metadata before HSDP batching."""
    # pylint: disable=too-many-locals
    layout_spec = hsdp_group.layout_spec
    shard_pgs = hsdp_group.shard_pgs
    if layout_spec is None or not layout_spec.shard_axes or not hsdp_group.records:
        return {}

    current_shapes = [tuple(to_local_if_dtensor(record.param.data).shape) for record in hsdp_group.records]
    stage_metas_per_param: List[List[ParamShardStageMeta]] = [[] for _ in hsdp_group.records]
    local_device = to_local_if_dtensor(hsdp_group.records[0].param.data).device

    for (_, tensor_dim), shard_pg in zip(layout_spec.shard_axes, shard_pgs):
        shard_size = dist.get_world_size(shard_pg) if shard_pg is not None else 1
        gathered_shapes = _allgather_shapes_for_group(current_shapes, shard_pg, shard_size, local_device)
        cur_rank = dist.get_rank(shard_pg) if shard_pg is not None and shard_size > 1 else 0
        next_shapes: List[Tuple[int, ...]] = []

        for record_idx, _ in enumerate(hsdp_group.records):
            rank_shapes = tuple(
                tuple(gathered_shapes[rank_idx][record_idx])
                for rank_idx in range(shard_size)
            )
            tensor_dim_norm = tensor_dim % len(rank_shapes[0])
            split_sizes = tuple(shape[tensor_dim_norm] for shape in rank_shapes)
            pad_shape = tuple(max(shape[dim] for shape in rank_shapes) for dim in range(len(rank_shapes[0])))
            stage_metas_per_param[record_idx].append(
                ParamShardStageMeta(
                    tensor_dim=tensor_dim_norm,
                    shard_size=shard_size,
                    cur_rank_in_shard_group=cur_rank,

                    pad_shape=pad_shape,
                )
            )

            output_shape = list(rank_shapes[cur_rank])
            if shard_size > 1:
                output_shape[tensor_dim_norm] = sum(split_sizes)
            next_shapes.append(tuple(output_shape))

        current_shapes = next_shapes

    param_to_meta: Dict[torch.Tensor, ParamShardMeta] = {}
    updated_records = []
    for record, stage_metas in zip(hsdp_group.records, stage_metas_per_param):
        shard_meta = make_param_shard_meta(record.param, stage_metas)
        param_to_meta[record.param] = shard_meta
        updated_records.append(record.__class__(index=record.index, param=record.param, shard_meta=shard_meta))

    hsdp_group.records = updated_records
    return param_to_meta


def build_pad_ns_inputs(
        comm_params: List[torch.Tensor],

        param_to_ns_input: Dict[torch.Tensor, torch.Tensor],
        param_shard_metadata: Optional[Dict[torch.Tensor, ParamShardMeta]] = None,
) -> List[torch.Tensor]:
    """Build aligned local NS inputs for all communication params."""
    if not comm_params:
        return []

    ref_tensor = next(iter(param_to_ns_input.values()), None)
    default_dtype = ref_tensor.dtype if ref_tensor is not None else to_local_if_dtensor(comm_params[0].data).dtype
    local_inputs: List[torch.Tensor] = []

    for param in comm_params:
        ns_input = param_to_ns_input.get(param)
        if ns_input is not None:
            local_inputs.append(ns_input)
            continue

        shard_meta = None
        if param_shard_metadata is not None:
            shard_meta = param_shard_metadata.get(param)
        if shard_meta is None and hasattr(param, "shard_meta"):
            shard_meta = getattr(param, "shard_meta")
        if shard_meta is None:
            local_shape = tuple(to_local_if_dtensor(param.data).shape)
        else:
            local_shape = shard_meta.local_shape

        local_inputs.append(
            torch.zeros(
                local_shape,
                dtype=default_dtype,
                device=to_local_if_dtensor(param.data).device,

                device=to_local_if_dtensor(param.data).device,
            )
        )

    return local_inputs


def chunk_update_by_layout(
        global_update: torch.Tensor,

        layout_spec: ParamLayoutSpec,
        param_shard_meta: Optional[ParamShardMeta] = None,
) -> torch.Tensor:
    """Slice a full update back to the local shard using narrow."""
    if not hasattr(param, "device_mesh") or layout_spec is None or not layout_spec.shard_axes:
        return global_update

    device_mesh = param.device_mesh
    mesh_coordinates = device_mesh.get_coordinate()

    shard_axes = layout_spec.shard_axes
    local_update = global_update

    # ``stage_metas`` are built while all-gathering local shards back to the
    # global tensor shape, so slicing a global update back to the local shard
    # must reverse that order.
    for axis_idx, (mesh_dim, tensor_dim) in reversed(list(enumerate(shard_axes))):
        num_chunks = device_mesh.size(mesh_dim)

        if num_chunks <= 1:
            continue

        local_rank = mesh_coordinates[mesh_dim]
        if param_shard_meta is not None and axis_idx < len(param_shard_meta.stage_metas):
            split_sizes = param_shard_meta.stage_metas[axis_idx].split_sizes
        else:
            full_chunk_size = (local_update.size(tensor_dim) + num_chunks - 1) // num_chunks
            split_sizes = tuple(
                min(full_chunk_size, max(local_update.size(tensor_dim) - rank * full_chunk_size, 0))
                for rank in range(num_chunks)
            )
        local_update = local_update.narrow(tensor_dim, sum(split_sizes[:local_rank]), split_sizes[local_rank])

    if not local_update.is_contiguous():
        local_update = local_update.contiguous()

    return local_update


def _get_or_alloc_buffer(
        cache: Optional[Dict],

        dtype: torch.dtype,
        device: torch.device,
) -> torch.Tensor:
    """Return a cached buffer with at least `numel` elements."""
    if cache is not None and key in cache:
        buf = cache[key]
        if buf.numel() >= numel:
            return buf
        buf = torch.empty(numel, dtype=dtype, device=device)
        cache[key] = buf
        return buf

    buf = torch.empty(numel, dtype=dtype, device=device)
    if cache is not None:
        cache[key] = buf
    return buf


def _early_return_tensors(
        local_tensors: List[torch.Tensor],

        local_tensors: List[torch.Tensor],
        keep_indices: Optional[set],
) -> List[Optional[torch.Tensor]]:
    """Return tensors directly when no communication is needed."""
    if keep_indices is None:
        return list(local_tensors)
    return [t if i in keep_indices else None for i, t in enumerate(local_tensors)]


@dataclass(frozen=True)
class _GatherParamMeta:

        shard_size: int,
        stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None,
) -> Tuple[List[torch.Tensor], List[_GatherParamMeta], int]:
    """Move shard dim to dim0 and compute padding metadata."""
    gather_inputs: List[torch.Tensor] = []
    param_meta: List[_GatherParamMeta] = []
    total_padded_numel = 0

    for idx, gather_input in enumerate(current_tensors):
        normalized_dim = tensor_dim % gather_input.dim()
        if normalized_dim != 0 or not gather_input.is_contiguous():
            gather_input = gather_input.movedim(normalized_dim, 0).contiguous()
        stage_meta = stage_metas[idx] if stage_metas is not None else None
        if stage_meta is None:
            split_sizes = tuple(gather_input.shape[0] for _ in range(shard_size))
            padded_numel_raw = gather_input.numel()
        else:
            split_sizes = stage_meta.split_sizes
            pad_shape = list(stage_meta.pad_shape)
            if normalized_dim != 0:
                pad_shape[0], pad_shape[normalized_dim] = pad_shape[normalized_dim], pad_shape[0]
            padded_numel_raw = math.prod(pad_shape)
        param_meta.append(
            _GatherParamMeta(
                offset=total_padded_numel,
                actual_numel=gather_input.numel(),
                padded_numel=(

                rest_shape=tuple(gather_input.shape[1:]),
                split_sizes=split_sizes,
            )
        )
        gather_inputs.append(gather_input)
        total_padded_numel += param_meta[-1].padded_numel

    return gather_inputs, param_meta, total_padded_numel


def _pack_and_allgather(
        gather_inputs: List[torch.Tensor],

        shard_size: int,
        buffer_cache: Optional[Dict],
) -> torch.Tensor:
    """Pack local shards into one buffer, all-gather, return gathered view."""
    pack_buffer = _get_or_alloc_buffer(
        buffer_cache, ("fused_allgather", axis_idx, dtype, device), total_padded_numel,
        dtype, device,
    )[:total_padded_numel]

    pack_buffer.zero_()
    for gi, meta in zip(gather_inputs, param_meta):
        pack_buffer[meta.offset:meta.offset + meta.actual_numel].copy_(gi.view(-1))

    gathered_buffer = _get_or_alloc_buffer(
        buffer_cache,
        ("fused_allgather_out", axis_idx, dtype, device),
        total_padded_numel * shard_size,
        dtype, device,

        total_padded_numel * shard_size,
        dtype, device,
    )[:total_padded_numel * shard_size]

    dist.all_gather_into_tensor(gathered_buffer, pack_buffer, group=shard_pg)
    return gathered_buffer.view(shard_size, total_padded_numel)


def _unpack_gathered_results(
        gathered_view: torch.Tensor,

        is_last_axis: bool,
        keep_indices: Optional[set],
) -> List[Optional[torch.Tensor]]:
    """Slice gathered buffer back to per-parameter full tensors."""
    new_tensors: List[Optional[torch.Tensor]] = []
    for i in range(n_params):
        if is_last_axis and keep_indices is not None and i not in keep_indices:
            new_tensors.append(None)
            continue

        meta = param_meta[i]
        rest_shape = meta.rest_shape
        rest_numel = math.prod(rest_shape) if rest_shape else 1
        param_slice = gathered_view[:, meta.offset:meta.offset + meta.padded_numel]

        param_chunks = []
        for rank in range(shard_size):
            valid_numel = meta.split_sizes[rank] * rest_numel
            if valid_numel == 0:
                continue
            param_chunks.append(param_slice[rank, :valid_numel].contiguous().view(-1))

        if param_chunks:
            param_data = torch.cat(param_chunks, dim=0)
        else:
            param_data = gather_inputs[i].new_empty((0,))

        result = param_data.view(sum(meta.split_sizes), *rest_shape)
        tensor_dim_norm = tensor_dim % current_tensors[i].dim()
        if tensor_dim_norm == 0:
            new_tensors.append(result)
        else:
            new_tensors.append(result.movedim(0, tensor_dim_norm))

    return new_tensors


def fused_allgather_dtensor_params(
        local_tensors: List[torch.Tensor],

        buffer_cache: Optional[Dict] = None,
        keep_indices: Optional[set] = None,
) -> List[Optional[torch.Tensor]]:
    """Fuse many parameter shards into one all-gather per shard axis."""
    if not shard_pgs or not local_tensors:
        return _early_return_tensors(local_tensors, keep_indices)

    n_params = len(local_tensors)
    device = local_tensors[0].device
    dtype = local_tensors[0].dtype
    alignment_bytes = 512
    element_size = local_tensors[0].element_size()
    alignment_elements = max(1, alignment_bytes // element_size)

    active_axes = []
    for axis_idx, ((_, tensor_dim), shard_pg) in enumerate(zip(layout_spec.shard_axes, shard_pgs)):
        if shard_pg is None:
            continue
        shard_size = dist.get_world_size(shard_pg)
        if shard_size <= 1:
            continue
        active_axes.append((axis_idx, tensor_dim, shard_pg, shard_size))

    if not active_axes:
        return _early_return_tensors(local_tensors, keep_indices)

    current_tensors: List[torch.Tensor] = list(local_tensors)

    for active_pos, (axis_idx, tensor_dim, shard_pg, shard_size) in enumerate(active_axes):
        is_last_axis = active_pos == len(active_axes) - 1
        stage_metas: Optional[List[Optional[ParamShardStageMeta]]] = None
        if param_shard_metadata is not None:
            stage_metas = []
            for meta in param_shard_metadata:
                if meta is None or axis_idx >= len(meta.stage_metas):
                    stage_metas.append(None)
                else:
                    stage_metas.append(meta.stage_metas[axis_idx])

        gather_inputs, param_meta, total_padded_numel = _prepare_gather_inputs(
            current_tensors, tensor_dim, alignment_elements, shard_size, stage_metas,
        )

        gathered_view = _pack_and_allgather(
            gather_inputs, param_meta, total_padded_numel,
            axis_idx, dtype, device, shard_pg, shard_size, buffer_cache,
        )

        new_tensors = _unpack_gathered_results(
            gathered_view, gather_inputs, param_meta,
            current_tensors, tensor_dim, shard_size,
            n_params, is_last_axis, keep_indices,
        )

            current_tensors, tensor_dim, shard_size,
            n_params, is_last_axis, keep_indices,
        )

        if is_last_axis:
            return new_tensors

        current_tensors = new_tensors  # type: ignore[assignment]

    return current_tensors

        # Cache Hit Check
        if cache_key in self._split_sub_pg_cache:
            sub_pg_map = self._split_sub_pg_cache[cache_key]
            for sub_pg in sub_pg_map.values():
                if sub_pg is not None:
                    try:
                        dist.get_rank(group=sub_pg)
                        return sub_pg

def _get_rank() -> int:
    try:
        if torch.distributed.is_initialized():
            return torch.distributed.get_rank()
    except (ImportError, RuntimeError):
        pass
    return int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")))


def info_rank0(self, msg, *args, **kwargs) -> None:

def warning_rank0(self, msg, *args, **kwargs) -> None:
    """Log a warning message only on rank zero."""
    if _get_rank() == 0:
        kwargs.setdefault("stacklevel", 2)
        self.warning(msg, *args, **kwargs)


def debug_rank0(self, msg, *args, **kwargs) -> None:
    """Log a debug message only on rank zero."""
    if _get_rank() == 0:
        kwargs.setdefault("stacklevel", 2)
        self.debug(msg, *args, **kwargs)


def get_device_count() -> int:
    """Return the active accelerator count, defaulting to 1."""

def get_device_count() -> int:
    """Return the active accelerator count, defaulting to 1."""
    npu = getattr(torch, "npu", None)
    if npu is not None and npu.is_available():
        return npu.device_count()
    if torch.cuda.is_available():
        return torch.cuda.device_count()
    return 1


def get_current_device() -> torch.device:

def get_current_device() -> torch.device:
    """Return the current accelerator device, or CPU when none is available."""
    npu = getattr(torch, "npu", None)
    if npu is not None and npu.is_available():
        return torch.device("npu", npu.current_device())
    if torch.cuda.is_available():
        return torch.device("cuda", torch.cuda.current_device())
    return torch.device("cpu")


def empty_accelerator_cache() -> None:

def empty_accelerator_cache() -> None:
    """Clear the active accelerator cache when supported."""
    npu = getattr(torch, "npu", None)
    if npu is not None and npu.is_available():
        npu.empty_cache()
    elif torch.cuda.is_available():
        torch.cuda.empty_cache()


_INSTALLED = False

def _install_logger_methods() -> None:
    global _INSTALLED
    if _INSTALLED:
        return
    logging.Logger.info_rank0 = info_rank0
    logging.Logger.warning_rank0 = warning_rank0
    logging.Logger.debug_rank0 = debug_rank0
    _INSTALLED = True

    if any(end <= start for start, end in zip(boundaries[:-1], boundaries[1:])):
        raise ValueError("cu_seq_lens must be strictly increasing")

    if boundaries[-1] != micro_batch_size * seq_length:
        raise ValueError(
            "cu_seq_lens must cover the physical batch length "
            f"({micro_batch_size * seq_length}), but got {boundaries[-1]}"
        )

    swa_mask = None
    if sliding_window is not None:
        positions = torch.arange(seq_length, dtype=torch.int64, device=device)
        # Broadcast query and key positions to identify tokens outside the sliding window.
        outside_window = positions.unsqueeze(1) - positions.unsqueeze(0) > sliding_window
        swa_mask = attention_mask & ~outside_window.unsqueeze(0).unsqueeze(0)

    return attention_mask, swa_mask

        for field in ("input_ids", "labels"):
            values = [model_sample[field] for model_sample in model_samples]
            packed_batch[field] = torch.cat(values, dim=-1).unsqueeze(0)

        input_ids = packed_batch.get("input_ids")
        if input_ids is None or packed_batch.get("labels") is None:
            raise ValueError("packed batch requires input_ids and labels")

        packed_seq_len = input_ids.shape[-1]
        pad_len = (-packed_seq_len) % self.sequence_parallel_size
        if pad_len:
            input_padding = input_ids.new_zeros((1, pad_len))
            label_padding = packed_batch["labels"].new_full((1, pad_len), IGNORE_INDEX)
            packed_batch["input_ids"] = torch.cat((input_ids, input_padding), dim=-1)
            packed_batch["labels"] = torch.cat((packed_batch["labels"], label_padding), dim=-1)

        seq_lens = model_samples[0]["input_ids"].new_tensor(
            [model_sample["input_ids"].shape[-1] for model_sample in model_samples]

                data_sharding=getattr(dataloader_target, "data_sharding", False),
                seed=training_config.seed if training_config.seed is not None else default_seed,
            )

        dataloaders[split_index] = dataloader_target.build(
            dataset=dataset,
            collate_fn=collate_fn,
            batch_sampler=batch_sampler,
            batch_size=training_config.micro_batch_size,

        )
        batch_samplers[split_index] = batch_sampler

    logger.debug("Finished building train/valid/test DataLoaders")
    return tuple(dataloaders), tuple(batch_samplers)


class FixedBatchDataLoader(StatefulDataLoader):
    """Build fixed-sample batches while retaining Trainer iterator policy."""

    def __init__(self, source_dataset: Any) -> None:
        """Store a Dataset that emits source items with stable output indices."""
        self.source_dataset = source_dataset

    @staticmethod
    def put_source_item(source_item: Any, batcher: TextTokenBatcher) -> None:
        """Flatten one indexed source item and retain each sample index."""
        model_samples_item, output_index = source_item
        model_samples = _normalize_source_samples(model_samples_item)
        for sample_idx, model_sample in enumerate(model_samples):

        model_samples = _normalize_source_samples(model_samples_item)
        for sample_idx, model_sample in enumerate(model_samples):
            batcher.put_item(model_sample, (output_index, sample_idx))

    @staticmethod
    def get_buffer_state(batcher: TextTokenBatcher) -> list[Any]:
        """Return compact output-index entries for buffered ModelSamples."""
        if len(batcher.buffer) != len(batcher.buffer_output_indices):
            raise RuntimeError("Dynamic sample and output-index buffers are inconsistent")

        """Retain the source Dataset and an optional index replay adapter."""
        self.source_dataset = dataset
        self.replay_dataset = replay_dataset

    @staticmethod
    def put_source_item(source_item: Any, batcher: TextTokenBatcher) -> None:
        """Flatten one streaming source item into the runtime buffer."""
        model_samples = _normalize_source_samples(source_item)
        for model_sample in model_samples:
            batcher.put_item(model_sample)

        model_samples = _normalize_source_samples(source_item)
        for model_sample in model_samples:
            batcher.put_item(model_sample)

    @staticmethod
    def get_buffer_state(batcher: TextTokenBatcher) -> list[Any]:
        """Return full samples already pulled beyond the streaming cursor."""
        if len(batcher.buffer) != len(batcher.buffer_output_indices):
            raise RuntimeError("Dynamic sample and output-index buffers are inconsistent")

        parallel_batch = self.tp_broadcaster.broadcast(
            cp_local_batch,
            cu_seq_lens,
        )
        self._log_batch_flow(canonical_batch, parallel_batch)

        position_ids = self._build_local_position_ids(
            parallel_batch["input_ids"],
            parallel_batch["cu_seq_lens"],

        model_inputs, loss_inputs = self._split_model_and_loss_inputs(parallel_batch)

        return model_inputs, loss_inputs

    def _log_batch_flow(
            self,
            canonical_batch: Mapping[str, Any] | None,
            parallel_batch: Mapping[str, Any],
    ) -> None:

            canonical_batch: Mapping[str, Any] | None,
            parallel_batch: Mapping[str, Any],
    ) -> None:
        """Log the resolved source and parallel batch shapes once."""
        if not self._batch_flow_logged:
            source_shape = None if canonical_batch is None else tuple(canonical_batch["input_ids"].shape)
            local_shape = tuple(parallel_batch["input_ids"].shape)
            num_boundaries = 0 if parallel_batch["cu_seq_lens"] is None else parallel_batch["cu_seq_lens"].numel()
            logger.debug(
                "Parallel batch flow: source=%s, tp_rank=%d/%d, cp_rank=%d/%d, "
                "source_owner=%s, source_shape=%s, local_shape=%s, global_boundaries=%d",
                self.source_type,
                self.parallel_context.tp_rank,

                local_shape,
                num_boundaries,
                enabled=True,
            )
            self._batch_flow_logged = True

    def _read_source_batch(self, data_iterator: Any) -> Mapping[str, Any] | None:
        """Read one complete batch on TP rank zero of each CP coordinate."""
        if self.parallel_context.build_on_rank():

class OnlineBoundaryResolver:
    """Read global cumulative sequence boundaries emitted by Online packing."""

    @staticmethod
    def resolve(canonical_batch: Mapping[str, Any]) -> Any:
        """Read leading-zero ``cu_seq_lens`` from an Online batch.

        Args:
            canonical_batch: Normalized batch produced by the Online packing collator.

    # 4. Prefer the numeric shard order encoded by the corpus filenames.
    # Unknown naming schemes remain usable in their original walk order.
    try:
        def numeric_shard_key(path: str) -> tuple[int, int]:
            """Extract numeric corpus and shard identifiers from a path."""
            path_parts = os.path.basename(path).split("_")
            return int(path_parts[2]), int(path_parts[-1])

        file_paths.sort(key=numeric_shard_key)
    except (IndexError, TypeError, ValueError):
        logger.warning("Cannot sort indexed files with the numeric filename rule; keeping walk order")

    # 5. Automatically discovered sources use equal blend weights.

        """Delegate GPT Dataset option normalization."""
        config = build_gpt_dataset_config(data_paths, self.data_config)
        return config

    @staticmethod
    def _select_dataset_type(config: GPTDatasetConfig) -> type:
        """Select GPT, Mock GPT, or MR GPT Dataset."""
        if config.mock:
            return MockGPTDataset

    def _build_interleaved_locations(self) -> list[tuple[int, int]]:
        """Alternate over non-empty Datasets until every sample is exposed."""
        max_size = max(len(dataset) for dataset in self.datasets)
        locations = []
        for sample_id in range(max_size):
            for dataset_id, dataset in enumerate(self.datasets):
                if sample_id < len(dataset):
                    locations.append((dataset_id, sample_id))
        return locations

                local_cu_seq_lens = torch.empty((num_boundaries,), dtype=torch.int32, device=self.device)

        # Broadcast only model fields that cannot be regenerated locally.
        for field in ("input_ids", "labels"):
            tensor = parallel_batch.get(field)
            if tensor is None:
                raise ValueError(f"parallel batch requires {field!r}")
            dist.broadcast(tensor, group=tp_group, group_src=0)

        # Packed boundaries remain global across CP ranks and variable in size.
        if local_cu_seq_lens is not None:
            dist.broadcast(local_cu_seq_lens, group=tp_group, group_src=0)

        """
        self.tokenizer = tokenizer
        self.chat_template = chat_template

    @staticmethod
    def __call__(sample: Any) -> Any:
        """Return the input sample without modification."""
        return sample

"""Lazy plaintext or conversation transforms for Online LLM sources."""

from __future__ import annotations

from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Any, TypeAlias

from torch.utils.data import IterableDataset  # pylint: disable=forbidden-backend-import

        )
        logger.debug("Wrapped mapping Dataset with transform=%s", type(transform).__name__)
        return transformed_dataset

    if isinstance(dataset, Iterable):
        transformed_dataset = _LLMIterableTransformDataset(
            dataset,
            transform,
            skip_invalid_samples=skip_invalid_samples,

import argparse
import glob
import gzip
import json
import logging
import math
import multiprocessing
import os
import sys

# Store generated samples in the indexed ``.bin/.idx`` format.
from hyper_parallel.data.tools import io as indexed_dataset

logger = logging.getLogger(__name__)


class CustomLanguageVars(PunktLanguageVars):
    """Preserve newline runs when Punkt detects sentence boundaries."""

    for name in names:
        process = multiprocessing.Process(target=_split_sentences_worker, args=(args, workers, name, queue))
        process.start()
        processes.append(process)
    split_results = _wait_for_processes(processes, queue)
    if any(result is not None for result in split_results):
        raise RuntimeError("Sentence-splitting workers returned unexpected results")


def _encode_partitions(
    args: argparse.Namespace,

    performance = {}
    input_files = _resolve_input_files(args.dataset_name_or_path)
    for workers in worker_candidates:
        logger.info("Processing data with %s workers", workers)
        workers_per_partition = workers // args.partitions

        if args.split_sentences:
            nltk.download("punkt", quiet=True, download_dir=os.environ.get("NLTK_DATA"))