Diff Coverage

        """Run all registered callbacks at the start of a training step."""
        for callback in self._callbacks:
            callback.on_step_begin(self.state, **kwargs)

    def on_micro_step_begin(self, micro_batch: Dict[str, Any], **kwargs: Any) -> None:
        """Run all registered callbacks before one forward-backward micro step.

        Args:
            micro_batch: Prepared inputs for the current micro step.

        Args:
            micro_batch: Prepared inputs for the current micro step.
            **kwargs: Additional callback context.
        """
        for callback in self._callbacks:
            callback.on_micro_step_begin(self.state, micro_batch, **kwargs)

    def on_step_end(
        self,
        loss: Optional[float] = None,

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

    @classmethod
    def _batch_tokens(cls, batch: Mapping[str, Any]) -> Any:
        """Count text tokens without synchronizing a device scalar to the host."""
        token_count = batch.get("token_count")
        if token_count is not None:
            return token_count

        labels = batch.get("labels")
        if labels is not None and callable(getattr(labels, "sum", None)):
            return (labels != IGNORE_INDEX).sum()

        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 attention_mask.sum()

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

        return float(reduced)

    def _accumulate_batches(self, value: Any) -> None:
        """Accumulate batch metrics without retaining input tensor references."""
        for batch in self._micro_batches(value):
            token_count = self._batch_tokens(batch)
            if callable(getattr(token_count, "detach", None)):
                token_count = token_count.detach()
            if self._local_step_tokens is None:
                if callable(getattr(token_count, "clone", None)):
                    token_count = token_count.clone()
                self._local_step_tokens = token_count
            else:
                self._local_step_tokens = self._local_step_tokens + token_count
            self._local_step_samples += self._batch_samples(batch)

    def _global_samples(self) -> int:
        """Reduce samples across DP+CP while removing CP replicas."""
        cp_size = int(getattr(self.trainer.mesh, "cp_size", 1))
        if cp_size < 1:
            raise ValueError(f"mesh.cp_size must be positive, but got {cp_size}")
        reduced_samples = self._reduce(self._local_step_samples, op="sum")
        global_samples = reduced_samples / cp_size
        if not global_samples.is_integer():
            raise ValueError(
                "Reduced sample count must be divisible by cp_size, "
                f"but got reduced_samples={reduced_samples} and cp_size={cp_size}"
            )
        return int(global_samples)

    def _current_lr(self) -> float:
        """Return the maximum learning rate across scheduler or optimizer groups."""
        schedulers = self.trainer.lr_scheduler

            micro_batches: Batches available before the step starts.
            **kwargs: Unused callback context.
        """
        del state, kwargs
        self._local_step_tokens = None
        self._local_step_samples = 0
        self._step_start_time = time.perf_counter()
        self._accumulate_batches(micro_batches)

    def on_micro_step_begin(
        self,
        state: TrainerState,

            state: Current training progress.
            micro_batch: Prepared inputs and lightweight metric metadata.
            **kwargs: Unused callback context.
        """
        del state, kwargs
        self._accumulate_batches(micro_batch)

    def on_step_end(
        self,
        state: TrainerState,

        """
        del state, kwargs
        step_time = max(time.perf_counter() - self._step_start_time, 0.0)
        global_step_time = self._reduce(step_time, op="max")
        local_step_tokens = 0 if self._local_step_tokens is None else self._local_step_tokens
        global_tokens = int(self._reduce(local_step_tokens, op="sum"))
        global_samples = self._global_samples()
        self._local_step_tokens = None
        self._local_step_samples = 0
        self._consumed_tokens += global_tokens
        self._consumed_samples += global_samples

        train_metrics = {

        initial: int,
        stream: Any,
    ) -> None:
        """Initialize non-interactive progress state without emitting a line."""
        self._tqdm_type = tqdm_type
        self._initial = initial
        self._start_time = time.monotonic()
        self._postfix: dict[str, str] = {}
        self._pending_message: str | None = None
        self.total = total
        self.n = initial
        self.fp = stream

    def set_postfix(self, postfix: dict[str, str], refresh: bool = True) -> None:
        """Store the final metrics for the next completed-step snapshot.

        Args:
            postfix: Formatted metric names and values.
            refresh: Accepted for compatibility with the tqdm API.
        """
        del refresh
        self._postfix = dict(postfix)

    def update(self, amount: int) -> None:
        """Advance progress and write exactly one newline-terminated snapshot.

        Args:
            amount: Number of newly completed steps.
        """
        self.n += amount
        elapsed = max(time.monotonic() - self._start_time, 0.0)
        completed = self.n - self._initial
        rate = completed / elapsed if elapsed > 0 else None
        postfix = ", ".join(f"{name}={value}" for name, value in self._postfix.items())
        progress_line = self._tqdm_type.format_meter(
            n=self.n,
            total=self.total,
            elapsed=elapsed,
            prefix="Training",

            unit="step",
            rate=rate,
            postfix=postfix or None,
        )
        if self._pending_message:
            progress_line = f"{progress_line} | {self._pending_message}"
        self.fp.write(f"{progress_line}\n")
        self.fp.flush()
        self._pending_message = None

    def write(self, message: str, file: Any = None) -> None:
        """Attach one structured message to the next progress snapshot.

        Args:
            message: Structured log message for the current step.
            file: Accepted for compatibility with the tqdm API.
        """
        del file
        self._pending_message = message

    def close(self) -> None:
        """Close without writing a duplicate final progress line."""

            if not self._missing_dependency_warned:
                logger.warning("TqdmCallback: 'tqdm' is not installed; progress bar disabled")
                self._missing_dependency_warned = True
            return None
        stream = sys.stderr
        isatty = getattr(stream, "isatty", None)
        if not callable(isatty) or not isatty():
            return _LogProgressBar(
                tqdm_type=tqdm,
                total=total,
                initial=initial,
                stream=stream,

            loss_dict: Named loss values for the completed step.
            grad_norm: Gradient norm measured before the optimizer update.
            **kwargs: Additional callback arguments.
        """
        del loss, loss_dict, grad_norm, kwargs
        if self._progress_bar is None or state.global_step <= self._last_updated_step:
            return

        postfix = self._postfix()

            return

        postfix = self._postfix()
        if postfix:
            self._progress_bar.set_postfix(postfix, refresh=False)
        self._progress_bar.update(state.global_step - self._last_updated_step)
        self._last_updated_step = state.global_step

    def write(self, message: str) -> bool:

    def on_step_begin(self) -> None:
        """Dispatch the step-begin lifecycle hook."""
        self.base.on_step_begin()

    def on_micro_step_begin(self, micro_batch: dict[str, Any]) -> None:
        """Dispatch the micro-step-begin lifecycle hook.

        Args:
            micro_batch: Prepared inputs and lightweight metric metadata.

        Args:
            micro_batch: Prepared inputs and lightweight metric metadata.
        """
        self.base.on_micro_step_begin(micro_batch)

    def on_step_end(
            self,
            loss: Any = None,

        self.base.step_token_counts = {
            name: token_count * num_micro_steps
            for name, token_count in self.base.current_token_counts.items()
        }
        metric_inputs = {
            **model_inputs,
            **loss_inputs,
            "token_count": self.base.current_token_counts["foundation_tokens"],
        }
        self.on_micro_step_begin(metric_inputs)
        del metric_inputs
        loss, loss_dict = self.base.forward_backward_step(model_inputs, loss_inputs)

        return loss, loss_dict