Diff Coverage

        self.tensorboard_callback = TensorBoardCallback(self)
        self.progress_callback = ProgressCallback(self)
        self.moe_monitor_callback = MoEMonitorCallback(self)
        # Health + operability (no-ops unless enabled in cfg.train.debug / .memory_monitor).
        self.training_state_monitor_callback = TrainingStateMonitorCallback(self)
        self.gradient_health_callback = GradientHealthCallback(self)
        self.memory_monitor_callback = MemoryMonitorCallback(self)
        self.gc_callback = GCCallback(self)
        # ``user_callbacks`` lets external code append extra Callback instances

        """Dispatch on_train_begin to all callbacks."""
        # Memory monitor first so it captures the truly-initial peak.
        self.memory_monitor_callback.on_train_begin(self.state)
        self.moe_monitor_callback.on_train_begin(self.state)
        self.training_state_monitor_callback.on_train_begin(self.state)
        self.profiler_callback.on_train_begin(self.state)
        self.wandb_callback.on_train_begin(self.state)
        self.tensorboard_callback.on_train_begin(self.state)
        # Checkpoint runs after log writers are armed and before progress so

        """Dispatch on_train_end to all callbacks."""
        self.checkpoint_callback.on_train_end(self.state)
        self.hf_export_callback.on_train_end(self.state)
        self.progress_callback.on_train_end(self.state)
        self.training_state_monitor_callback.on_train_end(self.state)
        self.tensorboard_callback.on_train_end(self.state)
        self.wandb_callback.on_train_end(self.state)
        self.profiler_callback.on_train_end(self.state)
        for cb in self.user_callbacks:

    def on_pre_optimizer_step(self, grad_norm=None):
        """Dispatch on_pre_optimizer_step (after grad clip, before optimizer.step)."""
        # Monitor first so an abnormal step is captured before health checks abort.
        self.training_state_monitor_callback.on_pre_optimizer_step(
            self.state, grad_norm=grad_norm,
        )
        self.gradient_health_callback.on_pre_optimizer_step(
            self.state, grad_norm=grad_norm,

import torch
try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    SummaryWriter = None

from hyper_parallel import get_platform
from hyper_parallel.core.distributed_checkpoint import load as dcp_load, save as dcp_save
from hyper_parallel.core.distributed_checkpoint.offline_transform import (

        message = " | ".join(f"{k}={v}" for k, v in metrics.items())
        monitor_cb = getattr(self.trainer, 'training_state_monitor_callback', None)
        monitor_active = monitor_cb is not None and getattr(monitor_cb, 'active', False)
        if monitor_active:
            logger.info(message)
        else:
            logger.info_rank0(message)

        record = {

        super().__init__(trainer)
        train_cfg = getattr(trainer.args, "train", None)
        self.cfg = getattr(train_cfg, "monitor", None)
        if self.cfg is None:
            self.cfg = getattr(trainer.args, "monitor", None)

        self.enabled = bool(getattr(self.cfg, "monitor_on", False)) if self.cfg else False
        self.dump_path = getattr(self.cfg, "dump_path", "./dump") if self.cfg else "./dump"
        self.step_interval = int(getattr(self.cfg, "step_interval", 1) if self.cfg else 1)

        self.enabled = bool(getattr(self.cfg, "monitor_on", False)) if self.cfg else False
        self.dump_path = getattr(self.cfg, "dump_path", "./dump") if self.cfg else "./dump"
        self.step_interval = int(getattr(self.cfg, "step_interval", 1) if self.cfg else 1)
        if self.step_interval < 1:
            raise ValueError("train.monitor.step_interval must be >= 1.")

        self.local_loss_format = self._parse_formats("local_loss_format")
        self.device_local_loss_format = self._parse_formats("device_local_loss_format")
        self.local_norm_format = self._parse_formats("local_norm_format")

        self.device_local_norm_format = self._parse_formats("device_local_norm_format")

        raw_patterns = getattr(self.cfg, "target", None) if self.cfg else None
        if isinstance(raw_patterns, str):
            raw_patterns = [raw_patterns]
        self._target_patterns = (
            [re.compile(pattern) for pattern in raw_patterns]
            if raw_patterns else None
        )

        """Parse configured monitor output formats for a metric field."""
        value = getattr(self.cfg, field_name, None) if self.cfg else None
        if value is None:
            return ()
        if isinstance(value, str):
            formats = (value,)
        else:
            formats = tuple(value)
        unknown = sorted(set(formats) - self._SUPPORTED_FORMATS)
        if unknown:
            raise ValueError(
                f"train.monitor.{field_name} only supports "
                f"{sorted(self._SUPPORTED_FORMATS)}, got {unknown}."
            )
        return formats

    @staticmethod
    def _to_scalar(value) -> float:
        """Convert tensor-like values to a Python float."""

    @staticmethod
    def _to_scalar(value) -> float:
        """Convert tensor-like values to a Python float."""
        if value is None:
            return 0.0
        if hasattr(value, "to_local"):
            value = value.to_local()
        if hasattr(value, "detach"):
            value = value.detach()
        if hasattr(value, "float"):
            value = value.float()

        if hasattr(value, "float"):
            value = value.float()
        if hasattr(value, "item"):
            return float(value.item())
        return float(value)

    @staticmethod
    def _sanitize_tag(value: str) -> str:
        value = value.replace(".", "/").replace(" ", "_")

    def _write(self, formats: tuple[str, ...], tag: str, value: float,
               step: int, *, global_metric: bool = False) -> None:
        if not formats:
            return
        if "tensorboard" in formats:
            writer = self._global_writer if global_metric else self._rank_writer
            if writer is not None:
                writer.add_scalar(tag, value, step)
        if "log" in formats:
            self._pending_log_metrics[tag] = value

    @staticmethod
    def _format_monitor_value(value: float) -> str:
        return f"{float(value):.8f}"

    def _flush_step_log(self, step: int) -> None:
        """Print one compact rank-local monitor line for console output."""
        if not self._pending_log_metrics:
            return
        field_map = (
            ("loss", "loss/local_loss"),
            ("accum_loss", "loss/device_accum_local_loss"),
            ("grad_norm", "grad/device_local_norm"),
            ("grad_nan_count", "grad/device_nan_count"),

            ("grad_norm", "grad/device_local_norm"),
            ("grad_nan_count", "grad/device_nan_count"),
            ("grad_inf_count", "grad/device_inf_count"),
        )
        raw_tags = {tag for _, tag in field_map}
        parts = []
        for display_name, tag in field_map:
            if tag in self._pending_log_metrics:
                value = self._pending_log_metrics[tag]
                parts.append(f"{display_name}={self._format_monitor_value(value)}")
        for tag in sorted(self._pending_log_metrics):
            if tag in raw_tags:
                continue
            value = self._pending_log_metrics[tag]
            parts.append(f"{tag}={self._format_monitor_value(value)}")
        timestamp = time.strftime("%H:%M:%S")
        print(
            f"[{timestamp}][rank{self._rank}][INFO] local_state: "
            f"step={step} | " + " | ".join(parts),
            flush=True,
        )
        self._pending_log_metrics = {}

    def _abnormal_record_path(self, *, global_record: bool = False) -> str:
        configured = getattr(self.cfg, "abnormal_record_path", None) if self.cfg else None
        base_path = configured or os.path.join(self.dump_path, "abnormal_training_state.json")

    def _abnormal_record_path(self, *, global_record: bool = False) -> str:
        configured = getattr(self.cfg, "abnormal_record_path", None) if self.cfg else None
        base_path = configured or os.path.join(self.dump_path, "abnormal_training_state.json")
        if global_record:
            return base_path
        root, ext = os.path.splitext(base_path)
        suffix = ext or ".json"
        return f"{root}_rank_{self._rank}{suffix}"

    def _record_abnormal(self, record: dict, *, global_record: bool = False) -> None:
        """Append an abnormal training-state record to the JSON log."""
        if global_record and self._rank != 0:
            return
        path = self._abnormal_record_path(global_record=global_record)
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        records = []
        if os.path.isfile(path):
            try:
                with open(path, encoding="utf-8") as file:
                    loaded = json.load(file)
                if isinstance(loaded, list):
                    records = loaded
            except (OSError, json.JSONDecodeError):
                records = []
        records.append(record)
        with open(path, "w", encoding="utf-8") as file:
            json.dump(records, file, indent=2)

    def _remove_grad_hooks(self) -> None:
        """Remove parameter backward hooks registered by the monitor."""
        for handle in self._grad_hook_handles:
            if hasattr(handle, "remove"):
                handle.remove()
        self._grad_hook_handles = []

    def _register_grad_hooks(self) -> None:
        """Register parameter hooks that capture pre-communication local grads."""

                self.local_norm_format,
                self.device_local_norm_format,
                self.check_for_nan_in_loss_and_grad,
                self.embedding_local_norm_threshold is not None,))):
            return
        self._remove_grad_hooks()
        self._clear_hook_grad_stats()
        for name, param in self.trainer.model.named_parameters():
            if not self._should_record_param(name) or not getattr(param, "requires_grad", False):

            if not self._should_record_param(name) or not getattr(param, "requires_grad", False):
                continue
            register_hook = getattr(param, "register_hook", None)
            if not callable(register_hook):
                logger.warning(
                    "TrainingStateMonitor: parameter %s does not support register_hook, skip local grad monitor.",
                    name,
                )
                continue

            def hook_fn(grad, param_name=name):
                self._accumulate_hook_grad_stats(param_name, grad)
                return grad

            self._grad_hook_handles.append(register_hook(hook_fn))

    def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
        if not (self.active or self.health_active):
            return
        self._rank = platform.get_rank()
        if self._uses_tensorboard():
            if SummaryWriter is None:
                raise RuntimeError(
                    "train.monitor requested tensorboard output, but "
                    "torch.utils.tensorboard.SummaryWriter is unavailable. "
                    "Install tensorboard or remove 'tensorboard' from monitor formats."
                )

                    "torch.utils.tensorboard.SummaryWriter is unavailable. "
                    "Install tensorboard or remove 'tensorboard' from monitor formats."
                )

            tb_root = os.path.join(self.dump_path, "tensorboard")
            self._rank_writer = SummaryWriter(
                os.path.join(tb_root, f"rank_{self._rank}")
            )
            if self._rank == 0:
                self._global_writer = SummaryWriter(os.path.join(tb_root, "global"))
        logger.info_rank0(
            "TrainingStateMonitor enabled: dump_path=%s step_interval=%d health=%s",
            self.dump_path, self.step_interval, self.health_active,
        )

        )
        self._register_grad_hooks()

    def on_train_end(self, state: "TrainerState", **kwargs) -> None:
        self._remove_grad_hooks()
        self._clear_hook_grad_stats()
        for writer in (self._rank_writer, self._global_writer):
            if writer is not None:
                writer.close()
        self._rank_writer = None
        self._global_writer = None

    def on_substep_end(self, state: "TrainerState", **kwargs) -> None:
        if not (self.active or self.health_active) or state.global_step % self.step_interval != 0:
            return
        substep_info = getattr(state, "substep_info", {})
        raw_loss = kwargs.get("raw_loss", substep_info.get("raw_loss"))
        if raw_loss is None:
            return
        micro_step = kwargs.get("micro_step", substep_info.get("micro_step", 0))
        if int(micro_step) == 0:
            self._step_loss_sum = 0.0
            self._step_loss_count = 0

        self.last_local_loss = self._step_loss_sum / self._step_loss_count
        self.last_monitor_step = state.global_step

        self._check_nonfinite(state, "local_loss", loss_value)
        tag = "loss/local_loss"
        self._write(
            self.local_loss_format,
            tag,
            loss_value,
            state.global_step,

            if not self._should_record_param(name):
                continue
            stats = self._hook_grad_stats.get(name)
            if stats is None:
                continue
            sum_sq = float(stats["sum_sq"])
            norm = math.sqrt(sum_sq)
            abs_max = float(stats["abs_max"])
            nan_count = int(stats["nan_count"])

                )

    def on_pre_optimizer_step(self, state: "TrainerState", **kwargs) -> None:
        if not (self.active or self.health_active):
            return
        if state.global_step % self.step_interval != 0:
            self._clear_hook_grad_stats()
            return
        try:
            self._record_hook_grad_stats(state)
        finally:
            self._clear_hook_grad_stats()

            self._clear_hook_grad_stats()

    def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
                    grad_norm: Optional[float] = None, **kwargs) -> None:
        if not (self.active or self.health_active):
            return
        if state.global_step % self.step_interval == 0:
            if self._step_loss_count > 0:
                device_loss = self._step_loss_sum / self._step_loss_count
                self._check_nonfinite(state, "device_accum_local_loss", device_loss)
                self._write(
                    self.device_local_loss_format,
                    "loss/device_accum_local_loss",
                    device_loss,
                    state.global_step,

                    "loss/device_accum_local_loss",
                    device_loss,
                    state.global_step,
                )
            if self._rank == 0 and loss is not None:
                self._write(
                    ("tensorboard",) if self._uses_tensorboard() else (),
                    "loss/global_loss",
                    float(loss),
                    state.global_step,

                    float(loss),
                    state.global_step,
                    global_metric=True,
                )
            if self._rank == 0 and grad_norm is not None:
                self._write(
                    ("tensorboard",) if self._uses_tensorboard() else (),
                    "grad/global_grad_norm",
                    float(grad_norm),
                    state.global_step,

                    float(grad_norm),
                    state.global_step,
                    global_metric=True,
                )
            self._flush_step_log(state.global_step)

        self._step_loss_sum = 0.0
        self._step_loss_count = 0


class GradientHealthCallback(Callback):
    """Detect NaN / Inf grad_norm and raise / warn.

            int(getattr(monitor_cfg, "global_norm_spike_count_threshold", 1))
            if monitor_cfg is not None else 1
        )
        if self.global_norm_spike_count_threshold < 1:
            raise ValueError("train.monitor.global_norm_spike_count_threshold must be >= 1.")
        self._global_norm_spike_count = 0
        self.monitor_cfg = monitor_cfg

    def _abnormal_record_path(self, *, global_record: bool = False) -> str:

        )
        base_path = configured or os.path.join(dump_path, "abnormal_training_state.json")
        if global_record:
            return base_path
        root, ext = os.path.splitext(base_path)
        suffix = ext or ".json"
        return f"{root}_rank_{platform.get_rank()}{suffix}"

    def _record_abnormal(self, record: dict, *, global_record: bool = False) -> None:
        """Append an abnormal gradient-health record to the JSON log."""
        if global_record and platform.get_rank() != 0:

                with open(path, encoding="utf-8") as file:
                    loaded = json.load(file)
                if isinstance(loaded, list):
                    records = loaded
            except (OSError, json.JSONDecodeError):
                records = []
        records.append(record)
        with open(path, "w", encoding="utf-8") as file:
            json.dump(records, file, indent=2)

            self._check_global_norm_spike(state, grad_norm)

    def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
                    grad_norm: Optional[float] = None, **kwargs) -> None:
        if not self.nan_inf_check_enabled or loss is None:
            return
        if math.isnan(loss) or math.isinf(loss):
            record = {
                "step": state.global_step,
                "rank": platform.get_rank(),
                "metric": "loss",
                "value": loss,

                "rank": platform.get_rank(),
                "metric": "loss",
                "value": loss,
            }
            self._record_abnormal(record)
            logger.error(
                "GradientHealthCallback: loss=%s at step %d (NaN/Inf).",
                loss, state.global_step,
            )
            raise RuntimeError(
                f"Non-finite loss={loss} at step {state.global_step}. "
                "Disable train.monitor.check_for_nan_in_loss_and_grad or "
                "train.debug.check_nan_inf to skip this guard."
            )

    def _check_global_norm_spike(self, state: "TrainerState", grad_norm: float) -> None:
        """Check whether global grad norm exceeds the configured spike threshold."""
        if not self.global_norm_spike_enabled:
            return
        threshold = self.global_norm_spike_threshold
        if grad_norm >= threshold:
            self._global_norm_spike_count += 1
            record = {