Diff Coverage

        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,

            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:

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

    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)