A zero-element parameter loses its shape when it is bound to a flat buffer

[alanhuangyoo]

A zero-element trainable parameter loses its shape when it is bound to a flattened
parameter group, and the model's own forward then fails. This does not need ZeRO — every
wrapper that flattens the parameters does it.

nn.Linear(8, 0, bias=False) in a model, deepspeed.initialize, one step:

config	wrapper	empty.weight after init	step
fp16 + stage 0	FP16_Optimizer	(0,)	RuntimeError
bf16 + stage 0	FP16_Optimizer	(0,)	RuntimeError
bf16 + stage 1 + fp32 accum	BF16_Optimizer	(0,)	RuntimeError
bf16 + ZeRO-1	DeepSpeedZeroOptimizer	(0,)	RuntimeError
bf16 + ZeRO-2	DeepSpeedZeroOptimizer	(0,)	RuntimeError

RuntimeError: size mismatch, got input (1), mat (1x8), vec (0)

That addmv signature is the giveaway — F.linear only dispatches there when the weight
is 1-D. dense.weight next to it is still (8, 8).

Root cause

Each wrapper repoints its parameters at their slices of the flat buffer with the same two
lines, e.g. _update_model_bit16_weights:

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

The shapes handed to unflatten are right — torch.zeros_like(param.data, device="meta") —
but torch.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:

>>> a, b = torch.randn(8, 8), torch.randn(0, 8)
>>> metas = [torch.zeros_like(a, device="meta"), torch.zeros_like(b, device="meta")]
>>> [tuple(m.shape) for m in metas]
[(8, 8), (0, 8)]
>>> [tuple(t.shape) for t in _unflatten_dense_tensors(_flatten_dense_tensors([a, b]), metas)]
[(8, 8), (0,)]

(ATen's unflatten_dense_tensors returns at::zeros({0}, flat.options()) for numel == 0
instead of narrowing and viewing.)

Assigning that to p.data replaces the parameter. The binding runs after every step() as
well as at init, so restoring the shape once would not have held: the second iteration's
forward would break instead of the first.

The change

bind_flat_views(tensors, views) in runtime/utils.py skips a zero-element tensor, and the
four binding sites call it:

• zero/stage_1_and_2.py — _update_model_bit16_weights
• fp16/fused_optimizer.py — init, and after step_fused_adam
• bf16_optimizer.py — _update_storage_to_flattened_tensor

There is no slice of the flat buffer for such a tensor to point at, and what torch returns
is a fresh allocation rather than a view, so the assignment was not keeping anything in
sync either.

FP16_Optimizer.step has a fifth site that copies rather than rebinds
(p.data.copy_(q.data)). Once the earlier sites stop corrupting the shape, that copy would
start raising on the (0, 8) vs (0,) mismatch, so it takes the same skip inline.

Nothing changes for a parameter with elements: the zip order and the narrow/view path are
untouched.

Testing

tests/unit/runtime/zero/test_zero_numel_param_shape.py, parametrized over the five
configurations above — the shape after initialize, the shape after a step, and a second
step running at all (the case that would survive a fix applied only at init).

On an H20, torch 2.9.1+cu128:

this branch:  15 passed
master:       15 failed
                5 x AssertionError                 (shape is (0,) after initialize)
               10 x RuntimeError: size mismatch, got input (1), mat (1x8), vec (0)

Regression, same box, DS_SKIP_CUDA_CHECK=1 so the CPU-Adam builds —
test_stage2_flatten_on_gpu.py, test_zero_tensor_fragment.py,
test_zero_coalesce_grad_reduction.py: 174 passed, 75 skipped, 0 failed.

Related

#8280 and #8298 (issues #8279, #8297) fixed the ZeRO-1/2 reduction path for zero-element
parameters. This is the parameter-binding path, which those did not reach — a model matching
#8279's shape still fails before reduction is ever attempted.

ZeRO-3 is untouched: it partitions to a 1-D local shard and keeps the real shape in
ds_shape by design, so (0,) there is correct. Its own zero-element failure is the
all-gather gate in #8375, a different bug in a different file.

[alanhuangyoo]

Scope widened after opening this: ZeRO-1/2 was not the only place. FP16_Optimizer and BF16_Optimizer bind parameters to a flat buffer with the same two lines, and drop the same shape — 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.

Measured on an H20 against unmodified master (cbd303e):

fp16 + stage 0               FP16_Optimizer   shape after init (0,)   step: RuntimeError
bf16 + stage 0               FP16_Optimizer   shape after init (0,)   step: RuntimeError
bf16 + stage 1 + fp32 accum  BF16_Optimizer   shape after init (0,)   step: RuntimeError
bf16 + ZeRO-1                                 shape after init (0,)   step: RuntimeError
bf16 + ZeRO-2                                 shape after init (0,)   step: RuntimeError

09d6f71 moves the guard into bind_flat_views in runtime/utils.py so the reason is written once, and calls it from all four binding sites. FP16_Optimizer.step has a fifth site that copies rather than rebinds; without a guard there that copy would start raising on the (0, 8) vs (0,) mismatch once the earlier sites stop corrupting the shape, so it takes the same skip inline — worth a look, since it is the one place the fix could have introduced a new failure rather than removed one.

Test is now parametrized over those five configurations: 15 passed here, 15 failed on master. Regression on the same box with DS_SKIP_CUDA_CHECK=1 (test_stage2_flatten_on_gpu, test_zero_tensor_fragment, test_zero_coalesce_grad_reduction): 174 passed, 75 skipped, 0 failed.