Gather zero-sized parameters in ZeRO-3

[alanhuangyoo]

ZeRO-3 fails on a zero-sized trainable parameter.

nn.Linear(8, 0, bias=False) used in the loss, one step:

stage 3: AssertionError: {'id': 1, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0,
                          'shape': (0,), 'ds_shape': (0, 8), 'requires_grad': True, ...}

Correction to an earlier version of this description, which said stages 1 and 2 already
worked. They do not, on this model shape, with #8280 and #8298 in master:

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

That is a different bug in a different place — _update_model_bit16_weights drops a
zero-element parameter's shape when it repoints it at the ZeRO-1/2 flat buffer, so the
module's own forward breaks before any of the code in this PR runs. It is fixed separately;
this PR is scoped to the ZeRO-3 gather gate and its test is now parametrized on stage 3 only.

raised at partitioned_param_coordinator.py:413.

Root cause

fetch_sub_module decides whether to run the all-gather from the number of elements left to
fetch:

fetch_numel = sum(
    [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE])

if fetch_numel > 0:
    ...
    self.__all_gather_params(params_to_fetch, forward)

A submodule whose only parameter is zero-sized contributes 0 to that sum, so the gather is
skipped entirely. The parameter never leaves NOT_AVAILABLE, and the wait loop immediately
below asserts that it is AVAILABLE.

The element count is a proxy for "is there anything to fetch", and it is the wrong proxy: a
zero-sized parameter has no bytes to move but still needs its status transitioned.

The change

Gate on whether any parameter still needs gathering. fetch_numel is kept for the profiler
event, which is what it was actually for.

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

Nothing changes when at least one parameter has elements — the two conditions agree everywhere
except the all-zero case, which today cannot proceed at all.

Relation to the stage 1/2 fixes

Same class of failure, different place. #8280 (issue #8279) skipped zero-sized parameters before
ZeRO-1/2 gradient reduction and #8298 (issue #8297) skipped them in HP fragment mapping; both
are about not processing a parameter with no elements. Stage 3 is the opposite — it has to
process it, because the status machine tracks the parameter and not its bytes.

I did not find an issue or PR covering stage 3.

Verification

Reproducer run against a clean upstream/master worktree and against this branch, same
environment, one process and two gloo ranks:

                     master        this branch
world_size=1
  stage 1              OK              OK
  stage 2              OK              OK
  stage 3      AssertionError          OK
world_size=2
  stage 1              OK              OK
  stage 3      AssertionError          OK

The two-rank case matters on its own: at world_size=1 the fetch takes the
_no_gather_coalesced shortcut, so only the multi-rank run exercises the real all-gather behind
the gate.

tests/unit/runtime/zero/test_zero_empty_param.py covers all three stages at world_size 1 and
2. Note that #8280's description mentions a file of this name, but no test file landed with it,
so stages 1 and 2 have had no in-tree coverage either — this adds it alongside stage 3.

$ pytest tests/unit/runtime/zero/test_zero_empty_param.py     # this branch
3 passed, 3 skipped

$ pytest tests/unit/runtime/zero/test_zero_empty_param.py     # upstream/master
1 failed, 2 passed, 3 skipped
E   AssertionError: {'id': 1, 'status': 'NOT_AVAILABLE', 'numel': 0, ...}

$ yapf==0.40.0 --diff  /  flake8
(clean)

Is one gate enough?

Checked that this is the only thing standing in the way, rather than the first of several, by
running the same zero-sized-parameter model through the other stage-3 paths on this branch:

zero.Init context          OK
three steps (prefetch)     OK
CPU param + optimizer offload   OK
save_checkpoint + load_checkpoint   OK

All four fail at the same assert on master, and all four pass here, so the numel gate in
fetch_sub_module is the whole of it.

The 3 skips are the world_size=2 class, which this single-accelerator box cannot schedule;
those cases were run directly over gloo instead, and are the two-rank rows in the table above.

[sfc-gh-truwase]

What is the source of zero-sized parameters and how they are set to ZeroParamStatus.NOT_AVAILABLE status? I think a better solution would be to avoid them.

[alanhuangyoo]

Where they come from. Model code declares them. An in-tree example is transformers'
FP-Quant integration, which installs zero-sized placeholders for the quantized tensors:

# transformers/integrations/fp_quant.py:101
".weight":   torch.nn.Parameter(torch.zeros(0)),
".dqweight": torch.nn.Parameter(torch.zeros(0)),
".qweight":  torch.nn.Parameter(torch.zeros(0)),
".scales":   torch.nn.Parameter(torch.zeros(0)),

More generally any config that drives a dimension to zero produces one — a classifier head with
no labels, an adapter left at rank 0, a modality tower that is configured off. The repro in the
PR is the minimal version of that, not a synthetic special case.

How they end up NOT_AVAILABLE. Nothing special happens to them.
partition_parameters.py:324 sets NOT_AVAILABLE on every parameter it partitions, regardless
of size, and release_and_reset_all does the same at :1744. So a zero-sized parameter enters
fetch_sub_module in exactly the state every other parameter is in.

The mismatch is that the status is tracked per parameter while the gate is on a sum of
elements:

fetch_numel = sum(p.partition_numel() for p in params_to_fetch
                  if p.ds_status == ZeroParamStatus.NOT_AVAILABLE)
if fetch_numel > 0:
    ...   # this is what flips them to AVAILABLE

A submodule whose only ungathered parameter is zero-sized sums to 0, the block is skipped, and
nothing moves the parameter out of NOT_AVAILABLE. Twelve lines further down the same function
asserts that it did:

assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()

which is the AssertionError in the PR description. The fetch_numel value is still needed for
the trace and prefetch accounting, so the change keeps computing it and only moves the branch
onto "is anything still ungathered".

On avoiding them instead. That would be a change of policy rather than a smaller fix — the
runtime already tolerates them in every other place it meets them:

runtime/utils.py:51	"Filter out empty parameters (numel == 0) from optimizer params"
runtime/utils.py:245	if x.numel() == 0
runtime/engine.py:3804	if param.numel() == 0
runtime/zero/stage_1_and_2.py:1205, :1270	zero-sized guards on the ZeRO-1/2 paths
runtime/zenflow/engine_stage3.py:281	if param.selected_indices.numel() == 0
fp16/onebit/lamb.py:86	"Filter out empty parameters (numel == 0) to avoid NaN"

ZeRO-3's own test fixtures assume the same. tests/unit/v1/compile/test_z3_eager_fallback.py:23
builds a stage-3 module whose single parameter is exactly this:

param = torch.nn.Parameter(torch.empty(0))
param.ds_id = 7
param.ds_status = ZeroParamStatus.NOT_AVAILABLE

Avoiding them would mean either rejecting models that legitimately declare them, or dropping
them from the ZeRO-3 registry — and dropping them changes what state_dict() contains, so
checkpoints would no longer round-trip against the original module.

I also checked this is the only gate in the way rather than the first of several: on this branch
the same model runs through zero.Init, three steps of prefetch, CPU param + optimizer offload,
and save_checkpoint/load_checkpoint; all four raise the same assert on master and all four
pass here.

Happy to go the other way if you would rather ZeRO-3 reject these outright — just say which, and
I will redo it as a clear error at Init time instead.

[alanhuangyoo]

@fwerkor @sfc-gh-truwase — this is the ZeRO-3 half of #8280, which you two wrote and reviewed for ZeRO-1/2, so flagging it rather than letting it sit unreviewed.

Same failure, same cause, different stage: the all-gather is gated on total elements > 0, so a sub-module holding only zero-sized parameters is skipped entirely, its parameters stay at NOT_AVAILABLE, and the next assert kills the run. #8280 fixed the equivalent gate in ZeRO-1/2 and ZeRO-3 was left with it.

No rush if it is queued; I mostly want to be sure it did not fall through because it looks like a duplicate of #8280.

[alanhuangyoo]

Correction to the description, and a narrowing.

I ran the parametrized test rather than trusting the earlier claim in it, and stages 1 and 2 do not work on this model shape with #8280 and #8298 in master:

TestZeroSizedParameterSingleRank::test_step_completes[1]  RuntimeError: size mismatch, got input (1), mat (1x8), vec (0)
TestZeroSizedParameterSingleRank::test_step_completes[2]  RuntimeError: size mismatch, got input (1), mat (1x8), vec (0)
TestZeroSizedParameterPartitioned::test_step_completes[1] same
TestZeroSizedParameterPartitioned::test_step_completes[2] same
TestZeroSizedParameterSingleRank::test_step_completes[3]  PASSED   (with this PR)
TestZeroSizedParameterPartitioned::test_step_completes[3] PASSED   (with this PR)

That is a different bug in a different file: _update_model_bit16_weights drops a zero-element parameter's shape when it repoints it at the ZeRO-1/2 flat buffer, so (0, 8) becomes (0,) and the module's own forward breaks before any of the code in this PR runs. #8467 fixes that one; the two are independent and can land in either order.

So this PR is now scoped to the ZeRO-3 gather gate: d4ea27b narrows the parametrize to stage 3 and corrects the docstring, which claimed stages 1 and 2 already worked. The code change is unchanged.

[fwerkor]

I think these need to be addressed before merge:

1. Coalesced gather still fails. AllGatherCoalescedHandle.wait() gets an empty partitions list for a zero-sized parameter and then calls torch.cat(partitions), raising ValueError. The current test only covers the single-parameter sequential path. Please add a coalesced regression test and handle this case.

2. zero_quantized_weights=True still fails. Empty params reach CUDAQuantizer.quantize(), where groups == 0 leads to ZeroDivisionError. Please bypass or explicitly handle empty params in the quantized path, with coverage.

3. Prefetch still uses a numel gate. An all-zero prefetch set is popped from the queue but not gathered because submission still requires numel_prefetching > 0. Please make this consistent with the new parameter-based fetch condition.

[alanhuangyoo]

Thanks — all three are real. 889e4b5 handles them, with one correction and one thing I could not make observable.

Correction first: all three reproduce on unmodified master (cbd303e), so they are pre-existing rather than regressions from this PR. That does not change what needs doing — a zero-sized parameter cannot get through ZeRO-3 without them — but it is worth knowing they are not this PR's doing:

master, coalesced bucket (sized + zero-sized param on one module):
    ValueError: torch.cat(): expected a non-empty list of Tensors
master, zero_quantized_nontrainable_weights + frozen zero-sized param:
    ZeroDivisionError: integer modulo by zero

The first one already bites on master because a bucket mixing sized and zero-sized parameters passes the fetch_numel > 0 gate on the sized ones, and then the handle chokes on the zero-sized one. This PR only widens the set of shapes that reach it.

---

1. Coalesced gather. Confirmed, exactly as you describe. Fixed by handing the parameter an empty tensor of its own shape rather than concatenating nothing.

The dtype is the part worth a look. The sized parameters of a bucket come out at param.ds_tensor.dtype on the plain path, but under quantization the partition is stored as int8 and the sized parameters are dequantized to param.dtype on the way out. Using ds_tensor.dtype unconditionally is what my first attempt did, and it hands the module an int8 parameter:

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: c10::BFloat16 != signed char

So the branch picks param.dtype if self.quantization else param.ds_tensor.dtype. The test asserts the returned parameter's dtype matches its sized sibling's, which is the property the module's own forward depends on.

2. zero_quantized_weights. Confirmed. The trigger is narrower than "empty params reach quantize()": in the coalesced path the bucket's ds_tensors are concatenated before quantizing, so a zero-sized parameter alongside a sized one contributes nothing and nothing breaks. It needs a partition that is entirely zero-sized — a submodule whose parameters are all zero-sized, or zero_quantized_nontrainable_weights with a frozen zero-sized parameter, which quantizes per-partition in _partition_param. Both are covered.

The fix is inside CUDAQuantizer: return the empty (int8, fp32 scale) pair rather than bypassing at the call sites, so ds_quant_scale is still set on the partition and the hasattr(params[0].ds_tensor, "ds_quant_scale") checks in _all_gather_coalesced stay uniform across a mixed bucket. dequantize gets the mirror guard.

3. Prefetch gate. Changed to gate on the parameter set, as you asked. But I could not construct a case where it changes behaviour, and I would rather say so than imply the test covers it:

• reverting only that line leaves all six tests passing;
• tracing __all_gather_params over a model with four all-zero-sized submodules and stage3_prefetch_bucket_size set gives the identical {'calls': 10, 'all_zero_calls': 8, 'all_zero_params': 16} with and without it.

The reason looks to be that the parameters are popped from the queue but fetch_sub_module still gathers them when their submodule is reached, so what the count costs there is the prefetch overlap, not the step. I have kept the change — it is the same proxy this PR removes from the fetch, and leaving the two inconsistent invites the next bug — but the test named for it only claims what it proves, which is the fetch gate under a prefetch configuration. If you know a shape where the prefetch gate alone breaks a step, I will add it.

---

Testing. 2×H20, torch 2.9.1+cu128, world_size=2:

this branch:  6 passed
master:       6 failed
    ValueError: torch.cat(): expected a non-empty list of Tensors   (coalesced)
    ZeroDivisionError: integer modulo by zero                        (quantized non-trainable)
    AssertionError: {'status': 'NOT_AVAILABLE', 'numel': 0, ...}     (the other four)

The quantized tests skip when QuantizerBuilder is unavailable, following test_coalesced_collectives.py.