[BUG] model.load_checkpoint out of memory

System Info

accelerate 0.20.3
python 3.10
numpy 1.24.3
torch 2.0.1
accelerate config:
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_config: {}
fsdp_config: {}
main_training_function: main
megatron_lm_config: {}
mixed_precision: bf16
num_machines: 1
num_processes: 8
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

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

I pretrain a 27B model from scratch with deepspeed stage 3 (no cpu offload) in 8*80G A100, batch size of each gpu is 2. And I use 'accelerator.save_state()' to save the optimizer/lr scheduler. The saved files are as follows:

When I want to resume from above saved checkpoint, I use the following code:
model, train_dataloader, optimizer, lr_scheduler = accelerator.prepare(
model, train_dataloader, optimizer, lr_scheduler
)
accelerator.load_state('/pretrained_model/xxx')
But I got the error of CUDA out of memory. Can you help me? Thanks a lot.

Expected behavior

I expected it cannot caused the cuda out of memory.

Start with the accelerator.prepare and accelerator.load_state calls in the reproduction, then inspect how the saved optimizer and scheduler checkpoint is restored under DeepSpeed ZeRO stage 3. Reproduce with the stated 27B model and 8×80G A100 configuration; done means resuming the checkpoint without a CUDA out-of-memory error.