Fix rank split when a full optimizer state is missing

deepspeed/runtime/zero/stage3.py

@@ -18,6 +18,7 @@
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from deepspeed.runtime.base_optimizer import ZeROOptimizer
 from deepspeed.utils import logger
+from deepspeed.utils.tensor_fragment import _count_fragment_owners
 from deepspeed.utils.torch import register_grad_hook, required_torch_version
 from deepspeed.utils.pin_memory_tracker import pinned_memory_summary
 from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler
@@ -3190,6 +3191,18 @@ def get_full_hp_param(self, param, optim_state_key=None) -> Tensor:
         if not param.requires_grad:
             return None
 
+        if optim_state_key is not None:
+            group_idx = self.grad_position[self.get_param_id(param)][0]
+            fp32_param = self.fp32_partitioned_groups_flat[group_idx]
+            state = self.optimizer.state[fp32_param] if fp32_param in self.optimizer.state else None
+            ready = state is not None and torch.is_tensor(state.get(optim_state_key))
+            group = self._get_param_partition_group(param)
+            device = torch.device(get_accelerator().current_device_name())
+            # Count first so every rank either all-gathers or raises.
+            owners, ready_count = _count_fragment_owners(True, ready, group, device)
+            if ready_count != owners:
+                raise ValueError(f'{optim_state_key} not found in optimizer state fragment')
+
         fp32_opt_state, group_idx = self._get_fp32_opt_state_partition(param,
                                                                        release_swap_buffers=True,
                                                                        optim_state_key=optim_state_key)

deepspeed/utils/tensor_fragment.py

@@ -34,10 +34,9 @@ def update_lp(self):
         self.lp_fragment.data.copy_(self.hp_fragment.data)
 
     def get_optim_state_fragment(self, key):
-        if key in self.optim_fragment:
+        if self.optim_fragment is not None and key in self.optim_fragment:
             return self.optim_fragment[key]
-        else:
-            raise ValueError(f'{key} not found in optimizer state fragment')
+        raise ValueError(f'{key} not found in optimizer state fragment')
 
     def set_optim_state_fragment(self, flat_hp_partition, optim_fragment):
         self.optim_fragment = {
@@ -84,8 +83,42 @@ def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys):
         optim_state[hp_param][key] = buffer
 
 
+def _count_fragment_owners(owns_fragment, fragment_ready, group, device):
+    # Every rank has to see the same counts before the data collective.
+    if not owns_fragment:
+        fragment_ready = False
+    status = torch.tensor([1.0 if owns_fragment else 0.0, 1.0 if fragment_ready else 0.0],
+                          dtype=torch.float32,
+                          device=device)
+    dist.all_reduce(status, group=group)
+    return int(status[0]), int(status[1])
+
+
+def _optim_state_ready(mapping, optim_state_key):
+    if mapping is None or mapping.optim_fragment is None:
+        return False
+    return optim_state_key in mapping.optim_fragment
+
+
+def _grad_fragment_ready(mapping):
+    if mapping is None:
+        return False
+    if mapping.use_offload:
+        gradient_dict = mapping.offload_gradient_dict
+    else:
+        gradient_dict = mapping.gradient_dict
+    return gradient_dict.get(mapping.param_group_index) is not None
+
+
 def get_full_hp_param(self, optim_state_key=None):
     reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    if optim_state_key is not None:
+        owns_fragment = self._hp_mapping is not None
+        ready = _optim_state_ready(self._hp_mapping, optim_state_key)
+        owners, ready_count = _count_fragment_owners(owns_fragment, ready, self._dp_group, reduce_buffer.device)
+        if owners > 0 and ready_count != owners:
+            raise ValueError(f'{optim_state_key} not found in optimizer state fragment')
+
     if self._hp_mapping is not None:
         lp_frag_address = self._hp_mapping.lp_fragment_address
         reduce_fragment = torch.narrow(reduce_buffer, 0, lp_frag_address.start, lp_frag_address.numel)
@@ -105,6 +138,12 @@ def set_full_hp_param(self, value, optim_state_key=None):
 
 def get_full_hp_grad(self):
     reduce_buffer = torch.zeros_like(self, dtype=torch.float32).flatten()
+    owns_fragment = self._hp_mapping is not None
+    ready = _grad_fragment_ready(self._hp_mapping)
+    owners, ready_count = _count_fragment_owners(owns_fragment, ready, self._dp_group, reduce_buffer.device)
+    if owners > 0 and ready_count != owners:
+        raise ValueError("Gradients are only available immediately after backward and before engine step")
+
     if self._hp_mapping is not None:
         lp_grad_fragment = self._hp_mapping.get_lp_grad_fragment(self._index_in_param_group)
         hp_grad_fragment = lp_grad_fragment.to(torch.float32).flatten()

tests/unit/runtime/zero/test_zero_tensor_fragment.py

@@ -50,6 +50,9 @@ def test_full_gradient_reconstruction_uses_only_the_owned_fragment(monkeypatch,
         peer_contribution[peer_start:peer_start + 2] = reference[peer_start:peer_start + 2]
 
         def sum_peer_fragment(tensor, *, group, _peer=peer_contribution):
+            # The readiness count is a 2-element reduction and is already local here.
+            if tensor.shape != _peer.shape:
+                return
             tensor.add_(_peer)
 
         monkeypatch.setattr(dist, "all_reduce", sum_peer_fragment)
@@ -607,3 +610,133 @@ def after_step_validate_func(model):
             pass
 
         run_fragmented_model(model, config_dict, hidden_dim, dtype, after_bwd_validate_func, after_step_validate_func)
+
+
+class _WeightAndLora(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(4, 4))
+        # One element stays inside a single ZeRO-1/2 partition after padding.
+        self.lora = torch.nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        return (self.weight * x).sum() + self.lora.sum()
+
+
+def _fragment_config(stage):
+    config = {
+        "train_micro_batch_size_per_gpu": 1,
+        "steps_per_print": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-3,
+                "torch_adam": True,
+            },
+        },
+        "bf16": {
+            "enabled": True
+        },
+        "zero_optimization": {
+            "stage": stage,
+        },
+    }
+    if stage == 3:
+        config["zero_optimization"]["param_persistence_threshold"] = 0
+    return config
+
+
+def _owner_count(param):
+    # NCCL rejects a CPU tensor. current_device_name() is 'cpu' on the CPU accelerator.
+    device = get_accelerator().current_device_name()
+    owned = torch.zeros(1, dtype=torch.float32, device=device)
+    if getattr(param, "_hp_mapping", None) is not None:
+        owned[0] = 1
+    dist.all_reduce(owned)
+    return int(owned.item())
+
+
+def _assert_same_across_ranks(value):
+    reference = value.detach().clone()
+    dist.broadcast(reference, src=0)
+    torch.testing.assert_close(value, reference)
+
+
+class TestMissingFragmentAgrees(DistributedTest):
+    world_size = 2
+
+    def test_zero2_missing_state_and_grad_raise_on_every_rank(self):
+        if not get_accelerator().is_bf16_supported():
+            pytest.skip("requires bf16")
+
+        model = _WeightAndLora()
+        # One group so the step links state for every parameter. The 1-element
+        # parameter still sits on a single partition.
+        param_groups = [{"params": [model.weight, model.lora]}]
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=param_groups, config=_fragment_config(2))
+        try:
+            weight = engine.module.weight
+            lora = engine.module.lora
+            assert _owner_count(lora) == 1
+            assert _owner_count(weight) == dist.get_world_size()
+
+            with pytest.raises(ValueError, match="exp_avg"):
+                safe_get_full_optimizer_state(lora, "exp_avg")
+            with pytest.raises(ValueError, match="exp_avg"):
+                safe_get_full_optimizer_state(weight, "exp_avg")
+            with pytest.raises(ValueError, match="Gradients are only available"):
+                safe_get_full_grad(lora)
+
+            x = torch.ones(4, 4, dtype=weight.dtype, device=engine.device)
+            engine.backward(engine(x))
+            _assert_same_across_ranks(safe_get_full_grad(lora))
+            _assert_same_across_ranks(safe_get_full_grad(weight))
+
+            engine.step()
+            lora_state = safe_get_full_optimizer_state(lora, "exp_avg")
+            weight_state = safe_get_full_optimizer_state(weight, "exp_avg")
+            assert lora_state.shape == lora.shape
+            assert weight_state.shape == weight.shape
+            _assert_same_across_ranks(lora_state)
+            _assert_same_across_ranks(weight_state)
+
+            with pytest.raises(ValueError, match="Gradients are only available"):
+                safe_get_full_grad(lora)
+
+            if dist.get_rank() == 0:
+                weight._hp_mapping.optim_fragment.pop("exp_avg")
+            with pytest.raises(ValueError, match="exp_avg"):
+                safe_get_full_optimizer_state(weight, "exp_avg")
+        finally:
+            engine.destroy()
+
+    def test_zero3_missing_state_raises_on_every_rank(self):
+        if not get_accelerator().is_bf16_supported():
+            pytest.skip("requires bf16")
+
+        model = _WeightAndLora()
+        param_groups = [{"params": [model.weight]}, {"params": [model.lora]}]
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=param_groups, config=_fragment_config(3))
+        try:
+            weight = engine.module.weight
+            with pytest.raises(ValueError, match="exp_avg"):
+                safe_get_full_optimizer_state(weight, "exp_avg")
+
+            x = torch.ones(4, 4, dtype=weight.dtype, device=engine.device)
+            engine.backward(engine(x))
+            engine.step()
+            state = safe_get_full_optimizer_state(weight, "exp_avg")
+            assert state is not None
+            assert tuple(state.shape) == tuple(getattr(weight, "ds_shape", weight.shape))
+            _assert_same_across_ranks(state)
+
+            zero = engine.optimizer
+            group_idx = zero.grad_position[zero.get_param_id(weight)][0]
+            fp32_param = zero.fp32_partitioned_groups_flat[group_idx]
+            if dist.get_rank() == 0:
+                del zero.optimizer.state[fp32_param]["exp_avg"]
+            with pytest.raises(ValueError, match="exp_avg"):
+                safe_get_full_optimizer_state(weight, "exp_avg")
+        finally:
+            engine.destroy()