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:

            cb.on_step_begin(self.state)

    def on_step_end(self, loss=None, grad_norm=None):
        """Dispatch on_step_end to callbacks with monitor output first."""
        self.training_state_monitor_callback.on_step_end(
            self.state, loss=loss, grad_norm=grad_norm,
        )
        for cb in self._all_callbacks():
            if cb is self.training_state_monitor_callback:
                continue
            cb.on_step_end(self.state, loss=loss, grad_norm=grad_norm)

    def on_substep_end(self):
        """Dispatch on_substep_end (after each micro-batch forward/backward)."""

    def on_substep_end(self):
        """Dispatch on_substep_end (after each micro-batch forward/backward)."""
        self.moe_monitor_callback.on_substep_end(self.state)
        self.training_state_monitor_callback.on_substep_end(self.state)
        for cb in self.user_callbacks:
            cb.on_substep_end(self.state)

    def on_pre_optimizer_step(self, grad_norm=None):

    def on_pre_optimizer_step(self, grad_norm=None):
        """Dispatch on_pre_optimizer_step (after grad clip, before optimizer.step)."""
        # Monitor first so scalar stats are 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,

        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
        )

        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."""
        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 _should_record_param(self, name: str) -> bool:
        patterns = self._target_patterns
        if not patterns:
            matched = True
        else:
            matched = any(pattern.search(name) for pattern in patterns)
        return not matched if self._invert else matched

    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

    def _flush_step_log(self, step: int) -> None:

            self._pending_log_metrics[tag] = value

    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}={float(value):.8f}")
        for tag in sorted(self._pending_log_metrics):
            if tag in raw_tags:
                continue
            value = self._pending_log_metrics[tag]
            parts.append(f"{tag}={float(value):.8f}")
        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 = {}

    @staticmethod
    def _compute_grad_stats(grad) -> dict:
        """Compute scalar stats for a hook-captured local gradient."""

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

            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.enabled:
            return
        self._rank = platform.get_rank()
        if self._uses_tensorboard():
            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",
            self.dump_path, self.step_interval,
        )

        )
        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.enabled 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

            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.enabled:
            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.enabled:
            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._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.

        super().__init__(trainer)
        train_cfg = getattr(trainer.args, "train", None)
        debug_cfg = getattr(train_cfg, "debug", None) if train_cfg is not None else None
        if debug_cfg is None:
            debug_cfg = getattr(trainer.args, 'debug', None)
        self.enabled = (
            getattr(debug_cfg, 'check_nan_inf', False) if debug_cfg else False
        )