Gather zero-sized parameters in ZeRO-3

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:

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)
@@ -475,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):

tests/unit/runtime/zero/test_zero_empty_param.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""A zero-sized trainable parameter must survive a full ZeRO-3 step.
+
+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
+import torch
+
+from unit.common import DistributedTest
+
+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=HIDDEN):
+        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()
+
+
+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",
+            "params": {
+                "lr": 1e-3
+            }
+        },
+        "zero_optimization": {
+            "stage": stage,
+            **zero_overrides
+        },
+        "bf16": {
+            "enabled": True
+        },
+    }
+
+
+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
+
+
+class TestZeroSizedParameterSingleRank(DistributedTest):
+    world_size = 1
+
+    @pytest.mark.parametrize("stage", [3])
+    def test_step_completes(self, stage):
+        engine = _run_steps(EmptyTailModel(), _config(stage))
+        assert engine.global_steps == 1
+
+
+class TestZeroSizedParameterPartitioned(DistributedTest):
+    world_size = 2
+
+    @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.
+        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