Fix silent gradient-bucket drop on CPU in ZeRO-1/2 non-contiguous reduction - #8382
Fix silent gradient-bucket drop on CPU in ZeRO-1/2 non-contiguous reduction#8382delock wants to merge 2 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3757a3b933
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # building them from the accelerator device name matches nothing on CPU and | ||
| # silently drops every gradient bucket, skipping the all-reduce entirely. | ||
| # Compare dtypes directly so the buckets are device-independent. | ||
| dtypes = [torch.half, torch.float, torch.double, torch.bfloat16] |
There was a problem hiding this comment.
This non-merge commit's raw message contains no Signed-off-by trailer, so it violates the repository's mandatory DCO requirement and may be rejected during integration. Recreate the commit with git commit --signoff before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
… harness split_half_float_double() built legacy type strings from the accelerator device name, but CPU tensors report 'torch.FloatTensor' without any device prefix, so no tensor ever matched: with contiguous_gradients=False and reduce_scatter=False, ZeRO-1/2 silently skipped the gradient all-reduce on CPU and each rank updated on local gradients. Mean configurations masked the bug because identical per-rank gradients make the unreduced local value equal the mean; only sum configurations expose it. Group buckets by dtype so bucketing is device-independent. reduce_boolean_flags() in tests/unit/common.py passed get_accelerator().current_device() to torch, which is a bare rank index (a LOCAL_RANK string on CPU) that torch cannot parse as a device, so any cross-rank comparison on CPU crashed with 'Invalid device string'. Use current_device_name(), valid on every accelerator, and carry the flag as a 1-element tensor since gloo rejects 0-dim all_gather_into_tensor inputs. Verified locally (CPU/gloo, world sizes 2-3): - TestGradientAllreduceOp 18/18 pass (sum+predivide cases previously wrong) - TestGradientAllreduceOpTraining and unmanaged offload comparisons now complete instead of crashing in the comparison helper Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
3757a3b to
ad3c27e
Compare
| buckets = [] | ||
| for i, dtype in enumerate(dtypes): | ||
| bucket = [t for t in tensors if t.type() == dtype] | ||
| bucket = [t for t in tensors if t.dtype == dtype] |
There was a problem hiding this comment.
t.type() encodes layout as well as dtype, so this also changes membership for sparse gradients, not just on CPU. sparse.type() is torch.sparse.FloatTensor, which matches none of the four legacy strings on any device, so sparse grads were dropped from every bucket. t.dtype == torch.float does match them, and allreduce_bucket opens with self.flatten(bucket), which is _flatten_dense_tensors.
>>> sparse = torch.sparse_coo_tensor(torch.tensor([[0, 2]]), torch.tensor([1., 2.]), (4,))
>>> sparse.type(), sparse.dtype
('torch.sparse.FloatTensor', torch.float32)
>>> _flatten_dense_tensors([torch.zeros(4), sparse])
RuntimeError: unsupported memory format option Contiguous
That is torch 2.14.0+cpu in a clean container, against this branch at ad3c27e.
I have not checked whether a sparse grad can actually reach here. ZeRO 1/2 has no sparse-tensor handling of its own, and the is_sparse routing at engine.py:3824 is on the non-ZeRO path, so it may be unreachable. If it is reachable, an nn.Embedding(sparse=True) model on ZeRO-1/2 with contiguous_gradients=False moves from a silently skipped all-reduce to a hard failure.
Adding and not t.is_sparse on line 74 would preserve the old membership exactly, if you would rather keep this scoped to the device prefix.
split_half_float_double() now matches dtypes directly, which admits sparse gradient layouts on every device; the legacy type strings never matched them anywhere (e.g. torch.cuda.sparse.FloatTensor). Keep them out of the buckets so membership stays identical to the legacy behavior on all accelerators. ZeRO-1/2 does not support sparse gradients and still fails later in get_flat_partition() for such configurations. The direct unit test pins both the CPU membership regression (the legacy type strings matched nothing on CPU) and the sparse exclusion. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
2 participants
Problem
While investigating multi-rank CPU CI (#8381),
TestGradientAllreduceOpfailed for two variants (sum+prescale_gradients=True+ ZeRO-1/2 +contiguous_gradients=False+reduce_scatter=False) with results off by a factor of world_size, andTestUnmanagedGradientAccumulationOffloadshowed a numeric mismatch. Separately, any test comparing tensors across ranks on CPU crashed withRuntimeError: Invalid device stringor a gloo allgather error.Root cause
split_half_float_double()built legacy type strings from the accelerator device name (torch.{device}.FloatTensor). CPU tensors reporttorch.FloatTensorwith no device prefix, so no tensor ever matched: the bucket list came back empty and the gradient all-reduce was silently skipped — every rank updated on its local gradients.gradient_allreduce_op=meanthis stays invisible: identical per-rank gradients make the unreduced local value equal the mean.sumthe result is wrong by a factor of world_size.The same variants pass on GPU CI (all 18
TestGradientAllreduceOpvariants and 5TestGradientAllreduceOpTrainingvariants are green onmodal-torch-latest), becausetorch.cuda.FloatTensormatches on CUDA — the bug is CPU-specific, and CPU multi-rank runs were the first to reach this code.The harness crashes were two stacked bugs in
reduce_boolean_flags()(tests/unit/common.py):current_device()returns a bare rank index (aLOCAL_RANKstring on CPU) that torch cannot parse as a device, and gloo rejects 0-dim inputs toall_gather_into_tensor.Fix
t.dtypeinstead of legacy type strings, making bucketing device-independent (behavior-equivalent on CUDA, where the legacy strings matched).current_device_name()(valid on every accelerator) and carry the flag as a 1-element tensor.Verification (CPU/gloo, world sizes 2–3)
TestGradientAllreduceOp(18 variants)TestGradientAllreduceOpTraining(5)TestUnmanagedGradientAccumulationOffload+InactiveParams(6)TestZero3ParamPartitioningBase/TestZeroToFP32(regression smoke)GPU behavior is unchanged by construction (same buckets, same order); GPU CI will re-confirm.
Follow-up
With the comparison helper fixed,
test_unmanaged_varying_backward_count[3](ZeRO-3) can now reach its numeric comparison and reveals a pre-existing unmanaged-vs-managed mismatch on CPU that was previously masked by the harness crash. Out of scope here.