Lightning-AI/pytorch-lightning · error · TypeError

AMP and the LBFGS optimizer are not compatible.

Error message

AMP and the LBFGS optimizer are not compatible.

What it means

GradScaler-based AMP cannot wrap LBFGS because LBFGS's step performs multiple function evaluations, which the scaler's unscale/skip-if-inf/nan logic cannot handle. Lightning therefore raises TypeError when optimizer_step is called with an LBFGS optimizer while a scaler is active (i.e. '16-mixed').

Source

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

        return apply_to_collection(data, function=_convert_fp_tensor, dtype=Tensor, dst_type=torch.get_default_dtype())

    @override
    def backward(self, tensor: Tensor, model: Optional[Module], *args: Any, **kwargs: Any) -> None:
        if self.scaler is not None:
            tensor = self.scaler.scale(tensor)
        super().backward(tensor, model, *args, **kwargs)

    @override
    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)
        if isinstance(optimizer, LBFGS):
            raise TypeError("AMP and the LBFGS optimizer are not compatible.")
        # 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 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:

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Switch precision to 'bf16-mixed' (scaler is None, so the scaler path is skipped) or disable mixed precision
  2. Replace LBFGS with a comparable multi-step optimizer (Adam, L-BFGS via a library that supports AMP)
  3. Use a custom Precision plugin that implements optimizer_step for LBFGS without the scaler

Example fix

# before
fabric = Fabric(precision="16-mixed")  # + torch.optim.LBFGS

# after
fabric = Fabric(precision="bf16-mixed")  # no scaler, LBFGS works
Defensive patterns

Strategy: validation

Validate before calling

import torch
from lightning.fabric.plugins.precision.amp import MixedPrecision

def check_amp_compatible(optimizer, plugin: MixedPrecision):
    if plugin.scaler is not None and isinstance(optimizer, torch.optim.LBFGS):
        raise SystemExit("LBFGS requires bf16-mixed or full precision")

Type guard

def amp_lbfgs_ok(plugin, optimizer) -> bool:
    return plugin.scaler is None or not isinstance(optimizer, torch.optim.LBFGS)

Try / catch

try:
    step_out = plugin.optimizer_step(optimizer)
except TypeError as e:
    if "LBFGS" not in str(e):
        raise
    # switch precision and rebuild

Prevention

When it happens

Trigger: Using torch.optim.LBFGS with MixedPrecision(precision='16-mixed'); the error fires on the first optimizer step, not at setup.

Common situations: Porting LBFGS-based optimization (e.g. some classical/physics or GAN-style setups) into an AMP-enabled Fabric/Trainer run.

Related errors


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