From 146f918068770d8e1d01590344672a8e0e95276c Mon Sep 17 00:00:00 2001 From: Shreyash Bhatkar Date: Tue, 10 Nov 2026 19:25:39 +0000 Subject: [PATCH 1/2] Fix rank split when a full optimizer state is missing safe_get_full_optimizer_state() raised only on ranks that owned a fragment, and the other ranks went into all_reduce. Count whether the fragment is readable first, and raise the same error on every rank. safe_get_full_grad() and the ZeRO-3 full optimizer-state read had the same split. Signed-off-by: Shreyash Bhatkar Co-authored-by: Shreyash Bhatkar --- deepspeed/runtime/zero/stage3.py | 13 ++ deepspeed/utils/tensor_fragment.py | 45 +++++- .../runtime/zero/test_zero_tensor_fragment.py | 131 ++++++++++++++++++ 3 files changed, 186 insertions(+), 3 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index f3b972601058..4950579729ff 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/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) diff --git a/deepspeed/utils/tensor_fragment.py b/deepspeed/utils/tensor_fragment.py index 296517ea44a2..a433988fec1d 100644 --- a/deepspeed/utils/tensor_fragment.py +++ b/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() diff --git a/tests/unit/runtime/zero/test_zero_tensor_fragment.py b/tests/unit/runtime/zero/test_zero_tensor_fragment.py index 83c42cdad3bf..d3e8d1405126 100644 --- a/tests/unit/runtime/zero/test_zero_tensor_fragment.py +++ b/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,131 @@ 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): + owned = torch.zeros(1, dtype=torch.float32) + 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() From 57ea7a132f8cb787aba6d8e3d1b9a7066cdb025b Mon Sep 17 00:00:00 2001 From: Shreyash Bhatkar Date: Wed, 11 Nov 2026 05:39:05 +0000 Subject: [PATCH 2/2] Put test owner count on accelerator device Signed-off-by: Shreyash Bhatkar Co-authored-by: Shreyash Bhatkar --- tests/unit/runtime/zero/test_zero_tensor_fragment.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/runtime/zero/test_zero_tensor_fragment.py b/tests/unit/runtime/zero/test_zero_tensor_fragment.py index d3e8d1405126..0cbfbf579060 100644 --- a/tests/unit/runtime/zero/test_zero_tensor_fragment.py +++ b/tests/unit/runtime/zero/test_zero_tensor_fragment.py @@ -648,7 +648,9 @@ def _fragment_config(stage): def _owner_count(param): - owned = torch.zeros(1, dtype=torch.float32) + # 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)