Gather zero-sized parameters in ZeRO-3 - #8375
Conversation
fetch_sub_module decides whether to run the all-gather from the number of
elements still 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:
A submodule whose only parameter is zero-sized contributes nothing to that sum,
so the gather is skipped, the parameter never leaves NOT_AVAILABLE, and the wait
loop right below asserts that it is AVAILABLE:
AssertionError: {'id': 1, 'status': 'NOT_AVAILABLE', 'numel': 0, ...}
Gate on whether any parameter still needs gathering instead. The element count
is kept for the profiler, which is what it was for.
This is the ZeRO-3 counterpart of deepspeedai#8280 and deepspeedai#8298, which fixed the same class of
failure for stages 1 and 2 (issues deepspeedai#8279 and deepspeedai#8297).
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
| params_to_fetch = set(iter_params(current_submodule, recurse=is_leaf)) | ||
| fetch_numel = sum( | ||
| [p.partition_numel() for p in params_to_fetch if p.ds_status == ZeroParamStatus.NOT_AVAILABLE]) | ||
| # Gate on whether anything still has to be gathered, not on how many elements that is: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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 AVAILABLEA 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_AVAILABLEAvoiding 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.
2 participants
ZeRO-3 fails on a zero-sized trainable parameter, where stages 1 and 2 now work.
nn.Linear(8, 0, bias=False)used in the loss, one step:raised at
partitioned_param_coordinator.py:413.Root cause
fetch_sub_moduledecides whether to run the all-gather from the number of elements left tofetch:
A submodule whose only parameter is zero-sized contributes
0to that sum, so the gather isskipped entirely. The parameter never leaves
NOT_AVAILABLE, and the wait loop immediatelybelow 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_numelis kept for the profilerevent, which is what it was actually for.
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/masterworktree and against this branch, sameenvironment, one process and two gloo ranks:
The two-rank case matters on its own: at
world_size=1the fetch takes the_no_gather_coalescedshortcut, so only the multi-rank run exercises the real all-gather behindthe gate.
tests/unit/runtime/zero/test_zero_empty_param.pycovers all three stages atworld_size1 and2. 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.
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:
All four fail at the same assert on master, and all four pass here, so the numel gate in
fetch_sub_moduleis the whole of it.The 3 skips are the
world_size=2class, 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.