Lightning-AI / Lightning-AI/pytorch-lightning

[Fabric] FSDPStrategy ignores error_if_nonfinite in gradient clipping

Open
#21,928 1 comment 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.