Diff Coverage

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, NamedTuple, Optional

import torch  # pylint: disable=forbidden-backend-import
import torch_npu

import torch  # pylint: disable=forbidden-backend-import
import torch_npu


class _AttentionInputs(NamedTuple):
    """Inputs normalized for the NPU fusion-attention operator."""

    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
    input_layout: str
    attention_mask: Optional[torch.Tensor]
    sparse_mode: int


def _npu_attention_mask(attention_mask: torch.Tensor) -> torch.Tensor:
    """Convert the Transformers attention mask to the NPU mask convention."""

            npu_mask = torch.ones((2048, 2048), dtype=torch.bool, device=query.device).triu(diagonal=1)
            sparse_mode = 3
        else:
            npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
        return _AttentionInputs(query, key, value, "TND", npu_mask, sparse_mode)
    if attention_mask is None and is_causal:
        return _AttentionInputs(query, key, value, "BNSD", _causal_attention_mask(query, key, sliding_window), 0)
    npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
    return _AttentionInputs(query, key, value, "BNSD", npu_mask, sparse_mode)


def npu_fusion_attention_forward(
    module: torch.nn.Module,

    target_flat = target.flatten()

    target_mask = (target_flat >= vocab_start) & (target_flat < vocab_end)

    target_mask = target_mask & (target_flat != ignore_index)

    if reduction == "none":
        loss = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
    else:

    if target_mask.any():
        local_target = target_flat[target_mask] - vocab_start

        selected_log_probs = log_probs.reshape(-1, log_probs.shape[-1])[
            torch.where(target_mask)[0], local_target
        ]

        if weight is not None:
            sample_weights = weight[target_flat[target_mask]]
            selected_log_probs = selected_log_probs * sample_weights
            total_weight = sample_weights.sum().reshape(1)
        else:
            total_weight = torch.tensor(

            total_weight = torch.tensor(
                target_mask.sum().item(), dtype=log_probs.dtype, device=log_probs.device
            ).reshape(1)

        selected_log_probs = -selected_log_probs

        if reduction == "none":
            loss = torch.zeros(batch_size, dtype=log_probs.dtype, device=log_probs.device)
            loss[target_mask] = selected_log_probs
            loss = loss.reshape(target.shape)
        elif reduction == "sum":
            loss = selected_log_probs.sum().unsqueeze(0)
        else:
            loss = selected_log_probs.sum().unsqueeze(0)
    else:
        if reduction == "none":
            loss = torch.zeros(
                batch_size, dtype=log_probs.dtype, device=log_probs.device

    Returns:
        Summed MTP loss over all depths.
    """
    total_mtp_loss = torch.tensor(0.0, device=labels.device, dtype=torch.float32)
    for logits in mtp_per_depth_logits:
        logits_shifted = logits[..., :-1, :].contiguous()
        labels_shifted = labels[..., 1:].contiguous()
        depth_loss = loss_fn(
            logits_shifted.view(-1, logits_shifted.size(-1)),

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any, NamedTuple

# This package provides PyTorch-specific high-performance modules.
# pylint: disable=forbidden-backend-import
import torch  # pylint: disable=forbidden-backend-import

)
from hyper_parallel.components.functional.npu_fusion_attention import resolve_packed_sequence_lengths


class _ProjectedAttentionStates(NamedTuple):
    """Projected states shared by the DSA attention implementations."""

    q_resid: torch.Tensor
    absorbed_query: torch.Tensor
    kv_nope: torch.Tensor
    q_rot: torch.Tensor
    k_rot: torch.Tensor
    kv_weight: torch.Tensor


def apply_mome(
    hidden_states: torch.Tensor,

        batch_size, seq_length, -1
    )


def _reshape_dsa_attention_states(
    attention_states: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Flatten the batch and sequence dimensions for the DSA auxiliary loss."""
    return tuple(tensor.reshape(-1, tensor.shape[2], tensor.shape[3]) for tensor in attention_states)


@module_replacement
class DeepseekV32DSAAttention(nn.Module):

            output_dim=self.kv_lora_rank,
        )
        kv_nope = self.kv_a_layernorm(kv_nope).view(batch_size, seq_length, 1, self.kv_lora_rank)
        k_rot = k_rot.view(batch_size, seq_length, 1, self.qk_rope_head_dim)
        return _ProjectedAttentionStates(q_resid, absorbed_query, kv_nope, q_rot, k_rot, kv_weight)

    def _project_index_states(
        self,
        hidden_states: torch.Tensor,

    ) -> torch.Tensor:
        """Attach the DSA KL auxiliary loss when training enables it."""
        if not (self.training and not self.freeze_dsa and self.dsa_loss_coeff):
            return attn_output
        attention_states_tnd = _reshape_dsa_attention_states(attention_states)
        aux_loss = dsa_kl_loss(
            *index_states, attention_states_tnd[0], attention_states_tnd[1],
            topk_indices, *softmax_stats, attention_states_tnd[2], attention_states_tnd[3],
            actual_q_len, actual_kv_len, self.scaling, self.dsa_loss_coeff,

    ) -> tuple[torch.Tensor, None]:
        """Run causal or packed DSA with the NPU sparse-attention kernels."""
        batch_size, seq_length = hidden_states.shape[:-1]
        self._validate_forward_inputs(hidden_states, attention_mask, past_key_values, position_ids, kwargs)
        attention_states = self._project_attention_states(hidden_states)
        attention_states = (
            attention_states[:3]
            + _apply_attention_rope(
                attention_states[3],
                attention_states[4],

                interleaved=self.rotary_interleaved,
            )
            + attention_states[5:]
        )
        index_states = self._project_index_states(hidden_states, attention_states[0], position_embeddings)
        actual_lengths = dict(kwargs)
        actual_lengths["actual_seq_len"] = actual_seq_len
        actual_q_len, actual_lengths = resolve_packed_sequence_lengths(
            actual_lengths,
            batch_size * seq_length,
            batch_size * seq_length,
        )
        actual_lengths = tuple(
            self._get_actual_seq_len(lengths, batch_size, seq_length, hidden_states.device)
            for lengths in (actual_q_len, actual_lengths)
        )
        index_states = dsa_indexer(*index_states, *actual_lengths, self.index_topk)
        attn_output = dsa_sparse_attention(
            *attention_states[1:5], index_states[0], self.scaling, *actual_lengths
        )
        attn_output = self._apply_auxiliary_loss(
            attn_output[0],

            seq_length=seq_length,
            kv_lora_rank=self.kv_lora_rank,
            value_head_dim=self.v_head_dim,
        )
        return self.o_proj(attn_output), None


@module_replacement
class DSAAttention(nn.Module):

            output_dim=self.kv_lora_rank,
        )
        kv_nope = self.k_layernorm(kv_nope).view(batch_size, seq_length, 1, self.kv_lora_rank)
        k_rot = k_rot.view(batch_size, seq_length, 1, self.qk_rope_head_dim)
        return _ProjectedAttentionStates(q_resid, absorbed_query, kv_nope, q_rot, k_rot, kv_weight)

    def _project_index_states(
        self,
        hidden_states: torch.Tensor,

    ) -> torch.Tensor:
        """Attach the DSA KL auxiliary loss when training enables it."""
        if not (self.training and not self.freeze_dsa and self.dsa_loss_coeff):
            return attn_output
        attention_states_tnd = _reshape_dsa_attention_states(attention_states)
        aux_loss = dsa_kl_loss(
            *index_states, attention_states_tnd[0], attention_states_tnd[1],
            topk_indices, *softmax_stats, attention_states_tnd[2], attention_states_tnd[3],
            actual_q_len, actual_kv_len, self.qk_head_dim**-0.5, self.dsa_loss_coeff,

            actual_q_len, actual_kv_len, self.qk_head_dim**-0.5, self.dsa_loss_coeff,
        )
        return aux_loss_auto_scale(attn_output, aux_loss)

    def _project_forward_output(
        self,
        attn_output: torch.Tensor,
        return_bias: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:

        attn_output: torch.Tensor,
        return_bias: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Apply the output projection and select its bias result."""
        output, bias = self.linear_proj(attn_output)
        return (output, bias) if return_bias else (output, None)

    def _forward_impl(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
        actual_seq_len: torch.Tensor | Sequence[int] | None,

        kwargs: dict[str, Any],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Run the validated DSA forward path."""
        batch_size, seq_length = hidden_states.shape[:-1]
        attention_states = self._project_attention_states(hidden_states, mome_mask)
        attention_states = (
            attention_states[:3]
            + _apply_attention_rope(
                attention_states[3],
                attention_states[4],

                interleaved=self.rotary_interleaved,
            )
            + attention_states[5:]
        )
        index_states = self._project_index_states(hidden_states, attention_states[0], position_embeddings)
        actual_lengths = dict(kwargs)
        actual_lengths["actual_seq_len"] = actual_seq_len
        actual_q_len, actual_lengths = resolve_packed_sequence_lengths(
            actual_lengths,
            batch_size * seq_length,
            batch_size * seq_length,
        )
        actual_lengths = tuple(
            self._get_actual_seq_len(lengths, batch_size, seq_length, hidden_states.device)
            for lengths in (actual_q_len, actual_lengths)
        )
        index_states = dsa_indexer(*index_states, *actual_lengths, self.index_topk)
        attn_output = self._run_sparse_attention(
            attention_states[1:5], index_states[0], *actual_lengths, batch_size, seq_length
        )
        attn_output = self._apply_auxiliary_loss(
            attn_output[0],

                mome_mask,
                self.o_conv,
                fused=self.use_fused_mome,
            )
        return self._project_forward_output(attn_output, return_bias)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        attention_mask: torch.Tensor | None = None,

        mome_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Run DSA with its configured MOME and parameter-sink paths."""
        self._validate_forward_inputs(
            attention_mask, kv_reuse_states, past_key_values, cache_position, output_attentions
        )
        return self._forward_impl(
            hidden_states,
            position_embeddings,
            actual_seq_len,
            mome_mask,

# limitations under the License.
# ============================================================================
"""Collators that produce the micro-batch groups consumed by Trainer."""

__all__ = [
    "MakeMicroBatchCollator",
    "calculate_num_micro_batches",
]

# limitations under the License.
# ============================================================================
"""Dataloader components used by Trainer targets."""

__all__ = ["DataLoader"]

from collections.abc import Callable
from typing import Any, Optional

    for key, value in batch.items():
        if key == "qkv_format" or not isinstance(value, torch.Tensor) or value.ndim < 1:
            continue
        if key in ("seq_lens", "seq_lens_padded"):
            continue  # recomputed separately, not padded
        if key == "position_ids":
            # position_ids increment-pad: continue incrementing from the last value
            last = value[..., -1:].to(torch.long)
            increment = torch.arange(1, pad_len + 1, device=value.device,

# limitations under the License.
# ============================================================================
"""Build the VLM multimodal sample transform."""

__all__ = ["VLMChatTransform", "build_vlm_data_transform"]

import re
from typing import Any, Optional

        """Store the processor and target sequence length."""
        self.processor = processor
        self.max_seq_len = max_seq_len

    @staticmethod
    def _normalize_messages(messages: Any, images: Any = None) -> Any:
        """Split ``<image>``/``<video>`` string placeholders into content-list parts."""
        if images is None:
            images = []
        image_queue = list(images)

# limitations under the License.
# ============================================================================
"""Build the VLM processor."""

__all__ = ["build_processor"]

from typing import Any

from transformers import AutoProcessor

# limitations under the License.
# ============================================================================
"""Build the VLM micro-batch collator."""

__all__ = ["VLMCollator", "build_vlm_collator"]

from typing import Any, Optional

import torch

# limitations under the License.
# ============================================================================
"""Build the VLM dataset from a LLaVA-style JSON list."""

__all__ = ["VLMDataset", "build_vlm_dataset"]

import json
import os
from collections.abc import Callable

        EP_ARCHETYPE_SUGGESTIONS,
    )
    for fqn, spec in plan.modules.items():
        if not spec.is_boundary:
            continue
        if (getattr(spec, "_ep_size", 0)  # pylint: disable=protected-access
                and getattr(spec, "local_compute_fn", None) is None
                and getattr(spec, "region_dispatch", None) is not False):
            suggestion = ""

    fn_names = {param.name for param in fn_params}
    required = []
    for param in fwd_params:
        if param.default is not inspect.Parameter.empty:
            continue
        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
            continue
        required.append(param.name)
    missing = [name for name in required if name not in fn_names]
    if missing:
        raise TypeError(

            state_dict/optimizer visibility is unchanged; only ``F.linear``
            inside the region sees a bias-free Linear.
            """
            bias = __owner.bias
            __owner._parameters["bias"] = None  # pylint: disable=protected-access
            try:
                return __original(*args, **kwargs)
            finally:
                __owner._parameters["bias"] = bias  # pylint: disable=protected-access

    elif spec is not None and isinstance(
            getattr(spec, "inner_wrapper", None), str):
        source = "explicitly specified (registry)"
    else:
        source = "explicitly specified (Target)"
    logger.info("inner-wrap: %s target=%s <- wrapper %r (%s)",
                type(module).__name__, target_name, name, source)
    return name

def _rollback_inner_rewrite(target, saved_state, committed_secondaries):
    """Restore primary and secondary targets after a failed inner rewrite."""
    for secondary_target, secondary_state in reversed(committed_secondaries):
        for attr, saved in secondary_state.items():
            _restore_attr(secondary_target, attr, saved)
    for attr, saved in saved_state.items():
        _restore_attr(target, attr, saved)

            holder.register_parameter(name, p)
        *path, leaf = parent_path.split(".")
        obj = module
        for seg in path:
            obj = obj[int(seg)] if seg.isdigit() else getattr(obj, seg)
        setattr(obj, leaf, holder)   # replace the original ModuleList (original expert params freed)


def _stack_moe_experts(module: nn.Module, ep_stack: Dict[str, List[str]]) -> None:

            shape = param_shapes.get(prefix + sources[0]) if sources else None
            if shape is not None:
                shape = (len(sources), *shape)
        if shape is None:
            return  # created later (inner_target factory) — nothing to check
        for axis, p in (placement or {}).items():
            if not isinstance(p, Shard):
                continue
            size = axis_sizes.get(axis, 1)

                continue
            axis_name = getattr(axis, "value", axis)  # MeshAxisName → "tp"
            dim = p.dim + len(shape) if p.dim < 0 else p.dim
            if dim >= len(shape):
                raise ValueError(
                    f"plan-time shard check failed: {full!r} has shape "
                    f"{tuple(shape)} but boundary {fqn!r} declares "
                    f"{{{axis_name}: Shard({p.dim})}} — dim {p.dim} is out "
                    f"of range for a {len(shape)}D parameter; fix the "

    flattened_states = hidden_states.reshape(-1, hidden_states.shape[-1])
    source_indices = torch.arange(
        flattened_states.shape[0], device=flattened_states.device
    ).repeat_interleave(topk_indices.shape[1])
    topk_indices = topk_indices.reshape(-1)
    topk_weights = topk_weights.reshape(-1).to(flattened_states.dtype)
    destination_ranks = torch.div(topk_indices, local_expert_count, rounding_mode="floor")
    dispatch_order = (destination_ranks * global_expert_count + topk_indices).argsort()
    dispatched_states = flattened_states[source_indices[dispatch_order]].contiguous()
    dispatched_indices = topk_indices[dispatch_order].unsqueeze(-1).contiguous()
    send_counts_tensor = torch.bincount(destination_ranks, minlength=ep_size)
    receive_counts_tensor = torch.empty_like(send_counts_tensor)
    dist.all_to_all_single(receive_counts_tensor, send_counts_tensor, group=ep_group)
    return _EPDispatch(

    def configure_fsdp_gradient_sync(self, micro_step: int, num_micro_steps: int) -> None:
        """Configure FSDP gradient synchronization for an external training loop."""
        self._configure_fsdp_gradient_sync(micro_step, num_micro_steps)

    def step_optimizers_and_schedulers(self) -> None:
        """Step optimizers and schedulers after gradient accumulation."""
        optimizers = self.optimizer if isinstance(self.optimizer, list) else [self.optimizer]
        for optimizer in optimizers:
            with SkipDTensorDispatch():
                optimizer.step()
            optimizer.zero_grad()

        schedulers = (
            self.lr_scheduler
            if isinstance(self.lr_scheduler, list)
            else ([self.lr_scheduler] if self.lr_scheduler is not None else [])
        )
        for scheduler in schedulers:
            scheduler.step()

    def train_step(
            self,
            data_iterator: Any,

            config.training.max_grad_norm,
        )

        # Optimizer and scheduler step
        self.step_optimizers_and_schedulers()

        grad_norm_value = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else float(grad_norm)
        self.on_step_end(loss=total_loss, loss_dict=total_loss_dict, grad_norm=grad_norm_value)
        return {

# limitations under the License.
# ============================================================================
"""Training and environment metric collection callback."""

import time
from collections.abc import Mapping, Sequence
from typing import Any

from hyper_parallel.trainer.runtime.distributed import get_world_size_safe
from hyper_parallel.trainer.runtime.distributed import all_reduce
from hyper_parallel.data.constants import IGNORE_INDEX
from hyper_parallel.trainer.runtime.device import get_device_type, get_torch_device

from .base import Callback, TrainerState


class EnvironMeterCallback(Callback):
    """Collect structured training, throughput, and memory metrics.

    The callback is the single producer of ``trainer.step_train_metrics`` and
    ``trainer.step_env_metrics``. Presentation and remote logging callbacks

    ``trainer.step_env_metrics``. Presentation and remote logging callbacks
    consume those dictionaries without recalculating or reducing metrics.
    """

    def __init__(self, trainer: Any) -> None:
        """Initialize per-step and cumulative counters.

        Args:
            trainer: Trainer that owns the callback lifecycle.

        Args:
            trainer: Trainer that owns the callback lifecycle.
        """
        super().__init__(trainer)
        self._step_start_time = 0.0
        self._local_step_tokens = 0
        self._local_step_samples = 0
        self._consumed_tokens = 0
        self._consumed_samples = 0
        self.trainer.step_train_metrics = {}
        self.trainer.step_env_metrics = {}

    @staticmethod
    def _scalar(value: Any, name: str) -> float:
        """Convert a scalar or scalar tensor-like value to ``float``.

        Args:
            value: Scalar value to convert.

        Raises:
            ValueError: If the value cannot be converted to a scalar float.
        """
        item = getattr(value, "item", None)
        if callable(item):
            value = item()
        try:
            return float(value)
        except (TypeError, ValueError) as exc:
            raise ValueError(f"Metric {name!r} must be scalar, but got {value!r}") from exc

    @staticmethod
    def _tensor_numel(value: Any) -> int | None:
        """Return ``value.numel()`` when it exposes a tensor-like interface."""
        numel = getattr(value, "numel", None)
        if not callable(numel):
            return None
        return int(numel())

    @classmethod
    def _batch_tokens(cls, batch: Mapping[str, Any]) -> int:
        """Count text tokens in one micro-batch without mutating it."""
        labels = batch.get("labels")
        if labels is not None and callable(getattr(labels, "sum", None)):
            return int((labels != IGNORE_INDEX).sum().item())

        attention_mask = batch.get("attention_mask")
        attention_mask_shape = getattr(attention_mask, "shape", ())
        if (
            len(attention_mask_shape) <= 2
            and attention_mask is not None
            and callable(getattr(attention_mask, "sum", None))
        ):
            return int(attention_mask.sum().item())

        input_ids = batch.get("input_ids")
        input_numel = cls._tensor_numel(input_ids)
        if input_numel is not None:
            return input_numel
        return 0

    @staticmethod
    def _batch_samples(batch: Mapping[str, Any]) -> int:
        """Count logical samples in one micro-batch."""
        value = batch.get("input_ids")
        if value is None:
            value = batch.get("labels")
        shape = getattr(value, "shape", None)
        if shape is None or len(shape) == 0:
            return 0
        if len(shape) == 1:
            return 1
        return int(shape[0])

    @staticmethod
    def _batch_mapping(value: Any) -> Mapping[str, Any] | None:
        """Return metric inputs for a mapping or prepared runtime batch."""
        if isinstance(value, Mapping):
            return value
        loss_count_inputs = getattr(value, "loss_count_inputs", None)
        if not callable(loss_count_inputs):
            return None
        metric_inputs = loss_count_inputs()
        return metric_inputs if isinstance(metric_inputs, Mapping) else None

    @classmethod
    def _micro_batches(cls, value: Any) -> list[Mapping[str, Any]]:
        """Normalize callback input into a list of mapping micro-batches."""
        if value is None:
            return []
        batch = cls._batch_mapping(value)
        if batch is not None:
            return [batch]
        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
            batches = []
            for item in value:
                batch = cls._batch_mapping(item)
                if batch is not None:
                    batches.append(batch)
            return batches
        return []

    def _metric_group(self) -> Any:
        """Return the DP+CP process group used by loss normalization."""
        dp_cp_mesh = getattr(self.trainer.mesh, "dp_cp_mesh", None)
        if dp_cp_mesh is None:
            return None
        return dp_cp_mesh.get_group()

    def _reduce(self, value: float | int, op: str) -> float:
        """Reduce one scalar metric, with a single-process no-op fallback."""
        if get_world_size_safe() <= 1:
            return float(value)
        reduced = all_reduce(value, op=op, group=self._metric_group())
        return float(reduced)

    def _current_lr(self) -> float:
        """Return the maximum learning rate across scheduler or optimizer groups."""
        schedulers = self.trainer.lr_scheduler
        if schedulers is not None:
            scheduler_list = schedulers if isinstance(schedulers, list) else [schedulers]
            learning_rates = []
            for scheduler in scheduler_list:
                for learning_rate in scheduler.get_last_lr():
                    learning_rates.append(float(learning_rate))
            if learning_rates:
                return max(learning_rates)

        optimizers = self.trainer.optimizer
        optimizer_list = optimizers if isinstance(optimizers, list) else [optimizers]
        learning_rates = []
        for optimizer in optimizer_list:
            if optimizer is None:
                continue
            for param_group in optimizer.param_groups:
                learning_rates.append(float(param_group["lr"]))
        return max(learning_rates, default=0.0)

    def _memory_metrics(self) -> dict[str, float]:
        """Collect maximum accelerator memory metrics, if available."""
        if get_device_type() == "cpu":
            return {}
        device = get_torch_device()
        allocated = self._reduce(device.max_memory_allocated(), op="max")
        reserved = self._reduce(device.max_memory_reserved(), op="max")
        gibibyte = 1024 ** 3
        return {
            "memory/device_max_allocated_gb": allocated / gibibyte,
            "memory/device_max_reserved_gb": reserved / gibibyte,
        }

    def state_dict(self) -> dict[str, int]:
        """Return cumulative metric state for future checkpoint integration."""
        return {
            "consumed_tokens": self._consumed_tokens,
            "consumed_samples": self._consumed_samples,
        }

    def load_state_dict(self, state_dict: dict[str, int]) -> None:
        """Restore cumulative metric state.

        Args:
            state_dict: Mapping produced by :meth:`state_dict`.

        Raises:
            ValueError: If required counters are missing or negative.
        """
        try:
            consumed_tokens = int(state_dict["consumed_tokens"])
            consumed_samples = int(state_dict["consumed_samples"])
        except (KeyError, TypeError, ValueError) as exc:
            raise ValueError(
                "EnvironMeterCallback state must contain integer consumed_tokens and consumed_samples"
            ) from exc
        if consumed_tokens < 0 or consumed_samples < 0:
            raise ValueError("EnvironMeterCallback cumulative counters must be non-negative")
        self._consumed_tokens = consumed_tokens
        self._consumed_samples = consumed_samples

    def on_step_begin(
        self,
        state: TrainerState,
        micro_batches: list[dict[str, Any]] | None = None,
        **kwargs: Any,

        micro_batches: list[dict[str, Any]] | None = None,
        **kwargs: Any,
    ) -> None:
        """Start timing and count local input tokens and samples."""
        del state, kwargs
        batches = self._micro_batches(micro_batches)
        self._local_step_tokens = sum(self._batch_tokens(batch) for batch in batches)
        self._local_step_samples = sum(self._batch_samples(batch) for batch in batches)
        self._step_start_time = time.perf_counter()

    def on_step_end(
        self,
        state: TrainerState,
        loss: float,
        loss_dict: dict[str, float] | None,

        grad_norm: float,
        **kwargs: Any,
    ) -> None:
        """Reduce and publish metrics for one completed optimizer step."""
        del state, kwargs
        step_time = max(time.perf_counter() - self._step_start_time, 0.0)
        global_step_time = self._reduce(step_time, op="max")
        global_tokens = int(self._reduce(self._local_step_tokens, op="sum"))
        global_samples = int(self._reduce(self._local_step_samples, op="sum"))
        self._consumed_tokens += global_tokens
        self._consumed_samples += global_samples

        train_metrics = {
            "training/total_loss": self._reduce(self._scalar(loss, "total_loss"), op="mean"),
            "training/grad_norm": self._reduce(self._scalar(grad_norm, "grad_norm"), op="mean"),
            "training/lr": self._current_lr(),
        }
        for name, value in sorted((loss_dict or {}).items()):
            metric_name = name if name.startswith("training/") else f"training/{name}"
            train_metrics[metric_name] = self._reduce(self._scalar(value, name), op="mean")

        env_metrics = {
            **train_metrics,
            "performance/step_time": global_step_time,
            "performance/tokens_per_second": global_tokens / global_step_time if global_step_time > 0 else 0.0,
            "data/step_tokens": float(global_tokens),

            "data/step_samples": float(global_samples),
            "data/consumed_samples": float(self._consumed_samples),
            **self._memory_metrics(),
        }
        self.trainer.step_train_metrics = train_metrics
        self.trainer.step_env_metrics = env_metrics

        if callable(tqdm_write) and tqdm_write(message):
            return
        logger.info("%s", message)

    def on_step_end(
        self,
        state: TrainerState,
        loss: float,
        loss_dict: dict[str, float],

        grad_norm: float,
        **kwargs: Any,
    ) -> None:
        """Log all shared environment metrics at the configured cadence."""
        del loss, loss_dict, grad_norm, kwargs
        if self.logging_steps <= 0:
            return
        if getattr(self.trainer, "global_rank", 0) != 0:
            return
        if state.global_step % self.logging_steps != 0:
            return
        if self._last_logged_step == state.global_step:
            return

        metrics = getattr(self.trainer, "step_env_metrics", {})
        if not metrics:

that model construction uses when no explicit device option is given — and
re-exported here so Trainer consumers have one import site.
"""

__all__ = [
    "IS_CUDA_AVAILABLE",
    "IS_NPU_AVAILABLE",
    "get_device_type",
    "get_device_name",

plan §1195-1196). The mesh objects themselves are AutoModels-side in
``hyper_parallel.distributed.mesh``.
"""

__all__ = ["all_gather", "all_reduce"]

import logging
import os
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union

        "sum": dist.ReduceOp.SUM,
        "max": dist.ReduceOp.MAX,
        "min": dist.ReduceOp.MIN,
    }
    reduce_op = reduce_ops.get(op)
    if reduce_op is None:
        raise ValueError("op must be one of: mean, sum, max, min")
    dist.all_reduce(data, op=reduce_op, group=group)
    if op == "mean":  # ReduceOp.AVG is not supported by the NPU backend
        data /= dist.get_world_size(group=group)

    if data.numel() == 1:

    return dist


def _resolve_data_parallel_sizes(
    world_size: int,
    tp_size: int,
    cp_size: int,
    pp_size: int,

            "DP+CP size "
            f"{fsdp_data_parallel_size} is not divisible by FSDP shard size "
            f"{dp_shard_size}"
        )
    return dp_size, fsdp_data_parallel_size // dp_shard_size


def _populate_mesh_ranks(mesh_context: MeshContext, dp_shard_size: int) -> None:
    """Populate local mesh ranks after device-mesh construction."""
    def _local_rank(dim: str) -> int:
        if (
            mesh_context.device_mesh is None

        else 0
    )


def create_distributed_setup_from_config(cfg: Any) -> DistributedSetup:
    """Create DistributedSetup and build the configured mesh domains."""
    accel = cfg.accelerator if cfg is not None and hasattr(cfg, "accelerator") else None
    if accel is None:
        return DistributedSetup(mesh_context=MeshContext())

    fsdp_config = cfg.fsdp_config
    dp_shard_size = max(1, fsdp_config.dp_shard_size)
    edp_shard_size = max(1, fsdp_config.edp_shard_size)
    tp_size = max(1, accel.tp_size)
    cp_size = max(1, accel.cp_size)
    pp_size = max(1, accel.pp_size)
    ep_size = max(1, accel.ep_size)

    world_size = dist.get_world_size() if dist.is_initialized() else 1
    dp_size, dp_replicate_size = _resolve_data_parallel_sizes(
        world_size,
        tp_size,
        cp_size,
        pp_size,

        cp_size,
        pp_size,
        dp_shard_size,
    )
    mesh_context = MeshContext(
        dp_size=dp_size,
        dp_replicate_size=dp_replicate_size,
        dp_shard_size=dp_shard_size,
        edp_shard_size=edp_shard_size,

        ep_size=ep_size,
        sequence_parallel=bool(accel.sequence_parallel),
        loss_parallel=bool(accel.loss_parallel),
    )
    if dist.is_initialized():
        mesh_context, _ = _build_device_mesh_from_accelerator(
            accel,
            dp_shard_size,
            dp_replicate_size,
            world_size,

            world_size,
            edp_shard_size,
        )

    _populate_mesh_ranks(mesh_context, dp_shard_size)

    fsdp_enabled = dist.is_initialized() and (
        dp_shard_size > 1 or dp_replicate_size > 1 or edp_shard_size > 1
    )

Split out of the former ``auto_models/components/utils/helper.py`` in stage 7
(05 §10.4); function names and signatures are unchanged.
"""

__all__ = [
    "create_logger",
    "disable_warning",
    "enable_third_party_logging",
    "setup_logging",

    builtin_logging.basicConfig(level=builtin_logging.ERROR)
    warnings.simplefilter("ignore")
    LoggingMetricsReporter()
    reporter_logger = builtin_logging.getLogger(LoggingMetricsReporter.__name__)
    # PyIceberg exposes no public logger setter; retain its class-level reporter binding.
    setattr(LoggingMetricsReporter, "_logger", reporter_logger)
    reporter_logger.setLevel(builtin_logging.WARNING)
    reporter_logger.propagate = False


if os.getenv("DISABLE_WARNINGS", "0").lower() in ["true", "1"]:
    disable_warning()

Split out of the former ``auto_models/components/utils/helper.py`` in stage 7
(05 §10.4); function names and signatures are unchanged.
"""

__all__ = [
    "empty_cache",
    "print_cpu_memory_info",
    "print_device_mem_info",
]

"""Global loss metrics reduction across data/context parallel ranks."""

from __future__ import annotations

__all__ = ["mean_global_loss"]

from typing import TYPE_CHECKING, Union

import torch

if TYPE_CHECKING:
    from hyper_parallel.distributed.mesh import MeshContext


def _reduce_weighted_loss(cur_loss, cur_token_len, all_reduced_len, loss_name, device_mesh, dp_cp_group):
    """Reduce a token-weighted loss while preserving the local backward graph."""
    if all_reduced_len != 0:
        local_weighted_loss = cur_loss * cur_token_len
        backward_loss = local_weighted_loss / all_reduced_len * device_mesh.dp_size * device_mesh.cp_size
        global_weighted_loss = all_reduce(
            local_weighted_loss.detach().item(),
            op="sum",
            group=dp_cp_group,
        )
        global_mean = cur_loss.new_tensor(global_weighted_loss / all_reduced_len)
        return backward_loss + global_mean - backward_loss.detach()

    if not torch.allclose(cur_loss, torch.zeros_like(cur_loss)):
        raise ValueError(f"The all_reduced_len for {loss_name}_tokens is 0, but the cur_loss is not 0: {cur_loss}")
    return cur_loss


def mean_global_loss(
    losses: Union[dict[str, torch.Tensor], torch.Tensor],