bf16-mixed causing issues with fused AdamW #21435
Description
Bug description
bfloat16-mixed is incompatible with fused optimizers on the lightning framework. After some digging I found PR #15555 as a response to issue [#15501] which makes it seem like this is a safety net introduced by the lightning team.
Reproduction Code:
import torch import torch.nn as nn from torch.utils.data import DataLoader, TensorDataset import lightning as L class SimpleModel(L.LightningModule): def __init__(self): super().__init__() self.layer = nn.Linear(10, 1) def forward(self, x): return self.layer(x) def training_step(self, batch, batch_idx): x, y = batch y_hat = self(x) loss = nn.functional.mse_loss(y_hat, y) return loss def configure_optimizers(self): return torch.optim.AdamW(self.parameters(), lr=1e-3, fused=True) def main(): X, y = torch.randn(100, 10), torch.randn(100, 1) dataloader = DataLoader(TensorDataset(X, y), batch_size=16) trainer = L.Trainer( accelerator="gpu", devices=1, precision="16-mixed", # Using bf16-true passes gradient_clip_val=1.0, max_epochs=1, logger=False, enable_checkpointing=False, enable_progress_bar=False, ) print(f"Scaler: {trainer.precision_plugin.scaler}") trainer.fit(SimpleModel(), dataloader) if __name__ == "__main__": main() Error:
RuntimeError: The current optimizer, AdamW, does not allow for gradient clipping because it performs unscaling of gradients internally. HINT: Are you using a 'fused' optimizer? Expected Solution:
I believe the key error is in this function handling. I'm a newbie to the pytorch-lightning OSS community and internal frameworks, but this is in the native-amp code, and is what's in traceback. I think it should have an additional condition which is basically and trainer.precision_plugin.scaler is not None (or however it's accessed in the internal API), which would prevent the unnecessary handle in the bf16 case.
if clip_val > 0 and _optimizer_handles_unscaling(optimizer): raise RuntimeError( f"The current optimizer, {type(optimizer).__qualname__}, does not allow for gradient clipping" " because it performs unscaling of gradients internally. HINT: Are you using a 'fused' optimizer?" ) The _step_supports_amp_scaling flag just indicates the optimizer can handle internal unscaling when a GradScaler is present. With bf16 (no scaler), there's nothing to unscale - gradient clipping works normally.
Lightning's check should be:
if clip_val > 0 and self.scaler is not None and _optimizer_handles_unscaling(optimizer):
raise RuntimeError(...)
Not just:
if clip_val > 0 and _optimizer_handles_unscaling(optimizer):
raise RuntimeError(...)
To verify that this is undesired behavior here's a PyTorch only code that works perfectly fine:
import torch import torch.nn as nn def main(): print("PyTorch:", torch.__version__) print() # Simple model model = nn.Linear(10, 1).cuda() # Fused AdamW optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, fused=True) print(f"_step_supports_amp_scaling: {optimizer._step_supports_amp_scaling}") for step in range(10): # Random data x = torch.randn(16, 10, device="cuda") y = torch.randn(16, 1, device="cuda") optimizer.zero_grad() # bf16 autocast (no GradScaler needed for bf16) with torch.autocast("cuda", dtype=torch.bfloat16): y_hat = model(x) loss = nn.functional.mse_loss(y_hat, y) # Backward loss.backward() # Gradient clipping - works fine! grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # Step optimizer.step() if __name__ == "__main__": main() Personal Setup:
RTX A5000
pytorch-lightning version: 2.5.6
What version are you seeing the problem on?
v2.5
Reproduced in studio
No response
How to reproduce the bug
Error messages and logs
# Error messages and logs here please Environment
Current environment
#- PyTorch Lightning Version (e.g., 2.5.0): #- PyTorch Version (e.g., 2.5): #- Python version (e.g., 3.12): #- OS (e.g., Linux): #- CUDA/cuDNN version: #- GPU models and configuration: #- How you installed Lightning(`conda`, `pip`, source): More info
No response