Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

Follow-ups for restoring fp32 masters without optimizer state

Open
  • Is something wrong?
#8,890 0 comments 0 reactions 1 assignee View on GitHub

Maintainers usually reply within 3 days

@tjruwase is already working on this.

Since Oct 22, 2026.

Assessment

This issue has not been assessed yet.

Description

A fresh-optimizer resume (load_optimizer_states=False) can restore the saved fp32 master weights through DeepSpeedEngine.load_checkpoint(..., load_fp32_weights=None). None inherits zero_optimization.load_from_fp32_weights (default True). load_module_only=True still rebuilds the masters from the loaded bf16/fp16 weights.

Two cases are unsupported when that flag stays True, because the masters live in the per-rank *_optim_states.pt shards alongside Adam's moments:

  • A parameter_alignment change, which is usually also a data-parallel size change. The saved partitions do not match the new buffers.
  • A checkpoint that has no *_optim_states.pt files.

Those resumes should keep failing. load_from_fp32_weights=False is the supported way to skip the shard and rebuild the masters. The alignment RuntimeError still tells the caller to disable optimizer-state loading, which is not sufficient: load_optimizer_states=False with the default config still opens the shard.

The items below do not fail current unit tests and can land separately.

ZeRO-3 serial checkpoint load can hang

zero_optimization.pipeline_loading_checkpoint receives the serial token in _load_zero_checkpoint before the data-parallel world-size check (deepspeed/runtime/engine.py, the dist.recv on every local rank except 0). A non-MoE resume that then returns early, because the saved fp32 partitions cannot be resized, never reaches the dist.send in stage3.py load_state_dict. Rank 0 returns and the next local rank waits.

tests/unit/checkpoint/test_zero_optimizer.py::test_pipeline_checkpoint_loading does not hit this. It reloads at the same world size with load_module_only=True.

ZeRO-3 MoE is not this case. A data-parallel size change with optimizer states off still loads the shard, because the expert weights live there.

Master-only tests accept a stale initial master

Several assertions only check that the masters sit on the bf16 grid and differ from the saved fp32 values. A master left at its initial ones passes both checks. Compare the loaded master to the bf16 rounding of the saved master.

Missing coverage

  • ZeRO-3 rigid and universal refresh_fp32_params when load_fp32_weights=False.
  • The dedicated BF16_Optimizer path (ZeRO stage 1, bf16 parameters, fp32 gradient accumulation).
  • fp16.
  • load_optimizer_states=True together with load_fp32_weights=False: moments load, masters are rebuilt.
  • A data-parallel size change that still requests optimizer states raises ZeRORuntimeException.

ZeRO-3 NVMe load_module_only=True refreshes an empty fp32 buffer

load_module_only=True calls refresh_fp32_params() inside _load_checkpoint before the NVMe branch in load_checkpoint. A swapped ZeRO-3 subgroup stores its fp32 partition as torch.empty(0). _restore_from_bit16_weights copies the fp16 partition into that empty tensor and raises, so the fp32 swap file is never copied.

A data-parallel size check on the NVMe branch does not cover this. The refresh throws first. A same-size NVMe load with load_module_only=False and load_optimizer_states=False still copies the fp32 swap file.

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 21h
Merged PRs (30d)
66

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.