From 50da5546541b41fff8b2470b43c4018a36e5f731 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Mon, 31 Aug 2026 22:37:13 +0800 Subject: [PATCH 1/3] Gather zero-sized parameters in ZeRO-3 fetch_sub_module decides whether to run the all-gather from the number of elements still to fetch: fetch_numel = sum(p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE) if fetch_numel > 0: A submodule whose only parameter is zero-sized contributes nothing to that sum, so the gather is skipped, the parameter never leaves NOT_AVAILABLE, and the wait loop right below asserts that it is AVAILABLE: AssertionError: {'id': 1, 'status': 'NOT_AVAILABLE', 'numel': 0, ...} Gate on whether any parameter still needs gathering instead. The element count is kept for the profiler, which is what it was for. This is the ZeRO-3 counterpart of #8280 and #8298, which fixed the same class of failure for stages 1 and 2 (issues #8279 and #8297). Signed-off-by: alanhuangyoo --- .../zero/partitioned_param_coordinator.py | 12 +-- .../runtime/zero/test_zero_empty_param.py | 79 +++++++++++++++++++ 2 files changed, 86 insertions(+), 5 deletions(-) create mode 100644 tests/unit/runtime/zero/test_zero_empty_param.py diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py index 68dd3e654d72..abb18b1c4349 100644 --- a/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -362,14 +362,16 @@ def _fetch_sub_module_impl(self, current_submodule: Module, forward: bool, is_le })) params_to_fetch = set(iter_params(current_submodule, recurse=is_leaf)) - fetch_numel = sum( - [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + # Gate on whether anything still has to be gathered, not on how many elements that is: + # a zero-sized parameter contributes 0 to the sum but still has to leave NOT_AVAILABLE, + # and the wait loop below asserts that it did. + params_to_gather = [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE] + fetch_numel = sum(p.partition_numel() for p in params_to_gather) - if fetch_numel > 0: + if params_to_gather: event_name = __class__.FORWARD_FETCH_SUBMIT if forward else __class__.BACKWARD_FETCH_SUBMIT self._dump_param_ids(event_name, current_submodule.ds_id, - [(p.ds_id, p.ds_shape) - for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) + [(p.ds_id, p.ds_shape) for p in params_to_gather]) # self._dump_params(event_name, current_submodule, [p for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) self.__profiler.start_event(event_name) diff --git a/tests/unit/runtime/zero/test_zero_empty_param.py b/tests/unit/runtime/zero/test_zero_empty_param.py new file mode 100644 index 000000000000..cfbf0be7c78c --- /dev/null +++ b/tests/unit/runtime/zero/test_zero_empty_param.py @@ -0,0 +1,79 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""A zero-sized trainable parameter must survive a full ZeRO step on every stage. + +Stages 1 and 2 were fixed in #8280 and #8298 (issues #8279, #8297). Stage 3 took the same +shape of failure from a different place: `fetch_sub_module` gates the all-gather on +`fetch_numel > 0`, and a submodule holding only a zero-sized parameter contributes nothing to +that sum, so the gather never runs, the parameter stays `NOT_AVAILABLE`, and the wait loop +immediately below asserts that it is `AVAILABLE`. +""" + +import pytest +import torch + +from unit.common import DistributedTest + +import deepspeed + + +class EmptyTailModel(torch.nn.Module): + """The shape from issue #8279: a trainable parameter with no elements, used in the loss.""" + + def __init__(self, hidden=8): + super().__init__() + self.dense = torch.nn.Linear(hidden, hidden, bias=False) + self.empty = torch.nn.Linear(hidden, 0, bias=False) + + def forward(self, x): + hidden = self.dense(x) + # `empty(hidden)` is (batch, 0); summing it keeps the parameter in the autograd graph. + return hidden.sum() + self.empty(hidden).sum() + + +def _run_one_step(stage, hidden=8): + config = { + "train_micro_batch_size_per_gpu": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-3 + } + }, + "zero_optimization": { + "stage": stage + }, + "fp16": { + "enabled": False + }, + } + model = EmptyTailModel(hidden) + engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + + loss = engine(torch.randn(1, hidden, device=engine.device, dtype=next(engine.parameters()).dtype)) + engine.backward(loss) + engine.step() + + return engine + + +class TestZeroSizedParameterSingleRank(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("stage", [1, 2, 3]) + def test_step_completes(self, stage): + engine = _run_one_step(stage) + assert engine.global_steps == 1 + + +class TestZeroSizedParameterPartitioned(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize("stage", [1, 2, 3]) + def test_step_completes(self, stage): + # With more than one rank the stage-3 path goes through the real all-gather rather than + # the single-rank shortcut, which is where the gate lives. + engine = _run_one_step(stage) + assert engine.global_steps == 1 From 5cda14e352329625aa7f42859e063a608993e4bd Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 9 Sep 2026 22:11:52 +0800 Subject: [PATCH 2/3] Scope this to ZeRO-3; stages 1 and 2 fail from a different place Running the parametrized test showed stages 1 and 2 still failing on this model shape with #8280 and #8298 in master: RuntimeError: size mismatch, got input (1), mat (1x8), vec (0) That is not the gather gate this PR is about. _update_model_bit16_weights drops a zero-element parameter's shape when it repoints it at the ZeRO-1/2 flat buffer, so (0, 8) comes back as (0,) and the module's own forward breaks before any of this code runs. Fixing it belongs in its own change. Narrow the parametrize to stage 3 and correct the docstring, which claimed stages 1 and 2 already worked. Signed-off-by: alanhuangyoo --- .../runtime/zero/test_zero_empty_param.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/unit/runtime/zero/test_zero_empty_param.py b/tests/unit/runtime/zero/test_zero_empty_param.py index cfbf0be7c78c..571f19dbc0d0 100644 --- a/tests/unit/runtime/zero/test_zero_empty_param.py +++ b/tests/unit/runtime/zero/test_zero_empty_param.py @@ -2,13 +2,16 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team -"""A zero-sized trainable parameter must survive a full ZeRO step on every stage. +"""A zero-sized trainable parameter must survive a full ZeRO-3 step. -Stages 1 and 2 were fixed in #8280 and #8298 (issues #8279, #8297). Stage 3 took the same -shape of failure from a different place: `fetch_sub_module` gates the all-gather on -`fetch_numel > 0`, and a submodule holding only a zero-sized parameter contributes nothing to -that sum, so the gather never runs, the parameter stays `NOT_AVAILABLE`, and the wait loop -immediately below asserts that it is `AVAILABLE`. +`fetch_sub_module` gates the all-gather on `fetch_numel > 0`, and a submodule holding only a +zero-sized parameter contributes nothing to that sum, so the gather never runs, the parameter +stays `NOT_AVAILABLE`, and the wait loop immediately below asserts that it is `AVAILABLE`. + +Stages 1 and 2 are deliberately not covered here. #8280 and #8298 (issues #8279, #8297) fixed +the reduction path they were reported against, but this model shape still fails on both from a +different place — `_update_model_bit16_weights` drops a zero-element parameter's shape when it +repoints it at the flat buffer — which is a separate fix. """ import pytest @@ -62,7 +65,7 @@ def _run_one_step(stage, hidden=8): class TestZeroSizedParameterSingleRank(DistributedTest): world_size = 1 - @pytest.mark.parametrize("stage", [1, 2, 3]) + @pytest.mark.parametrize("stage", [3]) def test_step_completes(self, stage): engine = _run_one_step(stage) assert engine.global_steps == 1 @@ -71,7 +74,7 @@ def test_step_completes(self, stage): class TestZeroSizedParameterPartitioned(DistributedTest): world_size = 2 - @pytest.mark.parametrize("stage", [1, 2, 3]) + @pytest.mark.parametrize("stage", [3]) def test_step_completes(self, stage): # With more than one rank the stage-3 path goes through the real all-gather rather than # the single-rank shortcut, which is where the gate lives. From 889e4b515dd1760325616a0b4589894a41c9d8f3 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Thu, 10 Sep 2026 01:57:53 +0800 Subject: [PATCH 3/3] Address review: coalesced gather, quantizer, prefetch gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three more places treat "has elements" and "needs handling" as the same question. All three reproduce on unmodified master, so they are pre-existing rather than introduced here — but they sit on the path this PR opens up, and a zero-sized parameter cannot get through ZeRO-3 without them. 1. `AllGatherCoalescedHandle.wait` concatenates the per-rank slices of each parameter. No rank holds a slice of a zero-sized one, so the list is empty: ValueError: torch.cat(): expected a non-empty list of Tensors Give the parameter an empty tensor of its own shape instead. The dtype has to be the one the sized parameters of the bucket end up with, which is the partition's dtype normally and the parameter's own under quantization, where the partition is stored as int8. 2. `CUDAQuantizer.quantize` derives its group count from numel, so a zero-sized partition makes `groups` 0 and the next `numel % (8 * groups * 2)` raises ZeroDivisionError. Return the empty (int8, fp32 scale) pair the callers expect, so `ds_quant_scale` is still set, with the mirror guard in `dequantize`. 3. The prefetch submit gates on `numel_prefetching > 0`, the same proxy this PR removed from the fetch. Gate on the parameter set instead. Tests: coalesced gather (shape and dtype of the parameter it hands back), `zero_quantized_weights`, `zero_quantized_nontrainable_weights` (both skipped when QuantizerBuilder is unavailable), and a prefetch-enabled model. 6 passed on this branch, 6 failed on master — ValueError for the coalesced case, ZeroDivisionError for the quantized non-trainable one, and the NOT_AVAILABLE assertion for the rest. The prefetch change is the one exception: reverting only that line still leaves all six passing, and tracing `__all_gather_params` shows the same 8 all-zero submissions either way, because the fetch gathers those parameters when their submodule is reached. It is a consistency fix with no case I could construct where it changes behaviour, so the test above does not claim to cover it. Signed-off-by: alanhuangyoo --- .../runtime/zero/partition_parameters.py | 25 ++- .../zero/partitioned_param_coordinator.py | 7 +- .../runtime/zero/test_zero_empty_param.py | 148 +++++++++++++++--- 3 files changed, 156 insertions(+), 24 deletions(-) diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index d91023afc500..3e82c30a2b5f 100644 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -780,8 +780,18 @@ def wait(self, handle_dependency=True) -> None: part_to_copy = self.partitions[rank].narrow(0, param_offset, min(param.ds_numel - param_start, ds_tensor_numel)) partitions.append(part_to_copy) - # Note that dtypes of param and partitions can be different (currently for torch.autocast support) - param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape).to(param.ds_tensor.dtype) + if not partitions: + # No rank holds a slice of a zero-element parameter, so the loop above + # appended nothing and there is nothing to concatenate. The gather still has + # to leave the parameter AVAILABLE in its own shape, and at the dtype the + # sized parameters of this bucket end up with: the partition's dtype + # normally, and the parameter's own under quantization, where the partition + # is stored as int8 and the sized params are dequantized on the way out. + empty_dtype = param.dtype if self.quantization else param.ds_tensor.dtype + param.data = torch.empty(param.ds_shape, dtype=empty_dtype, device=param.device) + else: + # Note that dtypes of param and partitions can be different (currently for torch.autocast support) + param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape).to(param.ds_tensor.dtype) param.ds_status = ZeroParamStatus.AVAILABLE if not get_accelerator().is_synchronized_device() and handle_dependency: for part_to_copy in partitions: @@ -854,6 +864,14 @@ def __init__(self) -> None: CUDAQuantizer.quantizer_cuda_module = deepspeed.ops.op_builder.QuantizerBuilder().load() def quantize(self, param, groups=None): + if param.numel() == 0: + # Nothing to quantize, and the group-size search below derives its divisor from + # numel: `groups` comes out 0 and the very next `numel % (8 * groups * 2)` raises + # ZeroDivisionError. Return the empty pair the callers expect, so `ds_quant_scale` + # is still set on a zero-element partition and dequantize round-trips it. + device = get_accelerator().device_name() + return (torch.empty(0, dtype=torch.int8, + device=device), torch.empty((0, 1), dtype=torch.float32, device=device)) if groups is None: try: groups = self.group_size_cache[param.numel()] @@ -883,6 +901,9 @@ def quantize(self, param, groups=None): return self.quantizer_cuda_module.quantize(param, groups, 8, self.quantizer_cuda_module.Symmetric) def dequantize(self, quantized_param, scale, dtype=None): + if quantized_param.numel() == 0: + # Mirror of the guard in quantize(): the kernel is given a zero group count here. + return torch.empty(0, dtype=dtype or torch.half, device=quantized_param.device) dequantized = self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8, self.quantizer_cuda_module.Symmetric) if dtype is not None and dequantized.dtype != dtype: diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py index abb18b1c4349..f47ca7796648 100644 --- a/deepspeed/runtime/zero/partitioned_param_coordinator.py +++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py @@ -477,7 +477,12 @@ def _is_currently_on_nvme(param): params_to_prefetch.add(param_in_trace.param) numel_prefetching += param_in_trace.param.ds_numel - if numel_prefetching > 0: + # Same reason as the fetch gate above: the element count is a proxy for + # "is there anything to gather". A set holding only zero-element parameters + # is popped off the queue and then never submitted; the fetch still gathers + # those parameters when their submodule is reached, so what the count costs + # here is the prefetch overlap rather than the step. + if params_to_prefetch: event_name = __class__.FORWARD_PREFETCH_SUBMIT if forward else __class__.BACKWARD_PREFETCH_SUBMIT self.__profiler.start_event(event_name) if logger.isEnabledFor(logging.DEBUG): diff --git a/tests/unit/runtime/zero/test_zero_empty_param.py b/tests/unit/runtime/zero/test_zero_empty_param.py index 571f19dbc0d0..26fc88b75318 100644 --- a/tests/unit/runtime/zero/test_zero_empty_param.py +++ b/tests/unit/runtime/zero/test_zero_empty_param.py @@ -4,14 +4,23 @@ # DeepSpeed Team """A zero-sized trainable parameter must survive a full ZeRO-3 step. -`fetch_sub_module` gates the all-gather on `fetch_numel > 0`, and a submodule holding only a -zero-sized parameter contributes nothing to that sum, so the gather never runs, the parameter -stays `NOT_AVAILABLE`, and the wait loop immediately below asserts that it is `AVAILABLE`. - -Stages 1 and 2 are deliberately not covered here. #8280 and #8298 (issues #8279, #8297) fixed -the reduction path they were reported against, but this model shape still fails on both from a -different place — `_update_model_bit16_weights` drops a zero-element parameter's shape when it -repoints it at the flat buffer — which is a separate fix. +Four places treat "has elements" and "needs handling" as the same question: + +* ``fetch_sub_module`` gates the all-gather on ``fetch_numel > 0``, so a submodule holding + only zero-sized parameters is skipped, the parameters stay ``NOT_AVAILABLE``, and the + wait loop immediately below asserts that they are ``AVAILABLE``; +* the prefetch submit gates on ``numel_prefetching > 0``, so an all-zero prefetch set is + popped off the queue and then never prefetched (the fetch above still gathers those + parameters when the submodule is reached, so this one costs the overlap, not the step); +* ``AllGatherCoalescedHandle.wait`` concatenates the per-rank slices, and no rank holds a + slice of a zero-sized parameter, so ``torch.cat`` is handed an empty list; +* ``CUDAQuantizer.quantize`` derives its group count from ``numel``, so a zero-sized + partition divides by zero. + +Stages 1 and 2 are deliberately not covered here. #8280 and #8298 (issues #8279, #8297) +fixed the reduction path they were reported against, but this model shape still fails on +both from a different place — ``_update_model_bit16_weights`` drops a zero-element +parameter's shape when it repoints it at the flat buffer — which is a separate fix. """ import pytest @@ -21,11 +30,13 @@ import deepspeed +HIDDEN = 8 + class EmptyTailModel(torch.nn.Module): """The shape from issue #8279: a trainable parameter with no elements, used in the loss.""" - def __init__(self, hidden=8): + def __init__(self, hidden=HIDDEN): super().__init__() self.dense = torch.nn.Linear(hidden, hidden, bias=False) self.empty = torch.nn.Linear(hidden, 0, bias=False) @@ -36,8 +47,55 @@ def forward(self, x): return hidden.sum() + self.empty(hidden).sum() -def _run_one_step(stage, hidden=8): - config = { +class SharedModuleBlock(torch.nn.Module): + """A sized and a zero-sized parameter on one module, so they are gathered together.""" + + def __init__(self, hidden=HIDDEN): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(hidden, hidden)) + self.empty = torch.nn.Parameter(torch.empty(0, hidden)) + + def forward(self, x): + return (x @ self.weight.t()).sum() + (x @ self.empty.t()).sum() + + +class OnlyEmptyBlock(torch.nn.Module): + """A submodule whose parameters are *all* zero-sized.""" + + def __init__(self, hidden=HIDDEN): + super().__init__() + self.first = torch.nn.Parameter(torch.empty(0, hidden)) + self.second = torch.nn.Parameter(torch.empty(0, hidden)) + + def forward(self, x): + return (x @ self.first.t()).sum() + (x @ self.second.t()).sum() + + +class SharedModuleModel(torch.nn.Module): + """Several blocks, each owning a sized and a zero-sized parameter.""" + + def __init__(self, hidden=HIDDEN, blocks=2): + super().__init__() + self.blocks = torch.nn.ModuleList([SharedModuleBlock(hidden) for _ in range(blocks)]) + + def forward(self, x): + return sum(block(x) for block in self.blocks) + + +class MixedModel(torch.nn.Module): + """A sized submodule next to one holding only zero-sized parameters.""" + + def __init__(self, hidden=HIDDEN, blocks=4): + super().__init__() + self.dense = torch.nn.Linear(hidden, hidden, bias=False) + self.only_empty = torch.nn.ModuleList([OnlyEmptyBlock(hidden) for _ in range(blocks)]) + + def forward(self, x): + return self.dense(x).sum() + sum(block(x) for block in self.only_empty) + + +def _config(stage=3, **zero_overrides): + return { "train_micro_batch_size_per_gpu": 1, "optimizer": { "type": "Adam", @@ -46,19 +104,22 @@ def _run_one_step(stage, hidden=8): } }, "zero_optimization": { - "stage": stage + "stage": stage, + **zero_overrides }, - "fp16": { - "enabled": False + "bf16": { + "enabled": True }, } - model = EmptyTailModel(hidden) - engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) - loss = engine(torch.randn(1, hidden, device=engine.device, dtype=next(engine.parameters()).dtype)) - engine.backward(loss) - engine.step() +def _run_steps(model, config, steps=1): + trainable = [p for p in model.parameters() if p.requires_grad] + engine, *_ = deepspeed.initialize(model=model, model_parameters=trainable, config=config) + for _ in range(steps): + loss = engine(torch.randn(1, HIDDEN, device=engine.device, dtype=torch.bfloat16)) + engine.backward(loss) + engine.step() return engine @@ -67,7 +128,7 @@ class TestZeroSizedParameterSingleRank(DistributedTest): @pytest.mark.parametrize("stage", [3]) def test_step_completes(self, stage): - engine = _run_one_step(stage) + engine = _run_steps(EmptyTailModel(), _config(stage)) assert engine.global_steps == 1 @@ -78,5 +139,50 @@ class TestZeroSizedParameterPartitioned(DistributedTest): def test_step_completes(self, stage): # With more than one rank the stage-3 path goes through the real all-gather rather than # the single-rank shortcut, which is where the gate lives. - engine = _run_one_step(stage) + engine = _run_steps(EmptyTailModel(), _config(stage)) + assert engine.global_steps == 1 + + def test_coalesced_gather_completes(self): + # One module owning both parameters puts them in a single coalesced gather, where + # the zero-sized one contributes no slice to concatenate. + engine = _run_steps(SharedModuleModel(), _config(), steps=2) + + assert engine.global_steps == 2 + block = engine.module.blocks[0] + assert block.empty.shape == torch.Size([0, HIDDEN]) + # The gather has to hand back a dtype the module's own forward can use. + assert block.empty.dtype == block.weight.dtype + + def test_completes_with_prefetch_enabled(self): + # Submodules holding only zero-sized parameters, with prefetch on. This covers the + # fetch gate under a prefetch configuration; it does not pin the prefetch submit + # gate itself, which costs the prefetch rather than the gather. + engine = _run_steps(MixedModel(), _config(stage3_prefetch_bucket_size=10000), steps=2) + assert engine.global_steps == 2 + + +class TestZeroSizedParameterQuantized(DistributedTest): + world_size = 2 + + def _skip_without_quantizer(self): + from deepspeed.ops.op_builder import QuantizerBuilder + if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]: + pytest.skip("QuantizerBuilder is not implemented") + + def test_quantized_weights(self): + self._skip_without_quantizer() + + engine = _run_steps(MixedModel(), _config(zero_quantized_weights=True)) + + assert engine.global_steps == 1 + + def test_quantized_nontrainable_weights(self): + self._skip_without_quantizer() + model = MixedModel() + for block in model.only_empty: + block.first.requires_grad = False + block.second.requires_grad = False + + engine = _run_steps(model, _config(zero_quantized_nontrainable_weights=True)) + assert engine.global_steps == 1