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 f05a53867c93..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,8 +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): - 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 new file mode 100644 index 000000000000..885335c8fc9d --- /dev/null +++ b/tests/unit/runtime/zero/test_zero_numel_param_shape.py @@ -0,0 +1,143 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""A zero-element parameter must keep its shape when it is bound to a flat buffer. + +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 parameters are 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 + +# 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.""" + + 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(case): + extra, _ = CONFIGS[case] + config = { + "train_micro_batch_size_per_gpu": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-3 + } + }, + **extra, + } + model = EmptyTailModel() + engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + return engine + + +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("case", list(CONFIGS)) +class TestZeroNumelParameterShape(DistributedTest): + world_size = 1 + + 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, case): + engine = _engine(case) + + _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, 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(case) + + _step(engine, case) + _step(engine, case) + + assert engine.global_steps == 2