[Fabric] FSDPStrategy ignores error_if_nonfinite in gradient clipping

[Fabric] FSDPStrategy ignores error_if_nonfinite in gradient clipping

Bug description

Fabric.clip_gradients(..., max_norm=..., error_if_nonfinite=True) raises for a non-finite gradient norm with the default strategy, but does not raise when the model is wrapped with FSDPStrategy.

The FSDP strategy accepts error_if_nonfinite, but delegates to PyTorch FSDP's clip_grad_norm_() without enforcing it. PyTorch FSDP's method does not accept an error_if_nonfinite argument. As a result, clipping returns a non-finite norm and training can continue into optimizer.step(), corrupting parameters instead of failing fast.

We observed this in a distributed HYBRID_SHARD training run: after the first non-finite gradient norm, error_if_nonfinite=True did not raise and subsequent optimizer/model values became non-finite.

What version are you seeing the problem on?

• Lightning: 2.5.6
• PyTorch: 2.10.0+cu130
• Python: 3.11
• Distributed strategy in the original run: FSDP HYBRID_SHARD

The FSDP implementation in Lightning 2.6.5 and the current master branch appears to retain the same behavior, based on source inspection.

How to reproduce the bug

Save the following as reproduce_lightning_fsdp_nonfinite.py:

import argparse

import torch
from lightning import Fabric
from lightning.fabric.strategies import FSDPStrategy
from torch import nn


parser = argparse.ArgumentParser()
parser.add_argument("--strategy", choices=("plain", "fsdp"), required=True)
args = parser.parse_args()

if args.strategy == "fsdp":
    fabric = Fabric(
        accelerator="cuda",
        devices=1,
        strategy=FSDPStrategy(use_orig_params=True),
    )
else:
    fabric = Fabric(accelerator="cpu", devices=1)

fabric.launch()

model = nn.Linear(4, 4, bias=False)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
model, optimizer = fabric.setup(model, optimizer)

x = torch.ones(2, 4, device=fabric.device)
loss = model(x).sum() * torch.tensor(float("nan"), device=fabric.device)
fabric.backward(loss)

try:
    norm = fabric.clip_gradients(
        model,
        optimizer,
        max_norm=1.0,
        error_if_nonfinite=True,
    )
except RuntimeError as ex:
    fabric.print(f"RAISED as expected: {ex}")
else:
    fabric.print(f"DID NOT RAISE: norm={norm}")
    optimizer.step()

Run the non-FSDP control:

python reproduce_lightning_fsdp_nonfinite.py --strategy plain

The control raises a RuntimeError reporting that the total gradient norm is non-finite.

Run the FSDP case on a machine with two CUDA GPUs:

torchrun --standalone --nproc-per-node=2 \
  reproduce_lightning_fsdp_nonfinite.py --strategy fsdp

The FSDP case does not raise and prints a non-finite norm. If the optimizer step is allowed to run, parameters can become non-finite.

Expected behavior

When error_if_nonfinite=True, all strategies that accept this argument should raise before an optimizer step if the total gradient norm is nan, inf, or -inf.

For distributed strategies, the failure should be coordinated across ranks so that one rank does not raise while peers continue into a collective.

Actual behavior

The default strategy raises, but FSDPStrategy silently ignores the requested fail-fast behavior and returns the non-finite norm.

Additional context

The relevant FSDP override accepts error_if_nonfinite but effectively does:

self.precision.unscale_gradients(optimizer)
return module.clip_grad_norm_(max_norm=max_norm, norm_type=norm_type)

PyTorch FSDP's clip_grad_norm_() has no error_if_nonfinite parameter. With a non-finite norm, its scaling coefficient is also non-finite, so applying the coefficient can propagate non-finite values through the gradients.

This report concerns Lightning Fabric's public gradient-clipping API, not the PyTorch Lightning Trainer gradient-clipping configuration.