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

FSDPPrecision.unscale_gradients calls scaler.unscale_(optimizer) when a GradScaler is active, but some optimizers (e.g. those with built-in gradient scoping like bnb optimizers) handle unscaling themselves (detected via _optimizer_handles_unscaling). Combining an external scaler with such an optimizer is unsupported, so a NotImplementedError is raised.

Source

Thrown at src/lightning/fabric/plugins/precision/fsdp.py:156

    def optimizer_step(
        self,
        optimizer: Optimizable,
        **kwargs: Any,
    ) -> Any:
        if self.scaler is None:
            # skip scaler logic, as bfloat16 does not require scaler
            return super().optimizer_step(optimizer, **kwargs)
        # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
        step_output = self.scaler.step(optimizer, **kwargs)  # type: ignore[arg-type]
        self.scaler.update()
        return step_output

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

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

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Switch precision to 'bf16-mixed' (or 32-true) so no scaler is created and unscale_gradients becomes a no-op
2. Use an optimizer that does not implement its own unscaling (plain torch.optim) with '16-mixed'
3. If you hit this during clipping only, disable the fabric-side clipping for that optimizer

Example fix

# before
Fabric(precision="16-mixed", strategy=FSDPPStrategy(), plugins=[...])  # with bnb optimizer
# after
Fabric(precision="bf16-mixed", strategy=FSDPPStrategy())  # bnb optimizer, no scaler

from lightning.fabric.plugins.precision.fsdp import _optimizer_handles_unscaling
uses_scaler = precision_plugin.scaler is not None
if uses_scaler and _optimizer_handles_unscaling(optimizer):
    raise RuntimeError("switch to bf16-mixed or a plain optimizer")

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

def needs_external_unscaling(optimizer) -> bool:
    return not _optimizer_handles_unscaling(optimizer)

try:
    fabric.strategy.precision_plugin.unscale_gradients(optimizer)
except NotImplementedError:
    pass  # optimizer handles unscaling internally

When it happens

Trigger: Using FSDPPrecision with '16-mixed' (so a ShardedGradScaler exists) together with an optimizer whose class implements its own unscaling — e.g. bitsandbytes optimizers — which triggers _optimizer_handles_unscaling(optimizer) to return True when unscale_gradients is called.

Common situations: Mixing bitsandbytes 8-bit/16-bit optimizers with FSDP 16-mixed training; QLoRA-style setups moved under Fabric FSDP.

Related errors

• `precision={precision!r}` does not use a scaler, found {scal
• The optimizer has references to the model's meta-device para
• `precision={precision!r})` is not supported in FSDP. `precis
• Found multiple FSDP models in the given state. Saving checkp
• Unknown state_dict_type: {self._state_dict_type}

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