From 35e294e035c474a4fc238b25180ec1067ec16083 Mon Sep 17 00:00:00 2001 From: kzahiri1 Date: Wed, 2 Sep 2026 12:48:10 -0700 Subject: [PATCH] Fix ZeRO checkpoint DP merge on scalar optimizer state merge_state concatenated every tensor it met, so reshaping a ZeRO 1/2 checkpoint to a smaller dp_degree died on torch's 0-dim step counter with "zero-dimensional tensor (at position 0) cannot be concatenated". Return replicated scalars unchanged, matching the value.dim() > 0 guard in stage_1_and_2.py. merge_state_dict also iterated only dict_b, silently dropping keys held by dict_a alone, and passed [str(key)] instead of the accumulated path, so the diagnostic printed the leaf key rather than the nested one. Merge from dict_a, then add dict_b's exclusive keys, and propagate key_list. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01DQughgr95y4B9H1jfaQH8o Signed-off-by: kzahiri1 --- deepspeed/checkpoint/reshape_utils.py | 14 ++++- .../checkpoint/test_reshape_checkpoint.py | 55 ++++++++++++++++++- 2 files changed, 63 insertions(+), 6 deletions(-) diff --git a/deepspeed/checkpoint/reshape_utils.py b/deepspeed/checkpoint/reshape_utils.py index 137607721ebf..9bf8e89dafa6 100644 --- a/deepspeed/checkpoint/reshape_utils.py +++ b/deepspeed/checkpoint/reshape_utils.py @@ -80,12 +80,16 @@ def _key_list_to_string(key_list): def merge_state_dict(dict_a, dict_b, key_list): merged_dict = type(dict_a)({}) - for key, value in dict_b.items(): - if key in dict_a.keys(): - merged_dict[key] = merge_state(dict_a[key], dict_b[key], [str(key)]) + for key, value in dict_a.items(): + if key in dict_b: + merged_dict[key] = merge_state(value, dict_b[key], key_list + [str(key)]) else: merged_dict[key] = value + for key, value in dict_b.items(): + if key not in dict_a: + merged_dict[key] = value + return merged_dict @@ -108,6 +112,10 @@ def merge_state(state_a, state_b, key_list=[]): elif type(state_a) in (list, tuple): return type(state_a)(merge_state_list(state_a, state_b, key_list)) elif torch.is_tensor(state_a): + # Scalars such as the optimizer step counter are replicated across ranks rather + # than partitioned, and torch.cat rejects 0-dim tensors outright. + if state_a.dim() == 0: + return state_a return torch.cat([state_a, state_b], 0) else: return state_a diff --git a/tests/unit/checkpoint/test_reshape_checkpoint.py b/tests/unit/checkpoint/test_reshape_checkpoint.py index d6edca485ee2..c4a52f0c4230 100644 --- a/tests/unit/checkpoint/test_reshape_checkpoint.py +++ b/tests/unit/checkpoint/test_reshape_checkpoint.py @@ -9,9 +9,12 @@ import pytest import torch -from deepspeed.checkpoint import DeepSpeedCheckpoint, ZeROCheckpoint, get_model_3d_descriptor, model_3d_desc -from deepspeed.checkpoint.constants import (AUTOEP_LAYERS_KEY, CHECKPOINT_PARALLEL_DIMS, CHECKPOINT_PP_DEGREE, - CHECKPOINT_TP_DEGREE, PARAM_SHAPES, UNIVERSAL_CHECKPOINT_INFO) +from deepspeed.checkpoint import (DeepSpeedCheckpoint, ZeROCheckpoint, get_model_3d_descriptor, merge_state, + model_3d_desc) +from deepspeed.checkpoint.constants import (AUTOEP_LAYERS_KEY, BASE_OPTIMIZER_STATE, CHECKPOINT_PARALLEL_DIMS, + CHECKPOINT_PP_DEGREE, CHECKPOINT_TP_DEGREE, GROUP_PADDINGS, + OPTIMIZER_STATE_DICT, PARAM_SHAPES, PARTITION_COUNT, + UNIVERSAL_CHECKPOINT_INFO) from deepspeed.checkpoint.ds_to_universal import _aggregate_autoep_zero12_metadata from deepspeed.runtime.engine import _checkpoint_parallel_metadata @@ -122,3 +125,49 @@ def tracked_load(path, *args, **kwargs): assert checkpoint.get_checkpoint_info(UNIVERSAL_CHECKPOINT_INFO) == {"source": "writer"} assert (checkpoint.pp_degree, checkpoint.tp_degree, checkpoint.dp_degree) == (2, 1, 2) assert model_loads == checkpoint.mp_rank_files + + +def _write_zero12_optim_checkpoint(tmpdir, dp_degree): + """Write a dp_degree-way ZeRO 1/2 checkpoint holding a real torch optimizer state_dict.""" + torch.save({PARAM_SHAPES: [{}], "ds_config": {}}, os.path.join(str(tmpdir), "mp_rank_00_model_states.pt")) + for dp_rank in range(dp_degree): + param = torch.nn.Parameter(torch.ones(8)) + optimizer = torch.optim.AdamW([param], lr=1e-3) + param.grad = torch.ones(8) + optimizer.step() + state = { + OPTIMIZER_STATE_DICT: { + BASE_OPTIMIZER_STATE: optimizer.state_dict(), + GROUP_PADDINGS: [0], + PARTITION_COUNT: [dp_degree], + } + } + torch.save(state, os.path.join(str(tmpdir), f"bf16_zero_pp_rank_{dp_rank}_mp_rank_00_optim_states.pt")) + + +def test_reshape_merges_zero12_optimizer_state_with_scalar_step(tmpdir): + _write_zero12_optim_checkpoint(tmpdir, dp_degree=2) + checkpoint = ZeROCheckpoint(str(tmpdir)) + checkpoint.reshape(model_3d_desc(pp_degree=1, tp_degree=1, dp_degree=1)) + + merged = checkpoint.get_state_for_rank(0, 0, 0)[OPTIMIZER_STATE_DICT][BASE_OPTIMIZER_STATE]["state"][0] + + assert merged["exp_avg"].numel() == 16 + assert merged["step"].dim() == 0 + + +def test_merge_state_keeps_keys_missing_from_the_second_dict(): + dict_a = {"exp_avg": torch.ones(4), "rank_0_only": torch.ones(2)} + dict_b = {"exp_avg": torch.ones(4)} + + merged = merge_state(dict_a, dict_b) + + assert list(merged) == ["exp_avg", "rank_0_only"] + assert merged["exp_avg"].numel() == 8 + + +def test_merge_state_reports_the_full_key_path_on_mismatch(capsys): + with pytest.raises(ValueError, match="Cannot merge lists of different lengths"): + merge_state({"opt": {"groups": [1, 2]}}, {"opt": {"groups": [1]}}) + + assert "opt.groups" in capsys.readouterr().out