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

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

ZeRO-1/2 had the same failure until #8277, which binds with narrow().view(meta.shape) and so keeps the shape. The tests still cover them.

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. FP16_Optimizer.__init__:

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

The shapes handed to unflatten are right (here they are the parameters themselves),
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
three binding sites call it:

• 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:

Re-run after merging master (b5e000c):

this branch:  15 passed
master:        9 failed, 6 passed (the zero1/zero2 cases)
                3 x AssertionError                 (shape is (0,) after initialize)
                6 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.

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.

[ebarkhordar]

I ran your tests at 990c49c and at the merge base b5e000c in a clean container (torch 2.8.0+cpu, deepspeed.__file__ printed each run). CPU, so this is the four non-fp16 cases.

Base b5e000c: 6 failed, 6 passed. bf16_stage0 and bf16_stage1_fp32_accum fail all three, with the signature from your description:

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

zero1 and zero2 pass at the base, so the #8277 attribution holds up. Head 990c49c: 12 passed. The parameter bind_flat_views skips also matches the bound one, empty.weight (0, 8) bfloat16 cpu beside dense.weight (8, 8) bfloat16 cpu, so skipping does not strand it on another dtype or device.

On "every wrapper that flattens the parameters": I parsed the package for unflatten call sites rather than grepping. 16, of which the ones that bind the result back are your three, stage_1_and_2.py:857 (already narrow().view() from #8277), and four in stage3.py. Those four bind fp16_partitioned_groups, which is [param.ds_tensor ...] at stage3.py:899 and 1-D, so there is no shape there to lose. The set looks closed.

One question. Stage 3 is the wrapper missing from your table, and EmptyTailModel under it fails at BOTH refs, identically:

AssertionError: {'id': 1, 'status': 'NOT_AVAILABLE', 'numel': 0,
 'ds_numel': 0, 'shape': (0,), 'ds_shape': (0, 8), 'persist': True}

ds_shape is intact, so that is a different bug and your diff does not touch stage3.py. Worth a line saying stage 3 is out of scope?

No GPU here, so fp16 and anything multi-rank is untested by me.

[ebarkhordar]

Answering my own question, since I should have checked before asking it: #8375 is already open for the stage 3 case, with the same assertion text, and its description says the ZeRO-1/2 shape loss is fixed separately. So the scope split here is deliberate and nothing needs adding to this body. The measured half of my comment stands.