Fix universal checkpoint resume across AutoTP sizes

deepspeed/runtime/engine.py

@@ -4779,8 +4779,15 @@ def _load_checkpoint(self,
 
         is_pipe_parallel = isinstance(self.module, PipelineModule)
 
-        load_path, checkpoint, _ = sd_loader.load(self.mp_world_size,
-                                                  self.checkpoint_mp_rank,
+        checkpoint_mp_world_size = self.mp_world_size
+        checkpoint_mp_rank = self.checkpoint_mp_rank
+        if self.load_universal_checkpoint():
+            # UC restores weights from zero/. Read metadata from the corresponding
+            # source rank without invoking Megatron's model-specific weight merger.
+            checkpoint_mp_rank = checkpoint_mp_rank * len(ckpt_list) // checkpoint_mp_world_size
+            checkpoint_mp_world_size = len(ckpt_list)
+        load_path, checkpoint, _ = sd_loader.load(checkpoint_mp_world_size,
+                                                  checkpoint_mp_rank,
                                                   is_pipe_parallel=is_pipe_parallel)
 
         if checkpoint is None:

deepspeed/utils/tensor_fragment.py

@@ -79,7 +79,7 @@ def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys):
                 hp_fragment_address = lp._hp_mapping.get_hp_fragment_address()
                 hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel)
                 hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data)
-                lp._hp_mapping.hp_fragment = hp_fragment
+                lp._hp_mapping.optim_fragment[key] = hp_fragment
 
         optim_state[hp_param][key] = buffer
 

tests/unit/checkpoint/test_autotp_uc_checkpoint.py

@@ -32,7 +32,7 @@
 from deepspeed.utils import RepeatingLoader, groups
 from deepspeed.module_inject.tp_shard import get_shard_size_list
 
-from unit.common import DistributedTest
+from unit.common import DistributedTest, DistributedFixture
 
 
 class _DummyAddress:
@@ -1567,3 +1567,181 @@ def test_strict_correctness(self, tmpdir):
 
         # 2c: optimizer state usable.
         _train_steps(restored_engine, hidden_dim, steps=1)
+
+
+class AffineResumeModel(torch.nn.Module):
+    """A column/row pair with a replicated, post-reduction row bias."""
+
+    def __init__(self):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(16, 16)
+        self.fc2 = torch.nn.Linear(16, 16)
+
+    def forward(self, x):
+        return self.fc2(torch.tanh(self.fc1(x)))
+
+
+def _affine_resume_engine(tp_size, load_universal=False):
+    torch.manual_seed(1234)
+    model = AffineResumeModel()
+    config = {
+        "train_micro_batch_size_per_gpu": 2,
+        "zero_optimization": {
+            "stage": 1
+        },
+        "zero_allow_untested_optimizer": True,
+        "checkpoint": {
+            "load_universal": load_universal
+        },
+    }
+    if tp_size > 1:
+        config["tensor_parallel"] = {
+            "autotp_size": tp_size,
+            "partition_config": {
+                "use_default_specs":
+                False,
+                "layer_specs": [
+                    {
+                        "patterns": [r".*fc1\.weight$"],
+                        "partition_type": "column"
+                    },
+                    {
+                        "patterns": [r".*fc2\.weight$"],
+                        "partition_type": "row"
+                    },
+                ],
+            },
+        }
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.8, 0.95), eps=1e-6)
+    engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config)
+    return engine
+
+
+def _affine_resume_batch(engine, step):
+    generator = torch.Generator().manual_seed(9000 + step)
+    x = torch.randn(2, 16, generator=generator).to(engine.device)
+    target = torch.randn(2, 16, generator=generator).to(engine.device)
+    return x, target
+
+
+def _affine_resume_full_tensor(name, value, tp_size):
+    if tp_size == 1:
+        return value.detach().cpu().clone()
+    shards = [torch.empty_like(value) for _ in range(tp_size)]
+    dist.all_gather(shards, value.contiguous(), group=groups.get_tensor_model_parallel_group())
+    if name == "fc2.bias":
+        for shard in shards[1:]:
+            torch.testing.assert_close(shard, shards[0])
+        full = shards[0]
+    else:
+        full = torch.cat(shards, dim=1 if name == "fc2.weight" else 0)
+    return full.detach().cpu().clone()
+
+
+def _affine_resume_state(engine, tp_size, gradients=False):
+    from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
+
+    result = {}
+    for name, param in engine.module.named_parameters():
+        if gradients:
+            values = {"grad": safe_get_full_grad(param)}
+        else:
+            values = {
+                "fp32": safe_get_full_fp32_param(param),
+                "exp_avg": safe_get_full_optimizer_state(param, "exp_avg"),
+                "exp_avg_sq": safe_get_full_optimizer_state(param, "exp_avg_sq"),
+            }
+        for key, value in values.items():
+            assert value is not None, (name, key)
+            result[f"{name}/{key}"] = _affine_resume_full_tensor(name, value, tp_size)
+    result["global_steps"] = engine.global_steps
+    if not gradients:
+        steps = [state["step"].item() for state in engine.optimizer.optimizer.state.values()]
+        assert steps and all(step == engine.global_steps for step in steps)
+        result["optimizer_step"] = steps[0]
+    return result
+
+
+def _affine_resume_step(engine, tp_size, step):
+    x, target = _affine_resume_batch(engine, step)
+    logits = engine(x)
+    loss = torch.nn.functional.mse_loss(logits, target)
+    engine.backward(loss)
+    result = _affine_resume_state(engine, tp_size, gradients=True)
+    result["logits"] = logits.detach().cpu().clone()
+    result["loss"] = loss.detach().cpu().clone()
+    engine.step()
+    result.update(_affine_resume_state(engine, tp_size))
+    return result
+
+
+class affine_resume_checkpoint(DistributedFixture):
+    world_size = 2
+
+    def run(self, tmpdir, affine_layout):
+        from deepspeed.checkpoint.affine import (AFFINE_MAP_FORMAT_VERSION, contiguous_split_map, replicated_map)
+        from deepspeed.checkpoint.constants import AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION
+
+        engine = _affine_resume_engine(self.world_size)
+        for step in range(4):
+            _affine_resume_step(engine, self.world_size, step)
+
+        # Phase 1 does not yet emit AutoTP affine metadata. Supply the known fixture
+        # layout through the checkpoint API; all weights and moments come from training.
+        # LinearAllreduce adds this bias AFTER reduction, so its real scale is one.
+        maps = {
+            r"^fc1\.weight$": contiguous_split_map((16, 16), [8, 8], 0),
+            r"^fc1\.bias$": contiguous_split_map((16, ), [8, 8], 0),
+            r"^fc2\.weight$": contiguous_split_map((16, 16), [8, 8], 1),
+            r"^fc2\.bias$": replicated_map((16, ), self.world_size),
+        }
+        uc_info = {
+            UNIVERSAL_CHECKPOINT_VERSION_KEY: UNIVERSAL_CHECKPOINT_VERSION_VALUE,
+            AFFINE_MAP: {
+                AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION,
+                AFFINE_MAP_PARAMS: {
+                    name: layout.to_dict()
+                    for name, layout in maps.items()
+                },
+            },
+        }
+        engine.save_checkpoint(tmpdir,
+                               tag="affine_resume",
+                               client_state={UNIVERSAL_CHECKPOINT_INFO: uc_info} if affine_layout else {})
+        dist.barrier()
+        if dist.get_rank() == 0:
+            _convert_to_universal(os.path.join(tmpdir, "affine_resume"), os.path.join(tmpdir, "affine_universal"))
+        dist.barrier()
+        reference = [_affine_resume_state(engine, self.world_size)]
+        for step in range(4, 8):
+            reference.append(_affine_resume_step(engine, self.world_size, step))
+        if dist.get_rank() == 0:
+            torch.save(reference, os.path.join(tmpdir, "affine_reference.pt"))
+        dist.barrier()
+        engine.destroy()
+
+
+@pytest.mark.parametrize("affine_layout", [False, True], ids=["legacy", "affine"])
+@pytest.mark.parametrize("world_size", [1, 2], ids=["tp1", "tp2"])
+class TestAffineUniversalCheckpointResume(DistributedTest):
+
+    def test_resume_matches_uninterrupted_training(self, affine_resume_checkpoint, tmpdir, affine_layout, world_size):
+        tp_size = dist.get_world_size()
+        reference = torch.load(os.path.join(tmpdir, "affine_reference.pt"), weights_only=False)
+        engine = _affine_resume_engine(tp_size, load_universal=True)
+        load_path, _ = engine.load_checkpoint(tmpdir, tag="affine_universal", load_optimizer_states=True)
+        assert load_path is not None
+        actual = [_affine_resume_state(engine, tp_size)]
+        for step in range(4, 8):
+            actual.append(_affine_resume_step(engine, tp_size, step))
+        # The uninterrupted job is independent of converter/loader geometry. Resetting
+        # moments, losing the row bias, or slicing along the wrong axis changes this trace.
+        for index, (restored, expected) in enumerate(zip(actual, reference)):
+            assert restored.keys() == expected.keys()
+            for name in expected:
+                torch.testing.assert_close(restored[name],
+                                           expected[name],
+                                           atol=2e-6,
+                                           rtol=2e-5,
+                                           msg=lambda message: f"snapshot {index}, {name}: {message}")
+        engine.destroy()