Lightning-AI/pytorch-lightning · error · NotImplementedError

Gradient clipping is not implemented for optimizers handling

Error message

Gradient clipping is not implemented for optimizers handling the unscaling.

What it means

unscale_gradients calls scaler.unscale_(optimizer), but some optimizers (e.g. OptimizerWithScalerWrapper variants / fused optimizers that handle unscaling internally in step) cannot be unscaled externally. Gradient clipping is not implemented for such optimizers under scaler-based AMP.

Source

Thrown at src/lightning/fabric/plugins/precision/amp.py:111

        return step_output

    @override
    def state_dict(self) -> dict[str, Any]:
        if self.scaler is not None:
            return self.scaler.state_dict()
        return {}

    @override
    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        if self.scaler is not None:
            self.scaler.load_state_dict(state_dict)

    @override
    def unscale_gradients(self, optimizer: Optimizer) -> None:
        scaler = self.scaler
        if scaler is not None:
            if _optimizer_handles_unscaling(optimizer):
                raise NotImplementedError("Gradient clipping is not implemented for optimizers handling the unscaling.")
            scaler.unscale_(optimizer)


def _optimizer_handles_unscaling(optimizer: Any) -> bool:
    """Determines whether a PyTorch optimizer handles unscaling gradients in the step method rather than through the
    :class:`torch.cuda.amp.GradScaler`.

    Since, the current implementation of this function checks a PyTorch internal variable on the optimizer, the return
    value will only be reliable for built-in PyTorch optimizers.

    """
    return getattr(optimizer, "_step_supports_amp_scaling", False)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Switch to 'bf16-mixed' so no scaler/unscale path is used
2. Use an optimizer that lets GradScaler.unscale_ handle unscaling (plain torch optimizers)
3. Disable gradient clipping or implement clipping inside the optimizer's step

Example fix

# before
fabric = Fabric(precision="16-mixed")
fabric.clip_gradients(model, optimizer, clip_val=1.0)  # optimizer self-uncales

# after
fabric = Fabric(precision="bf16-mixed")
fabric.clip_gradients(model, optimizer, clip_val=1.0)

from lightning.fabric.plugins.precision.amp import _optimizer_handles_unscaling

def clipping_supported(plugin, optimizer) -> bool:
    return plugin.scaler is None or not _optimizer_handles_unscaling(optimizer)

def can_clip_gradients(plugin, optimizer) -> bool:
    from lightning.fabric.plugins.precision.amp import _optimizer_handles_unscaling
    return plugin.scaler is None or not _optimizer_handles_unscaling(optimizer)

When it happens

Trigger: MixedPrecision(precision='16-mixed') with gradient clipping enabled (clip_gradients) while the optimizer's step handles unscaling itself (detected by _optimizer_handles_unscaling, e.g. optimizers with a built-in scaler).

Common situations: Using fused/custom optimizers that integrate the GradScaler (some fairseq/apex-style wrappers) together with Lightning's gradient clipping on fp16-mixed.

Related errors

• The current optimizer, {type(optimizer).__qualname__}, does
• Passed `{type(self).__name__}(precision={precision!r})`. Pre
• `precision='bf16-mixed'` does not use a scaler, found {scale
• AMP and the LBFGS optimizer are not compatible.
• You have set `Trainer(gradient_clip_val={self.trainer.gradie

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/d99e00bb4d758407. Report an issue: GitHub.