Loss spike when resuming from FSDP SHARDED_STATE_DICT checkpoint (possible optimizer-state mismatch)

John6666 · June 28, 2025, 7:05am

4

I know the known issues regarding this.

Final output is upcasted to float32 even though I used mixed_precision="fp16"

opened 03:30PM - 11 Jun 25 UTC

### System Info ```Shell - `Accelerate` version: 1.7.0 - Platform: Linux-6.1.12…3+-x86_64-with-glibc2.35 - `accelerate` bash location: /usr/local/bin/accelerate - Python version: 3.11.13 - Numpy version: 2.0.2 - PyTorch version: 2.6.0+cu124 - PyTorch accelerator: CUDA - System RAM: 12.67 GB - GPU type: Tesla T4 - `Accelerate` default config: Not found ``` ### Information - [ ] The official example scripts - [x] My own modified scripts ### Tasks - [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`) - [x] My own task or dataset (give details below) ### Reproduction Final output is upcasted to float32 even though I used mixed_precision="fp16". I don't understand why. Minimal example: ``` import torch import torch.nn as nn from accelerate import Accelerator from torch import amp class Net(nn.Module): def __init__(self): super().__init__() self.layer = nn.Conv2d(3, 3, kernel_size=3, stride=1, padding=1) def forward(self, x): x = self.layer(x) print("x2", x.dtype, x.device) return x accelerator = Accelerator(mixed_precision="fp16") model = Net() model = accelerator.prepare(model) model.eval() x = torch.rand((3, 1024, 1024), dtype=torch.float16, device="cuda") print("x1", x.dtype, x.device) with torch.no_grad(), amp.autocast(device_type='cuda', dtype=torch.float16): x = model(x) print("x3", x.dtype, x.device) ``` Output is : ``` x1 torch.float16 cuda:0 x2 torch.float16 cuda:0 x3 torch.float32 cuda:0 ``` Can you help please? Am I doing something wrong? ### Expected behavior Expected output is : ``` x1 torch.float16 cuda:0 x2 torch.float16 cuda:0 x3 torch.float16 cuda:0 ```

Topic	Replies	Views
Difficulty with checkpoint saving and loading (trainer+ FSDP accelerate) Beginners	556	April 1, 2024
Deepspeed resume training from saved states 🤗Accelerate	1263	September 8, 2022
Weird behavior when saving checkpoint in DDP 🤗Accelerate	47	August 11, 2024
FSDP training not saving the best checkpoint and load from checkpoint fails 🤗Transformers	778	January 23, 2024
Eval Loss spike Seq2seq Trainer Resume from Checkpoint 🤗Transformers	519	June 22, 2021

Loss spike when resuming from FSDP SHARDED_STATE_DICT checkpoint (possible optimizer-state mismatch)

Related topics