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.

        self.model_bwd_context = nullcontext()

    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,
        )