Diff Coverage

    graph-pass surface from the trainer topology without altering runtime
    objects. Accepting ``Any`` keeps this module importable without the
    trainer package.
    """
    accelerator = config.accelerator
    if fsdp_enabled is None:
        fsdp_enabled = (
            config.fsdp_config.dp_shard_size > 1
            or config.fsdp_config.edp_shard_size > 1
        )
    fsdp_degree = config.fsdp_config.dp_shard_size if fsdp_enabled else None
    return PassConfig(
        enable_overlap=enable_overlap,
        fsdp_enabled=fsdp_enabled,
        fsdp_degree=fsdp_degree,
        tp_size=accelerator.tp_size,

import torch

from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
from hyper_parallel.trainer.text_trainer import TextTrainer
from hyper_parallel.trainer.config import TrainerConfig

from .trainer import GraphTrainer


class GraphTextTrainer(TextTrainer):
    """Reuse TextTrainer and delegate graph execution to ``GraphTrainer``."""

    def __init__(
        self,
        config: TrainerConfig,
    ) -> None:
        """Build graph-mode text training on top of the eager TextTrainer stages."""
        super().__init__(config)
        self._graph_train_fn = self._default_train_fn
        self.base_graph_trainer = GraphTrainer(
            model=self.base.model,
            train_fn=self._graph_train_fn,
            trainer_config=config,
            device=self.base.device,

            mesh_context=self.base.mesh,
            manage_optimizer=False,
        )

    def _default_train_fn(
        self,
        model: torch.nn.Module,
        model_inputs: Mapping[str, Any],
        loss_inputs: Mapping[str, Any],

        model_inputs: Mapping[str, Any],
        loss_inputs: Mapping[str, Any],
    ) -> torch.Tensor:
        """Default graph trace function for Transformer-style text training."""
        outputs = model(**dict(model_inputs), use_cache=False)
        labels = loss_inputs.get("labels")
        loss = self.base.loss_fn(model_output=outputs, labels=labels)
        if isinstance(loss, dict):
            return torch.stack(list(loss.values())).sum()
        return loss

    def set_pytree_pre_hook(self, hook: Any) -> "GraphTextTrainer":
        """Register a tracer pre-hook on the underlying graph executor."""
        self.base_graph_trainer.set_pytree_pre_hook(hook)
        return self

    def forward_backward_step(
        self,
        data_iterator: Any,
        num_micro_steps: int,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

        data_iterator: Any,
        num_micro_steps: int,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        """Fetch one text batch and execute graph-mode forward/backward."""
        model_inputs, loss_inputs = self.base.get_batch(data_iterator)
        self.base.current_token_counts = count_loss_token(loss_inputs)
        self.base.step_token_counts = {
            name: token_count * num_micro_steps
            for name, token_count in self.base.current_token_counts.items()
        }
        loss = self.base_graph_trainer.train_step(model_inputs, loss_inputs)
        return loss, _loss_to_metrics(loss)


def _loss_to_metrics(loss: Any) -> dict[str, Any]:
    """Normalize graph loss output to a callback-friendly metrics mapping."""
    if isinstance(loss, dict):
        return {
            str(name): value.detach() if hasattr(value, "detach") else value
            for name, value in loss.items()
        }
    return {"graph_loss": loss.detach() if hasattr(loss, "detach") else loss}

__all__ = ["GraphTextTrainer"]

            f"  Got:    {list(model_state.keys())}"
        )

    state_flat, _ = torch.utils._pytree.tree_flatten({"model": model_state})
    user_inputs_flat, _ = torch.utils._pytree.tree_flatten(
        (input_batch, label_batch)
    )
    flat_inputs = list(state_flat) + list(user_inputs_flat)

    with torch.no_grad():
        outputs = joint_graph.graph_module(*flat_inputs)

                builds and steps the optimizer on the live model.
        """
        self.model = model
        self.train_fn = train_fn
        self.pass_config = self._resolve_pass_config(
            pass_config=pass_config,
            trainer_config=trainer_config,
        )
        self.pass_plan = pass_plan

        )
        self.pass_plan = pass_plan
        self.optimizer_config = optimizer_config or {}
        self._mesh_context = mesh_context
        self._manage_optimizer = manage_optimizer
        self.device = device or (
            torch.device("npu")
            if (hasattr(torch, "npu") and torch.npu.is_available())
            else torch.device("cpu")

            if (hasattr(torch, "npu") and torch.npu.is_available())
            else torch.device("cpu")
        )

        self.pass_config.validate()

        self._joint_graph = None
        self.optimizer = None
        # Optional hook run right before the first compile, for model-specific

        pass_config: Optional[PassConfig],
        trainer_config: Optional[Any],
    ) -> PassConfig:
        """Resolve the graph pass config from explicit or trainer-level input."""
        if pass_config is not None:
            return pass_config
        if trainer_config is not None:
            return build_pass_config_from_trainer_config(trainer_config)
        return PassConfig()

    def compile(self, sample_input: torch.Tensor, sample_label: torch.Tensor) -> None:
        """
        Compile model into parallel graph

        pipeline.run(joint_graph.graph_module, **pass_kwargs)

        self._joint_graph = joint_graph

        if self._manage_optimizer:
            self._init_optimizer()

    def _init_device_mesh(self, mesh_context: Optional[Any] = None):
        """Initialize the FSDP process group.

        return loss

    def optimizer_step(self) -> None:
        """Optimizer update"""
        if not self._manage_optimizer:
            raise RuntimeError(
                "optimizer_step() is disabled when manage_optimizer=False. "
                "Let the outer trainer runtime own optimizer stepping."
            )
        if self.optimizer is None:

    if isinstance(compile_config, dict):
        compile_config = CompileConfig(enabled=True, **compile_config)
    if compile_config is not None and not isinstance(compile_config, CompileConfig):
        raise TypeError("compile_config must be a CompileConfig, mapping, or None")
    graph_mode_requested = bool(
        compile_config is not None and compile_config.selects_graph_trainer()
    )
    compile_for_execution = bool(
        not validate_placement

        and compile_config is not None
        and compile_config.enabled
        and not graph_mode_requested
    )
    if graph_mode_requested:
        logger.info(
            "compile.graphtrainer_enabled=True selects graph-mode trainer; "
            "skipping decoder-layer compile",
        )
    if validate_placement and compile_config is not None and compile_config.enabled:

    peft: Optional[Any] = None

    def __post_init__(self) -> None:
        """Validate combinations that span multiple config sections."""
        if (
            self.compile.enabled
            and not self.compile.selects_graph_trainer()
            and self.accelerator.pp_size > 1
        ):