[BUG] ZeRO-3 ignores checkpoint tag when load_optimizer_states=False

Describe the bug

For a non-MoE ZeRO-3 model, load_checkpoint(root, tag="earlier", load_optimizer_states=False) can silently load the weights from a different checkpoint. If root/latest points to later, the returned path and client_state refer to earlier, but the actual model parameters match later.

If the requested checkpoint exists but there is no latest file, the same call instead raises ValueError: Unable to find 'latest' file. Both behaviors come from the same missing argument during FP32 weight reconstruction.

Suggested severity: high, based on silently using the wrong model for evaluation, rollback, or a new fine-tuning run. This assessment is limited to the affected ZeRO-3 reconstruction branch; it does not imply that the default full optimizer-state restore path is affected.

To Reproduce

Baseline: master at 71d316d608a56af2fcc27b84b14cf854b7052eff.

Save the script below as repro_tag.py and run it with DeepSpeed and its CPU dependencies installed:

DS_ACCELERATOR=cpu OMP_NUM_THREADS=1 python repro_tag.py

The script trains a small model, saves two distinct checkpoints, then requests the first checkpoint using a freshly initialized engine. It uses real training, checkpoint files, and parameter comparisons; no loader mocks are involved.

# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import os
import tempfile

import torch
import deepspeed
import deepspeed.comm as dist


def main():
    os.environ.setdefault('LOCAL_RANK', '0')
    torch.set_num_threads(1)
    with tempfile.TemporaryDirectory() as root:
        dist.init_distributed('gloo', auto_mpi_discovery=False,
                              init_method=f'file://{root}/rendezvous', rank=0, world_size=1)
        config = {
            'train_batch_size': 1,
            'zero_allow_untested_optimizer': True,
            'zero_optimization': {
                'stage': 3,
                'reduce_bucket_size': 1000,
                'stage3_prefetch_bucket_size': 1000,
            },
        }

        def make_engine():
            torch.manual_seed(123)
            model = torch.nn.Linear(4, 2)
            optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
            return deepspeed.initialize(model=model, optimizer=optimizer, config=config)[0]

        def snapshot(engine):
            with deepspeed.zero.GatheredParameters(list(engine.module.parameters())):
                return {k: v.detach().clone() for k, v in engine.module.state_dict().items()}

        source = make_engine()
        target = None
        try:
            saved = {}
            for tag in ('earlier', 'later'):
                loss = source(torch.ones(1, 4, device=source.device)).square().mean()
                source.backward(loss)
                source.step()
                saved[tag] = snapshot(source)
                source.save_checkpoint(root, tag=tag, client_state={'label': tag})
            target = make_engine()
            path, client = target.load_checkpoint(root, tag='earlier', load_optimizer_states=False)
            actual = snapshot(target)
            print('Returned path:', path)
            print('Returned label:', client['label'])
            print('Matches later:', all(torch.equal(actual[k], saved['later'][k]) for k in actual))
            for name, value in actual.items():
                torch.testing.assert_close(value, saved['earlier'][name], rtol=0, atol=0)
            print('PASS: all parameters match the requested checkpoint')
        finally:
            if target is not None:
                target.destroy()
            source.destroy()
            dist.destroy_process_group()


if __name__ == '__main__':
    main()

On the baseline, load_checkpoint() returns normally and reports earlier, but the script prints Matches later: True and fails its parameter assertion. That assertion is the reproducer's check, not an error raised by DeepSpeed.

To exercise the missing-latest variant, pass save_latest=False to save_checkpoint() in the same script; the explicit earlier tag still exists, but loading fails while the helper searches for latest.

Expected behavior

An explicit tag should identify the checkpoint used for both metadata and model weights. It should not require a latest file. Omitting the tag should continue to select latest.

Root cause and proposed fix

DeepSpeedEngine._load_checkpoint() calls:

get_fp32_state_dict_from_zero_checkpoint(load_dir)

The helper consequently resolves latest again. Passing the already resolved tag keeps FP32 reconstruction consistent with the checkpoint that the engine selected:

get_fp32_state_dict_from_zero_checkpoint(load_dir, tag=tag)

Reading FP32 weights from optimizer shard files is needed in this path even when optimizer state restoration is disabled. The helper already supports an explicit tag; its documented fallback to latest applies when that argument is omitted.

Validation

Current-base checks on 71d316d608a56af2fcc27b84b14cf854b7052eff plus this fix (2026-09-13):

• New regression: 6 passed on CPU/Gloo and 6 passed on RTX 3090/CUDA/NCCL. Each pytest case executes with both one and two ranks, giving 12 distributed executions per backend.
• Coverage includes explicit older tags with/without latest, both load_module_only values, omitted-tag selection, and a full optimizer-state restore control.
• Existing GPU regressions: TestZeROCheckpoint::test_not_load_optimizer_state[3-False-Adam] and TestZeROCheckpoint::test_load_module_only[3]: 2 passed, each with two ranks.
• All applicable pre-commit hooks on the two modified files and git diff --check passed.

Run the new regression from the source checkout with its test dependencies installed:

export PYTHONPATH="$PWD:$PWD/tests"
export OMP_NUM_THREADS=1 LOCAL_SIZE=2 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
DS_ACCELERATOR=cpu python -m pytest -p pytest_forked --forked \
  tests/unit/checkpoint/test_zero_optimizer.py::TestZeROCheckpointTag \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6
DS_ACCELERATOR=cuda CUDA_VISIBLE_DEVICES=0,1 python -m pytest -p pytest_forked --forked \
  tests/unit/checkpoint/test_zero_optimizer.py::TestZeROCheckpointTag \
  --torch_ver=2.12.1+cu126 --cuda_ver=12.6

The explicit version options describe the tested PyTorch installation; adjust them for another environment.

Earlier full-model integration validation, on base 29d0abbc21f11da806ca14970fa8b3ddb757a152 with the same one-line production fix, used the complete pretrained language model from Qwen3.5-0.8B: 24 decoder layers and 752,393,024 trainable parameters, BF16, ZeRO-3, PyTorch AdamW, sequence length 32, microbatch size 1 per GPU. Both one and four RTX 3090 runs reproduced the wrong-tag load. After the fix, all restored model parameters matched the requested checkpoint exactly, and training could continue. The vision branch was excluded. This was a short training/restoration check, not a convergence or throughput benchmark.

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
deepspeed_not_implemented  [NO] ....... [OKAY]
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
deepspeed_ccl_comm ..... [NO] ....... [OKAY]
deepspeed_shm_comm ..... [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
pin_memory ............. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['<python-environment>/lib/python3.11/site-packages/torch']
torch version .................... 2.12.1+cu126
deepspeed install path ........... ['<source-checkout>/deepspeed']
deepspeed info ................... 0.0.0, [none], [none]
deepspeed wheel compiled w. ...... torch 0.0
shared memory (/dev/shm) size .... 125.77 GB

System info

• Ubuntu 22.04.4 LTS; 2 × Intel Xeon Gold 6133, 80 logical CPUs.
• The current tests import the source checkout through PYTHONPATH. Build metadata was not generated, so ds_report displays 0.0.0 / [none]; the tested Git base is the full SHA above.
• Python 3.11.15; PyTorch 2.12.1+cu126; pytest 8.3.5.
• GPU validation: NVIDIA GeForce RTX 3090 24 GiB; driver 580.173.02; single-node NCCL. The earlier full-model run used Transformers 5.10.4.
• CPU reproduction uses Gloo and DeepSpeed's SHM extension. The bug does not require a GPU.

Launcher context

The standalone script initializes a one-rank Gloo process group directly. Regression tests use DeepSpeed's DistributedTest harness through pytest --forked. Earlier multi-GPU full-model checks used torchrun.

Docker context

Local Python virtual environment; no task-specific Docker image.

Additional context / duplicate search

Before filing, I checked all 225 open PRs by current head commit and searched related issues and PRs across open and closed states. I did not find a fix for this missing engine argument within that search scope.

• #3116 changes when ZeRO optimizer states are loaded; it does not pass the selected tag to this helper.
• #4089 adds a tag option to the offline conversion script; the engine call above still omits it.
• Open PR #8378 changes this same call to preserve frozen-parameter dtypes, but still does not pass tag. The changes address different contracts and may need to be combined when rebasing.

Start in deepspeed/runtime/engine.py at DeepSpeedEngine._load_checkpoint() and inspect the ZeRO-3 FP32 reconstruction path, then run the named TestZeROCheckpointTag regression in tests/unit/checkpoint/test_zero_optimizer.py. Done means an explicit checkpoint tag consistently selects the same model weights and metadata, while omitted-tag loading and full optimizer-state restore remain covered.