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 68dd3e654d72..f47ca7796648 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) @@ -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): 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..26fc88b75318 --- /dev/null +++ b/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