Lightning-AI / Lightning-AI/pytorch-lightning

Mixed precision autocast is disabled when running FSDP on CPU

Open
#21,815 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug ver: 2.5.x ver: 2.6.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

The `FSDPPrecision` plugin hardcodes the `"cuda"` device type in `forward_context` when using mixed precision (e.g., `bf16-mixed`).

When FSDP is run on CPU with mixed precision, this hardcoded `"cuda"` value causes PyTorch to print a warning and disable autocast. Because autocast is disabled, CPU operations are not cast to `bfloat16` during the forward pass, which can lead to type mismatches (e.g., `Float` inputs meeting `BFloat16` weights) or incorrect benchmarking results.

### What version are you seeing the problem on?

v2.5

### Reproduced in studio

_No response_

### How to reproduce the bug

```python
import torch
from lightning.fabric.plugins.precision.fsdp import FSDPPrecision

# Initialize precision plugin for bf16-mixed
precision = FSDPPrecision(precision="bf16-mixed")

# On a CPU-only environment, enter forward context
with precision.forward_context():
# This will trigger: UserWarning: CUDA is not available... Disabling autocast.
print("Autocast CPU enabled:", torch.is_autocast_enabled("cpu"))
# Expected: True (if we want bf16 mixed on CPU)
# Actual: False
```

### Error messages and logs

```
UserWarning: CUDA is not available or torch_xla is imported. Disabling autocast.
```
This can also lead to:
```
RuntimeError: expected scalar type Float but found BFloat16
```
or similar type mismatch errors when inputs are not auto-casted to match the FSDP sharded parameter type.

### Environment

Current environment

```
- PyTorch Lightning Version: master/latest
- PyTorch Version: 2.0+
- OS: Linux/macOS/Windows
- Accelerator: CPU
```

### More info

`FSDPPrecision` should be updated to accept or dynamically set the target `device_type` (e.g., propagating `self.root_device.type` from the FSDP strategy during setup) and pass it to `torch.autocast`.

cc @ethanwharris

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 in the FSDPPrecision implementation, especially forward_context, and trace how the plugin gets its device from FSDP strategy setup. Reproduce the issue with bf16-mixed on CPU and inspect how torch.autocast receives its device type. Done means CPU autocast is enabled without the CUDA warning while existing CUDA behavior remains correct.

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
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.