Diff Coverage

    state = _AUX_LOSS_STATE.get()
    if state is None or not state.active or not loss.requires_grad:
        return
    if loss.numel() != 1:
        raise ValueError("Auxiliary gradient scaling requires a scalar task loss")

    def save_scale(gradient: torch.Tensor) -> None:
        """Capture the actual multiplier after token and parallel weighting.

        return loss, loss_dict

    def _init_callbacks(self):
        """Initialize callbacks."""
        dp_cp_mesh = self.mesh.dp_cp_mesh
        self.moe_monitor = MoEMonitorCallback(
            self.model,
            dp_group=dp_cp_mesh.get_group() if dp_cp_mesh is not None else None,
            tp_group=self.device_mesh["tp"].get_group() if self.mesh.sequence_parallel else None,
        )

        Returns:
            Backward loss and named globally aggregated loss values.
        """
        local_loss = self.loss_fn(model_output=outputs, labels=labels)
        foundation_loss = local_loss.get("foundation_loss") if isinstance(local_loss, dict) else local_loss
        if foundation_loss is not None:
            bind_aux_loss_scale(foundation_loss)
        loss_dict: Dict[str, torch.Tensor] = mean_global_loss(
            local_loss,
            self.current_token_counts,
            self.step_token_counts,