Diff Coverage

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

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

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

            "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 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",
]

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],

            op="sum",
            group=dp_cp_group,
        )

        cur_loss = _reduce_weighted_loss(
            cur_loss,
            cur_token_len,
            all_reduced_len,
            loss_name,

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

__all__ = [
    "CACHE_DIR",
    "ProfilerWithMem",
    "create_profiler",
    "get_cache_dir",

import datetime
import logging
import os
from typing import Any, NamedTuple, Optional

import torch

from hyper_parallel.models.build_options import (

        """Advance the wrapped profiler by one step."""
        return self._p.step(*a, **kw)


class _ProfilerBackend(NamedTuple):
    """Backend-specific profiler objects."""

    module: Any
    activities: Any
    trace_handler: Any
    experimental_config: Any


def _create_profiler_backend(trace_dir: str) -> _ProfilerBackend:
    """Create backend-specific profiler objects."""
    if IS_NPU_AVAILABLE:
        profiler_module = torch_npu.profiler
        activities = [profiler_module.ProfilerActivity.CPU, profiler_module.ProfilerActivity.NPU]
        trace_handler = torch_npu.profiler.tensorboard_trace_handler(
            CACHE_DIR if trace_dir.startswith("hdfs://") else trace_dir
        )
        experimental_config = torch_npu.profiler._ExperimentalConfig(  # pylint: disable=protected-access
            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
            data_simplification=False,
        )

            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
            data_simplification=False,
        )
    else:
        profiler_module = torch.profiler
        activities = [profiler_module.ProfilerActivity.CPU, profiler_module.ProfilerActivity.CUDA]
        trace_handler = None
        experimental_config = None

    return _ProfilerBackend(profiler_module, activities, trace_handler, experimental_config)


def _create_profiler_schedule(profiler_module: Any, start_step: int, end_step: int) -> Any:
    """Create and log the profiler schedule."""
    warmup = 0 if start_step == 1 else 1
    wait = start_step - warmup - 1
    active = end_step - start_step
    logger.info(f"build profiler schedule - wait: {wait}, warmup: {warmup}, active: {active}.")  # pylint: disable=logging-fstring-interpolation

    return profiler_module.schedule(
        wait=wait,
        warmup=warmup,
        active=active,
        repeat=1,

        Args:
            p: The profiler instance that produced the trace.
        """
        time = int(datetime.datetime.now(datetime.timezone.utc).timestamp())

        trace_file_extention = "pt.trace.json.gz"
        gpu_memory_file_extension = "pkl"

            trace_file = os.path.join(trace_dir, f"veomni_rank{global_rank}_{time}.{trace_file_extention}")
            gpu_memory_file = os.path.join(trace_dir, f"veomni_rank{global_rank}_{time}.{gpu_memory_file_extension}")

        if IS_NPU_AVAILABLE:
            profiler_backend.trace_handler(p)
            trace_file = p.prof_if.prof_path
        elif IS_CUDA_AVAILABLE:
            p.export_chrome_trace(trace_file)
        logger.info(f"Profiling result saved at {trace_file}.")  # pylint: disable=logging-fstring-interpolation

                raise ValueError("hdfs_io.copy is required for an HDFS profiling trace directory")
            copy(trace_file, trace_dir)
            logger.info(f"Profiling result uploaded to {trace_dir}.")  # pylint: disable=logging-fstring-interpolation

    profiler_backend = _create_profiler_backend(trace_dir)
    schedule = _create_profiler_schedule(profiler_backend.module, start_step, end_step)
    base_profiler = profiler_backend.module.profile(
        activities=profiler_backend.activities,
        schedule=schedule,
        on_trace_ready=handler_fn,
        record_shapes=record_shapes,

# See the License for the specific language governing permissions and
# limitations under the License.
"""Text Trainer assembled from the shared BaseTrainer stages."""

__all__ = ["TextTrainer"]

from collections import defaultdict
from typing import Any, Dict

# See the License for the specific language governing permissions and
# limitations under the License.
"""VLM Trainer assembled from the shared BaseTrainer stages."""

__all__ = ["VLMTrainer"]

from collections import defaultdict
from typing import Any, Dict

            loss_dict=loss_dict,
            grad_norm=grad_norm,
        )

    def _forward_backward_micro_batches(
            self,
            training_batches: list[Any],
            num_micro_steps: int,
    ) -> tuple[float, Dict[str, float]]:

            total_loss += loss.item()
            for loss_name, loss_value in loss_dict.items():
                total_loss_dict[loss_name] += loss_value.item()

        return total_loss, total_loss_dict

    def train_step(self, data_iterator: Any) -> Dict[str, float]:
        """Execute one VLM training step."""
        config = self.base.config
        first_training_batch = self.base.get_batch(data_iterator)
        num_micro_steps = self.base.num_micro_batches
        training_batches = [first_training_batch]
        for _ in range(1, num_micro_steps):
            training_batches.append(self.base.get_batch(data_iterator))

        self.on_step_begin(
            micro_batches=[model_inputs for model_inputs, _ in training_batches]
        )
        synchronize()

        total_loss, total_loss_dict = self._forward_backward_micro_batches(
            training_batches,
            num_micro_steps,
        )

            self.base.model,
            config.training.max_grad_norm,
        )

        self.base.step_optimizers_and_schedulers()

        grad_norm_value = float(grad_norm)
        self.on_step_end(
            loss=total_loss,