From 0ea262c4f12841a51dc71edd16a4435ce0a81ff1 Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Sun, 19 Jul 2026 23:39:20 +0000 Subject: [PATCH 1/2] Fix ZeRO-3 all_reduce param fetch stride for padded parameters In the all_reduce reconstruction path, the flat buffer is sized sum(ds_numel_aligned) but the per-parameter offset advanced by the unaligned ds_numel. When a parameter's numel is not divisible by the partition world size the two differ, so each subsequent parameter's span started inside its predecessor's aligned span. The overlapping slots are the predecessor's partition padding, which is never written and therefore holds uninitialized memory, and the all_reduce SUM adds it into the next parameter's leading elements. Advance by ds_numel_aligned so each parameter owns the disjoint span the buffer sizing already reserves for it. Signed-off-by: Ehsan Barkhordar --- .../runtime/zero/partition_parameters.py | 2 +- .../zero/test_zero_allreduce_fetch_params.py | 71 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index 827e07341c3c..8f9b571c9177 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -1364,7 +1364,7 @@ def _all_gather_coalesced(params, world_size, rank_in_group, use_secondary_tenso start = start_param + param.ds_tensor.ds_numel * rank_in_group flat_tensor.narrow(0, start, param.ds_tensor.ds_numel).copy_(param.ds_tensor) - start_param += param.ds_numel + start_param += param.ds_numel_aligned handle = dist.all_reduce(flat_tensor, group=ds_process_group, async_op=True) diff --git a/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py b/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py new file mode 100644 index 000000000000..fd239f321309 --- /dev/null +++ b/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py @@ -0,0 +1,71 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import pytest +import torch + +import deepspeed +import deepspeed.runtime.zero.partition_parameters as partition_parameters +from unit.common import DistributedTest + +# Odd numels need a slot of partition padding under world_size=2, so ds_numel_aligned +# differs from ds_numel. Even numels make the two equal, which is the case the flat-buffer +# stride already handled; both are checked so the padded fix does not regress it. +PADDED_NUMELS = [5, 7] +ALIGNED_NUMELS = [4, 6] +POISON = 7777.0 + + +class ParamHolder(torch.nn.Module): + + def __init__(self, numels): + super().__init__() + for i, numel in enumerate(numels): + self.register_parameter(f"p{i}", torch.nn.Parameter(torch.arange(1, numel + 1, dtype=torch.float32))) + + +class TestAllReduceFetchParamsPadded(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize("numels", [PADDED_NUMELS, ALIGNED_NUMELS], ids=["padded", "aligned"]) + def test_params_reconstruct_exactly(self, numels): + config = { + "train_micro_batch_size_per_gpu": 1, + "zero_optimization": { + "stage": 3, + "stage3_use_all_reduce_for_fetch_params": True, + }, + } + + # The tail of the last rank's partition is never written (torch.empty, and the + # partial-copy branch only copies the elements that exist), so its contents are + # whatever the allocator returns. Fill new allocations with a sentinel so the test + # is deterministic instead of depending on whether the allocator hands back a + # freshly zeroed page. + real_empty = partition_parameters._orig_torch_empty + + def poisoned_empty(*args, **kwargs): + tensor = real_empty(*args, **kwargs) + if tensor.dtype.is_floating_point: + tensor.fill_(POISON) + return tensor + + expected = [torch.arange(1, numel + 1, dtype=torch.float32) for numel in numels] + + partition_parameters._orig_torch_empty = poisoned_empty + try: + with deepspeed.zero.Init(config_dict_or_path=config, mem_efficient_linear=False, enabled=True): + module = ParamHolder(numels) + finally: + partition_parameters._orig_torch_empty = real_empty + + params = [getattr(module, f"p{i}") for i in range(len(numels))] + params[0].all_gather_coalesced(params).wait() + + for i, param in enumerate(params): + gathered = param.data.detach().reshape(-1).cpu() + assert torch.equal( + gathered, expected[i]), (f"param p{i} was not reconstructed exactly: expected {expected[i].tolist()}, " + f"got {gathered.tolist()}") From 409ccd71328389a39940ae9f0a36ac67b977eea0 Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Sun, 19 Jul 2026 23:55:18 +0000 Subject: [PATCH 2/2] Restore the public torch.empty binding after the poisoned-allocator block Init.__exit__ restores torch.empty by reading _orig_torch_empty, so leaving the test's poisoned function in that global rebound torch.empty to it on the way out and left every later allocation in the worker process sentinel-filled. Signed-off-by: Ehsan Barkhordar --- .../unit/runtime/zero/test_zero_allreduce_fetch_params.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py b/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py index fd239f321309..6827db7d4153 100644 --- a/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py +++ b/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py @@ -54,12 +54,19 @@ def poisoned_empty(*args, **kwargs): expected = [torch.arange(1, numel + 1, dtype=torch.float32) for numel in numels] + # Init.__exit__ restores torch.empty by reading _orig_torch_empty, so on the way out + # it rebinds the public torch.empty to poisoned_empty. Restoring the module global + # alone would leave every later allocation in this worker process sentinel-filled, + # so put the public binding back too. + saved_torch_empty = torch.empty + partition_parameters._orig_torch_empty = poisoned_empty try: with deepspeed.zero.Init(config_dict_or_path=config, mem_efficient_linear=False, enabled=True): module = ParamHolder(numels) finally: partition_parameters._orig_torch_empty = real_empty + torch.empty = saved_torch_empty params = [getattr(module, f"p{i}") for i in range(len(numels))] params[0].all_gather_coalesced(params).wait()