Unable to save/load optimizer state after accelerator.prepare()

[fkggggg]

System Info

accelerate version==1.7.0, OS==Linux, python version==3.10.16, 
"deepspeed_config": {
    "zero_stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },q
  }

Information

• The official example scripts
• 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)
• My own task or dataset (give details below)

Reproduction

After preparing the optimizer with accelerator.prepare(), I'm unable to properly save and load its state:

1. When trying to save the prepared optimizer directly:

optimizer = torch.optim.Adam(model.parameters())
optimizer = accelerator.prepare(model, optimizer)[1]

# Saving
torch.save({
    'optimizer_state_dict': optimizer.state_dict(),
}, 'checkpoint.pt')

# Loading 
checkpoint = torch.load('checkpoint.pt')
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # Fails

This fails with:

KeyError: 'param_groups'

2.I've tried looking at the saved optimizer state:

ckpt = torch.load('checkpoint.pt', map_location='cpu')
print(ckpt['optimizer_state_dict'].keys())

Shows DeepSpeed format keys:

dict_keys(['loss_scaler', 'dynamic_loss_scale', 'overflow', 'clip_grad', 'base_optimizer_state', 'single_partition_of_fp32_groups', 'zero_stage', 'group_paddings', 'partition_count', 'ds_version', 'param_slice_mappings'])

So I tried to load from base_optimizer_state:

optimizer.load_state_dict(checkpoint['optimizer_state_dict']['base_optimizer_state'])

Fails with:

ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group

3.When trying to unwrap the optimizer first:

unwrapped_optimizer = accelerator.unwrap_model(optimizer)  # Fails

This fails with:

AttributeError: 'DeepSpeedOptimizerWrapper' object has no attribute '_modules'

Expected behavior

Should be able to save and load optimizer states when using Accelerate, similar to regular PyTorch training.

[S1ro1]

You can just use accelerator.save_state and accelerator.load_state which does work, I don't think saving optimizers/models like this is supposed to work with DeepSpeed.

As for why this is the case, you can notice that DeepSpeed saves num_dp_ranks optimizer state_dicts, because you're using ZERO2, then what DeepSpeed does on the background is that each dp rank holds only a optimizer of only the dp shard rank. That's why you're getting this issue: ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group, because you're loading into a different sized optimizer. DeepSpeed does a bunch of reordering/flattening of optimizer param groups, that's why it's hard to reproduce.

# divide the flat weights into near equal partition equal to the data parallel degree
# each process will compute on a different part of the partition
data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)

Copy of a relevant part from deepspeed/runtime/zero/zero_stage_1_and_2.py::DeepSpeedZeroOptimizer