[BUG] Python checkpoint writers fail to reload full ZeRO-3 state with default torch.load

Describe the bug

With a Python checkpoint writer configured, ZeRO-3 training can save a checkpoint successfully but fail to load that same checkpoint:

_pickle.UnpicklingError: Weights only load failed.
Unsupported global: GLOBAL deepspeed.runtime.zero.config.ZeroStageEnum

This affects both FastCheckpointEngine and DecoupledCheckpointEngine under PyTorch's default weights-only loading behavior. ZeroStageEnum is part of DeepSpeed's own saved optimizer state; the user does not need to add custom objects to trigger the failure. The asynchronous reproduction completes the pending write before attempting restoration.

Suggested severity: medium-high for the affected writer configurations because checkpoint restoration is blocked. The files are not permanently corrupted; loading trusted files through an existing full-state loading path is a workaround.

To Reproduce

Baseline: master at 71d316d608a56af2fcc27b84b14cf854b7052eff, using PyTorch 2.12.1+cu126 with no weights-only environment override or custom safe-global allowlist.

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

DS_ACCELERATOR=cpu OMP_NUM_THREADS=1 python repro_writer.py

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

import torch
import deepspeed
import deepspeed.comm as dist


def main():
    for name in ('LOCAL_RANK', 'CROSS_RANK'):
        os.environ.setdefault(name, '0')
    os.environ.setdefault('CROSS_SIZE', '1')
    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,
            },
            'checkpoint': {'writer': {'type': 'python'}},
        }

        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 train(engine):
            loss = engine(torch.ones(1, 4, device=engine.device)).square().mean()
            engine.backward(loss)
            engine.step()

        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:
            train(source)
            saved = snapshot(source)
            source.save_checkpoint(root, tag='resume')
            # Also commits a pending save if the writer is configured as decoupled.
            train(source)
            expected_next_step = snapshot(source)
            target = make_engine()
            target.load_checkpoint(root, tag='resume')
            for name, value in snapshot(target).items():
                torch.testing.assert_close(value, saved[name], rtol=0, atol=0)
            train(target)
            for name, value in snapshot(target).items():
                torch.testing.assert_close(value, expected_next_step[name], rtol=0, atol=0)
            print('PASS: restored parameters and the next Adam update match exactly')
        finally:
            if target is not None:
                target.destroy()
            source.destroy()
            dist.destroy_process_group()


if __name__ == '__main__':
    main()

The baseline fails inside the checkpoint engine's load() call with the error above. To exercise the asynchronous engine, add "decoupled": True beside "type": "python". The script already performs the next optimizer step to commit the pending asynchronous save.

Expected behavior

An engine should be able to restore the complete training state it saved. In this deterministic reproduction, the restored parameters and the next Adam update should match uninterrupted training exactly.

Root cause and proposed fix

FastCheckpointEngine.load() and DecoupledCheckpointEngine.load() call torch.load(path, map_location=map_location) without specifying weights_only.

PyTorch defaults to weights_only=True starting in 2.6 when pickle_module is not supplied. DeepSpeed's current complete ZeRO state includes an enum and loss-scaler objects that are not accepted by the default restricted loader. See the PyTorch serialization documentation.

The minimal compatibility fix is to pass weights_only=False in both calls, matching the existing TorchCheckpointEngine behavior. Both call sites belong to the same root cause.

This proposal is for self-produced or otherwise trusted training checkpoints. weights_only=False permits pickle code execution; it is not a safe-loading mechanism for untrusted files. An audited allowlist or a versioned tensor/primitive-only checkpoint format would be alternative designs with a wider compatibility scope.

Validation

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

• New regression: 4 passed on CPU/Gloo and 4 passed on RTX 3090/CUDA/NCCL. Each pytest case executes with both one and two ranks, giving 8 distributed executions per backend.
• The four cases cover synchronous/asynchronous Python writers and legacy/ZIP serialization, including exact saved-parameter restoration, the restored training step count, and the next Adam update against uninterrupted training.
• Existing GPU regression: TestOtherOptimizerCheckpoint::test_checkpoint_fp32_optimizer: 1 passed with two ranks.
• All applicable pre-commit hooks on the three 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_other_optimizer.py::TestCheckpointWriterResume \
  --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_other_optimizer.py::TestCheckpointWriterResume \
  --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 two-line production fix, covered synchronous/asynchronous Python writers × one/four RTX 3090 GPUs. It loaded the complete pretrained Qwen3.5-0.8B language model: 24 decoder layers, 752,393,024 trainable parameters, BF16, ZeRO-3, PyTorch AdamW, sequence length 32, microbatch size 1 per GPU. All four baseline scenarios failed at deserialization. After the fix, every restored parameter and every parameter after one additional update matched the uninterrupted reference exactly; the training step count also matched. These GPU runs used ZIP serialization. The vision branch was excluded.

These are short restoration checks. Multi-node training, AIO/GDS writers, NVMe offload, long-run convergence, and the full DeepSpeed suite were not validated.

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. Earlier full-model runs used Transformers 5.10.4.
• CPU reproduction uses Gloo and DeepSpeed's SHM extension. Both legacy and ZIP serialization are covered by regression tests.

Launcher context

The standalone script initializes a one-rank Gloo process group directly and supplies single-node CROSS_RANK / CROSS_SIZE values used by the writer. Regression tests use DeepSpeed's DistributedTest harness through pytest --forked; asynchronous tests use a real writer subprocess. 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 covering these two load calls within that search scope.

• #6751 specifies the loading mode at other checkpoint load sites, including TorchCheckpointEngine; these two engines still omit it.
• #7741 / #7742 concern asynchronous writer timeouts and process health, not the deserialization failure here.
• #7635 concerns passing a weights_only argument to the public DeepSpeedEngine.load_checkpoint() API. This report concerns the internal writer load calls and does not propose changing that public API.

Read FastCheckpointEngine.load() in deepspeed/runtime/checkpoint_engine/fast_checkpoint_engine.py and DecoupledCheckpointEngine.load() in deepspeed/runtime/checkpoint_engine/decoupled_checkpoint_engine.py, then run TestCheckpointWriterResume in tests/unit/checkpoint/test_other_optimizer.py. Done means synchronous and asynchronous Python writers restore the saved parameters, step count, and next Adam update under the reported PyTorch loading behavior.