From f9ec76324a2dbd6a891c2c44c258f2fbce511f66 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 9 Sep 2026 22:11:30 +0800 Subject: [PATCH 1/2] Keep a zero-element parameter's shape across the ZeRO-1/2 flat buffer `_update_model_bit16_weights` repoints every parameter at its slice of the flattened group: updated_params = self.unflatten(self.bit16_groups_flat[i], self.round_robin_bit16_meta[i]) for p, q in zip(self.round_robin_bit16_groups[i], updated_params): p.data = q.data torch's `unflatten_dense_tensors` special-cases a zero-element tensor and returns a freshly allocated 1-D `zeros({0})` instead of a view of the requested shape: >>> _unflatten_dense_tensors(_flatten_dense_tensors([a, b]), ... [torch.zeros_like(a, device="meta"), # (8, 8) ... torch.zeros_like(b, device="meta")]) # (0, 8) [(8, 8), (0,)] So a `nn.Linear(8, 0, bias=False)` weight came out of `deepspeed.initialize` with shape `(0,)` instead of `(0, 8)`, and the module's own forward then dispatched `F.linear` to `addmv`: RuntimeError: size mismatch, got input (1), mat (1x8), vec (0) The parameters are rebuilt from the flat buffer after every `step()` as well as at init, so restoring the shape once would not have held either. Skip the assignment for a zero-element parameter. There is no slice of the flat buffer for it to point at, and the tensor torch hands back is a fresh allocation rather than a view, so nothing is being kept in sync by the assignment. Stages 1 and 2 only. ZeRO-3 keeps the real shape in `ds_shape` and partitions to a 1-D local shard by design; its own zero-element failure is a different one. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zero/stage_1_and_2.py | 7 ++ .../zero/test_zero_numel_param_shape.py | 99 +++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 tests/unit/runtime/zero/test_zero_numel_param_shape.py diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index f05a53867c93..a5d1b762f5ae 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -805,6 +805,13 @@ def _configure_moe_settings(self): def _update_model_bit16_weights(self, group_index): updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index]) for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): + if p.numel() == 0: + # torch's unflatten_dense_tensors special-cases a zero-element tensor and + # hands back a freshly allocated 1-D `zeros({0})` instead of a view of the + # requested shape, so assigning it would replace e.g. a (0, 8) parameter + # with a (0,) one and break the module's own forward. There is nothing in + # the flat buffer to point such a parameter at anyway. + continue p.data = q.data # set model fp16 weight to slices of reordered flattened buffer diff --git a/tests/unit/runtime/zero/test_zero_numel_param_shape.py b/tests/unit/runtime/zero/test_zero_numel_param_shape.py new file mode 100644 index 000000000000..2eec4a469f4d --- /dev/null +++ b/tests/unit/runtime/zero/test_zero_numel_param_shape.py @@ -0,0 +1,99 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""A zero-element parameter must keep its shape across the ZeRO-1/2 flat buffer. + +`_update_model_bit16_weights` repoints every parameter at its slice of the flattened +group. torch's `unflatten_dense_tensors` special-cases a zero-element tensor and returns +a freshly allocated 1-D `zeros({0})` rather than a view of the requested shape, so a +`(0, 8)` parameter came back as `(0,)` and the module's own forward then dispatched +`F.linear` to `addmv`: + + RuntimeError: size mismatch, got input (1), mat (1x8), vec (0) + +The parameter is rebuilt on every `step()` as well as at init, so the shape did not +survive one iteration either. +""" + +import pytest +import torch + +from unit.common import DistributedTest + +import deepspeed + +HIDDEN = 8 + + +class EmptyTailModel(torch.nn.Module): + """A trainable parameter with no elements, kept in the autograd graph by 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 graph. + return hidden.sum() + self.empty(hidden).sum() + + +def _engine(stage): + config = { + "train_micro_batch_size_per_gpu": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-3 + } + }, + "zero_optimization": { + "stage": stage + }, + "bf16": { + "enabled": True + }, + } + model = EmptyTailModel() + engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + return engine + + +def _step(engine): + x = torch.randn(1, HIDDEN, device=engine.device, dtype=torch.bfloat16) + loss = engine(x) + engine.backward(loss) + engine.step() + + +@pytest.mark.parametrize("stage", [1, 2]) +class TestZeroNumelParameterShape(DistributedTest): + world_size = 1 + + def test_shape_survives_initialize(self, stage): + engine = _engine(stage) + + assert engine.module.empty.weight.shape == torch.Size([0, HIDDEN]) + # The sized parameter shares the flat buffer, which is what makes the + # zero-element one the special case rather than the rule. + assert engine.module.dense.weight.shape == torch.Size([HIDDEN, HIDDEN]) + + def test_shape_survives_a_step(self, stage): + engine = _engine(stage) + + _step(engine) + + assert engine.global_steps == 1 + assert engine.module.empty.weight.shape == torch.Size([0, HIDDEN]) + + def test_a_second_step_still_runs(self, stage): + # step() rebuilds the parameters from the flat buffer, so a shape lost there + # only shows up on the forward of the iteration after it. + engine = _engine(stage) + + _step(engine) + _step(engine) + + assert engine.global_steps == 2 From f23ef1da7153c9f6d20adf2ef1f8ea5cc17f063b Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 9 Sep 2026 22:55:15 +0800 Subject: [PATCH 2/2] Widen to every wrapper that binds parameters to a flat buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ZeRO-1/2 was not the only one. The same `p.data = q.data` over `unflatten_dense_tensors` output appears in FP16_Optimizer and BF16_Optimizer, and all three drop a zero-element parameter's shape: fp16 + stage 0 FP16_Optimizer (0,) step: RuntimeError bf16 + stage 0 FP16_Optimizer (0,) step: RuntimeError bf16 + stage 1 + fp32 accum BF16_Optimizer (0,) step: RuntimeError So the failure does not need ZeRO at all — any run with fp16 or bf16 enabled and a zero-element parameter breaks on the first forward after `initialize`. Move the guard into `bind_flat_views` in runtime/utils.py and call it from all four binding sites, so the reason is written once. FP16_Optimizer's third site copies rather than rebinds; without the guard that copy would start raising on the shape mismatch once the earlier sites stop corrupting the shape, so it takes the same skip inline. Test parametrized over the five configurations, one per wrapper. Signed-off-by: alanhuangyoo --- deepspeed/runtime/bf16_optimizer.py | 10 +- deepspeed/runtime/fp16/fused_optimizer.py | 14 ++- deepspeed/runtime/utils.py | 15 +++ deepspeed/runtime/zero/stage_1_and_2.py | 16 +-- .../zero/test_zero_numel_param_shape.py | 98 ++++++++++++++----- 5 files changed, 104 insertions(+), 49 deletions(-) diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py index 7b26ed043f39..b89f39d1243d 100644 --- a/deepspeed/runtime/bf16_optimizer.py +++ b/deepspeed/runtime/bf16_optimizer.py @@ -12,9 +12,10 @@ from deepspeed.runtime.base_optimizer import ZeROOptimizer from packaging import version as pkg_version from deepspeed.git_version_info import version -from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim, - align_dense_tensors, all_gather_dp_groups, is_model_parallel_parameter, - see_memory_usage, graph_process, get_norm_with_moe_layers) +from deepspeed.runtime.utils import (bind_flat_views, get_global_norm_of_tensors, clip_tensors_by_global_norm, + DummyOptim, align_dense_tensors, all_gather_dp_groups, + is_model_parallel_parameter, see_memory_usage, graph_process, + get_norm_with_moe_layers) from deepspeed.utils import link_hp_params, lazy_init_hp_params_optimizer_state, fragment_address, groups from deepspeed.moe.utils import is_moe_param, is_moe_param_group from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank @@ -293,8 +294,7 @@ def _split_flat_tensor(self, flat_tensor, num_elem_list): def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor): updated_params = self.unflatten(flat_tensor, tensor_list) - for p, q in zip(tensor_list, updated_params): - p.data = q.data + bind_flat_views(tensor_list, updated_params) def _flatten_dense_tensors_aligned(self, tensor_list, alignment): return self.flatten(align_dense_tensors(tensor_list, alignment)) diff --git a/deepspeed/runtime/fp16/fused_optimizer.py b/deepspeed/runtime/fp16/fused_optimizer.py index 09815d2c035c..6ca80c4e41da 100755 --- a/deepspeed/runtime/fp16/fused_optimizer.py +++ b/deepspeed/runtime/fp16/fused_optimizer.py @@ -10,7 +10,8 @@ import torch from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer -from deepspeed.runtime.utils import get_global_norm, get_flattened_grad_norm, CheckOverflow, get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter +from deepspeed.runtime.utils import (bind_flat_views, get_global_norm, get_flattened_grad_norm, CheckOverflow, + get_weight_norm, get_norm_with_moe_layers, is_model_parallel_parameter) from deepspeed.runtime.fp16.loss_scaler import LossScaleConfig, LossScaleProfile from deepspeed.utils import logger, log_dist from deepspeed.utils.torch import required_torch_version @@ -92,8 +93,7 @@ def __init__(self, self.fp16_groups_flat.append(_flatten_dense_tensors([p.clone().detach() for p in self.fp16_groups[i]])) # set model fp16 weight to slices of flattened buffer updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) - for p, q in zip(self.fp16_groups[i], updated_params): - p.data = q.data + bind_flat_views(self.fp16_groups[i], updated_params) # init master weight, flattened self.fp32_groups_flat.append(self.fp16_groups_flat[i].clone().float().detach()) # modify optimizer of have flat master weight @@ -187,8 +187,7 @@ def step_fused_adam(self, closure=None): # TODO: we probably don't need this? just to be safe for i in range(len(norm_groups)): updated_params = _unflatten_dense_tensors(self.fp16_groups_flat[i], self.fp16_groups[i]) - for p, q in zip(self.fp16_groups[i], updated_params): - p.data = q.data + bind_flat_views(self.fp16_groups[i], updated_params) return self.overflow def set_lr(self, lr): @@ -354,6 +353,11 @@ def step(self, closure=None): for i in range(len(self.fp16_groups)): updated_params = _unflatten_dense_tensors(self.fp32_groups_flat[i], self.fp16_groups[i]) for p, q in zip(self.fp16_groups[i], updated_params): + if p.numel() == 0: + # See bind_flat_views: `q` is a 1-D zeros({0}) here, not a view of + # `p`'s shape, so this copy would raise on the shape mismatch. There + # are no elements to copy either way. + continue p.data.copy_(q.data) self.has_executed_step = True if self.timers: diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index 54a8ddc60d26..8b7a39f2eefd 100644 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -823,6 +823,21 @@ def empty_cache(): get_accelerator().reset_peak_memory_stats() +def bind_flat_views(tensors, views): + """Point each tensor at its view of a flat buffer, skipping zero-element ones. + + torch's ``unflatten_dense_tensors`` special-cases ``numel == 0`` and returns a + freshly allocated 1-D ``zeros({0})`` rather than a view of the requested shape, + so assigning it would replace e.g. a ``(0, 8)`` parameter with a ``(0,)`` one and + break the owning module's own forward. There is no slice of the flat buffer for + such a tensor to point at either, so nothing is left unbound by skipping it. + """ + for tensor, view in zip(tensors, views): + if tensor.numel() == 0: + continue + tensor.data = view.data + + def see_memory_usage(message, force=False): if not force: return diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index a5d1b762f5ae..827a4412cf78 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -21,9 +21,9 @@ from deepspeed.runtime.base_optimizer import ZeROOptimizer from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler from deepspeed.runtime.torch_autocast import get_autocast_dtype, get_all_comm_dtypes, is_autocast_initialized, sort_dtypes -from deepspeed.runtime.utils import (empty_cache, see_memory_usage, has_inf_or_nan, inf, is_model_parallel_parameter, - align_dense_tensors, all_gather_dp_groups, mask_nan_or_inf_with_val_inplace, - count_used_parameters_in_backward) +from deepspeed.runtime.utils import (bind_flat_views, empty_cache, see_memory_usage, has_inf_or_nan, inf, + is_model_parallel_parameter, align_dense_tensors, all_gather_dp_groups, + mask_nan_or_inf_with_val_inplace, count_used_parameters_in_backward) from deepspeed.runtime.zero.config import ZeroStageEnum from deepspeed.runtime.zero.utils import get_norm_dtype from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum @@ -804,15 +804,7 @@ def _configure_moe_settings(self): def _update_model_bit16_weights(self, group_index): updated_params = self.unflatten(self.bit16_groups_flat[group_index], self.round_robin_bit16_meta[group_index]) - for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params): - if p.numel() == 0: - # torch's unflatten_dense_tensors special-cases a zero-element tensor and - # hands back a freshly allocated 1-D `zeros({0})` instead of a view of the - # requested shape, so assigning it would replace e.g. a (0, 8) parameter - # with a (0,) one and break the module's own forward. There is nothing in - # the flat buffer to point such a parameter at anyway. - continue - p.data = q.data + bind_flat_views(self.round_robin_bit16_groups[group_index], updated_params) # set model fp16 weight to slices of reordered flattened buffer for param_index, param in enumerate(self.bit16_groups[group_index]): diff --git a/tests/unit/runtime/zero/test_zero_numel_param_shape.py b/tests/unit/runtime/zero/test_zero_numel_param_shape.py index 2eec4a469f4d..885335c8fc9d 100644 --- a/tests/unit/runtime/zero/test_zero_numel_param_shape.py +++ b/tests/unit/runtime/zero/test_zero_numel_param_shape.py @@ -2,17 +2,17 @@ # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team -"""A zero-element parameter must keep its shape across the ZeRO-1/2 flat buffer. +"""A zero-element parameter must keep its shape when it is bound to a flat buffer. -`_update_model_bit16_weights` repoints every parameter at its slice of the flattened -group. torch's `unflatten_dense_tensors` special-cases a zero-element tensor and returns -a freshly allocated 1-D `zeros({0})` rather than a view of the requested shape, so a -`(0, 8)` parameter came back as `(0,)` and the module's own forward then dispatched -`F.linear` to `addmv`: +Every optimizer wrapper that flattens the parameters repoints each one at its slice of +the flat buffer. torch's `unflatten_dense_tensors` special-cases a zero-element tensor +and returns a freshly allocated 1-D `zeros({0})` rather than a view of the requested +shape, so a `(0, 8)` parameter came back as `(0,)` and the module's own forward then +dispatched `F.linear` to `addmv`: RuntimeError: size mismatch, got input (1), mat (1x8), vec (0) -The parameter is rebuilt on every `step()` as well as at init, so the shape did not +The parameters are rebuilt on every `step()` as well as at init, so the shape did not survive one iteration either. """ @@ -25,6 +25,54 @@ HIDDEN = 8 +# One case per wrapper that binds parameters to a flat buffer. +CONFIGS = { + "fp16_stage0": ({ + "fp16": { + "enabled": True, + "loss_scale": 1.0 + }, + "zero_optimization": { + "stage": 0 + } + }, torch.float16), + "bf16_stage0": ({ + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 0 + } + }, torch.bfloat16), + "bf16_stage1_fp32_accum": ({ + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 1 + }, + "data_types": { + "grad_accum_dtype": "fp32" + } + }, torch.bfloat16), + "zero1": ({ + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 1 + } + }, torch.bfloat16), + "zero2": ({ + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 2 + } + }, torch.bfloat16), +} + class EmptyTailModel(torch.nn.Module): """A trainable parameter with no elements, kept in the autograd graph by the loss.""" @@ -40,7 +88,8 @@ def forward(self, x): return hidden.sum() + self.empty(hidden).sum() -def _engine(stage): +def _engine(case): + extra, _ = CONFIGS[case] config = { "train_micro_batch_size_per_gpu": 1, "optimizer": { @@ -49,51 +98,46 @@ def _engine(stage): "lr": 1e-3 } }, - "zero_optimization": { - "stage": stage - }, - "bf16": { - "enabled": True - }, + **extra, } model = EmptyTailModel() engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) return engine -def _step(engine): - x = torch.randn(1, HIDDEN, device=engine.device, dtype=torch.bfloat16) - loss = engine(x) +def _step(engine, case): + _, dtype = CONFIGS[case] + loss = engine(torch.randn(1, HIDDEN, device=engine.device, dtype=dtype)) engine.backward(loss) engine.step() -@pytest.mark.parametrize("stage", [1, 2]) +@pytest.mark.parametrize("case", list(CONFIGS)) class TestZeroNumelParameterShape(DistributedTest): world_size = 1 - def test_shape_survives_initialize(self, stage): - engine = _engine(stage) + def test_shape_survives_initialize(self, case): + engine = _engine(case) assert engine.module.empty.weight.shape == torch.Size([0, HIDDEN]) # The sized parameter shares the flat buffer, which is what makes the # zero-element one the special case rather than the rule. assert engine.module.dense.weight.shape == torch.Size([HIDDEN, HIDDEN]) - def test_shape_survives_a_step(self, stage): - engine = _engine(stage) + def test_shape_survives_a_step(self, case): + engine = _engine(case) - _step(engine) + _step(engine, case) assert engine.global_steps == 1 assert engine.module.empty.weight.shape == torch.Size([0, HIDDEN]) - def test_a_second_step_still_runs(self, stage): + def test_a_second_step_still_runs(self, case): # step() rebuilds the parameters from the flat buffer, so a shape lost there # only shows up on the forward of the iteration after it. - engine = _engine(stage) + engine = _engine(case) - _step(engine) - _step(engine) + _step(engine, case) + _step(engine, case) assert engine.global_steps == 2