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..6827db7d4153 --- /dev/null +++ b/tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py @@ -0,0 +1,78 @@ +# 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] + + # 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() + + 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()}")