Lightning-AI / Lightning-AI/pytorch-lightning
[Fabric] FSDPStrategy ignores error_if_nonfinite in gradient clipping
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
# [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`:
```python
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:
```bash
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:
```bash
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:
```python
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the FSDPStrategy gradient-clipping override described in the issue and compare it with the default strategy using reproduce_lightning_fsdp_nonfinite.py. Verify the plain and FSDP commands, then add coverage for the non-finite gradient case; done means error_if_nonfinite=True raises before optimizer.step() and the distributed failure is coordinated across ranks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 70/100