Save/Load deepspeed checkpoint with cpu

[jisngprk]

I have some question.

I am using two GPU in ond node

1. I want to load the model with cpu
   How I could load the model ckpts with normal torch.load or deepspeed engine without distributed gpu environment setting?
   If there is some example in DeepspeedExample repo, please let me know.

2. When I save the ckpt , the ckpt is saved in two separate directory that are named with loss. Is it normal that the ckpt is saved separately ? - I am guessing because of loss in name of directory

  if step_cnt % 2000 == 0:
      fstring = 'epoch%d-step%d-loss%.2f' % (epoch, step_cnt, loss_avg)
      self.model_engine.save_checkpoint(wpath, fstring, client_state={
          'epoch': epoch,
          'step': step_cnt})
      logging.info("[Rank - %d, MODEL SAVE]: %s" % (args.rank, os.path.join(wpath, fstring)))

3. To load the model, does we need all of the checkpoints ?

• the deepspeed engine save two checkpoints

# ls checkpoints/test5/epoch0-step12000-loss4.23/
mp_rank_00_model_states.pt                zero_pp_rank_0_mp_rank_00optim_states.pt

# ls checkpoints/test5/epoch0-step12000-loss4.12/
zero_pp_rank_1_mp_rank_00optim_states.pt

Thank you!

[asigalov61]

@jisngprk I am also new to deep speed so I may be wrong but this is what works for many people:

model_engine.save('checkpoint_dir', f'{epoch}_{step}')

model_engine.load('checkpoint_dir')

As far as checkpoints go, I think one is the checkpoint, and another is the full model with state dic + optimizer state + other stuff. At least this is how it works with torch. For example:

Save:

torch.save(model.state_dict(), '/content/model.pth')

checkpoint = {'state_dict': model.state_dict(),'optimizer' :optim.state_dict()}
torch.save(checkpoint, '/content/model_sd_opt.pth')

Load:

torch.load('/content/model_sd_opt.pth')

And if you want to use just the CPU, then you usually just specify it.

I.e.

model.cuda()
model.cpu()

Hope this is helpful.

[jisngprk]

Thank you.

[kinthaiofficial]

Saving and loading DeepSpeed checkpoints to/from CPU is important for scenarios where GPU memory is limited. Here's what works:

Saving with CPU offload:

# In your training loop
if step % save_interval == 0:
    # Save checkpoint - DeepSpeed handles sharding automatically
    model_engine.save_checkpoint(
        save_dir="./checkpoints",
        tag=f"step_{step}",
        # Save optimizer states to CPU to avoid GPU memory spike
        exclude_frozen_parameters=True,
    )

Loading checkpoint:

# During initialization
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=ds_config,
)

# Load from checkpoint
_, client_state = model_engine.load_checkpoint(
    load_dir="./checkpoints",
    tag="step_1000",
    load_optimizer_states=True,
    load_lr_scheduler_states=True,
)

CPU checkpoint consolidation:
For ZeRO stage 3 (where model parameters are sharded across GPUs), you need to consolidate before loading on different hardware:

python -c "
from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
convert_zero_checkpoint_to_fp32_state_dict('./checkpoints/step_1000', './checkpoints/fp32_model.pt')
"

Common pitfall:
ZeRO stage 3 checkpoints are sharded — you need the same number of GPUs to resume, OR you consolidate first with the script above.

For distributed LLM serving:
After training, the consolidated FP32/FP16 checkpoint is what vLLM/Triton/etc. will consume. The ZeRO sharding format is training-only.