Fix ZeRO-3 all_reduce param fetch stride for padded parameters

[ebarkhordar]

Problem

In Init._all_gather_coalesced, the branch taken when
stage3_use_all_reduce_for_fetch_params is enabled sizes the flat buffer with the
aligned element count but walks it with the unaligned one:

flat_buffer_size = sum(p.ds_numel_aligned for p in params)   # aligned
...
    start_param += param.ds_numel                            # unaligned

ds_numel_aligned is ds_numel rounded up to a multiple of the partition world size, so
the two are equal only when ds_numel % world_size == 0. When a parameter is padded, every
following parameter's base offset lands padding_i slots inside its predecessor's aligned
span.

Those overlapping slots are the predecessor's partition padding. They are not harmless:
_partition_param allocates the partition with torch.empty and, on the one rank whose
partition extends past the end of the parameter, copies only elems_to_copy elements, so
the tail is never written and holds whatever the allocator returned. Because the fetch is
an all_reduce with SUM, that uninitialized tail is added into the next parameter's
leading elements. The sibling all_gather path writes the same padding but each parameter
narrows only its own ds_numel out of the result, so it is never read; the corruption is
specific to the SUM reconstruction.

Total stride is smaller than the allocated buffer, so there is no out-of-bounds access and
nothing crashes. The parameter simply comes back with wrong leading values.

Invariant

Each parameter owns a disjoint contiguous span of ds_numel_aligned slots in flat_tensor,
which is exactly what flat_buffer_size already reserves for it.

ds_numel_aligned has a single writer in the package, param.ds_numel_aligned = tensor_size
in _partition_param, set in the same block that creates ds_tensor. It has two readers:
the buffer sizing above and the line changed here. Since the buffer sizing already reads the
attribute for every param in this same loop, the change introduces no new requirement on
when the attribute must exist.

Fix

Advance start_param by ds_numel_aligned, matching the buffer sizing.

Verification

Added tests/unit/runtime/zero/test_zero_allreduce_fetch_params.py, parametrized over a
padded case (numels 5 and 7 under world_size=2) and an aligned case (numels 4 and 6):

case	master	this branch
padded	fails	passes
aligned	passes	passes

On master the padded case reports:

param p1 was not reconstructed exactly:
  expected [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
  got      [7778.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]

7778 is 7777 (the sentinel filling p0's uninitialized padding) plus 1.0 (p1's real
first element), which is the SUM overlap rather than a generically wrong value.

About that sentinel: the padding tail is uninitialized by construction, and a fresh
allocation usually reads back as zero, in which case the SUM adds zero and the bug is
invisible. A first attempt at this reproduction showed no corruption for exactly that
reason. The test therefore fills new allocations with a sentinel so the tail's contents are
deterministic rather than dependent on whether the allocator hands back a recycled block.
It does not inject anything into the buffer under test.

Run in a CPU-only container, world_size 2 over gloo, LOCAL_SIZE=2:

python -m pytest unit/runtime/zero/test_zero_allreduce_fetch_params.py

pre-commit run --files passes on both changed files.

What I did not verify

I have no multi-GPU machine here, so this was exercised only on CPU with gloo, not with
NCCL on real devices. I also did not measure the effect on a full training run, so I cannot
say how the corrupted leading elements affect convergence in practice; the claim here is
limited to the reconstructed parameter values being wrong.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0ea262c4f1

ℹ️ 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".

[chatgpt-codex-connector]

Restore torch.empty after poisoning _orig_torch_empty

When this test exits the deepspeed.zero.Init context, Init.__exit__ restores torch.empty from partition_parameters._orig_torch_empty; because this assignment is still pointing at poisoned_empty until the finally block runs, torch.empty is left globally patched to return sentinel-filled tensors for the remainder of the worker process. In contexts that reuse the distributed worker or add any later allocation in this test, unrelated code will see poisoned allocations, so the cleanup needs to restore the public torch.empty binding as well or avoid mutating _orig_torch_empty through the context teardown.

Useful? React with 👍 / 👎.

[ebarkhordar]

Good catch, this is real and I have pushed a fix in 409ccd7.

_remove_tensor_creation_wrappers (partition_parameters.py:659) reads the module global at exit time, so Init.__exit__ rebound the public torch.empty to the test's poisoned_empty before the finally block restored _orig_torch_empty. Measured it directly against this branch:

torch.empty is real_empty : False
torch.empty is poisoned   : True
post-block torch.empty(3) : [7777.0, 7777.0, 7777.0]

Every later allocation in that worker process would have come back sentinel-filled. The fix saves the public binding before the block and restores it alongside the module global; the same probe after the change reports torch.empty restored: True and a normal uninitialized tensor.

The poisoning inside the block is unchanged, so the padded/aligned differential the PR rests on still holds.

[tohtana]

Hi @ebarkhordar,
Great catch! It's surprising that we still have such a bug. I really appreciate your contribution.

[PKUWZP]

@ebarkhordar Thanks so much for this PR! In order to merge the PR, it would be ideal to run some tests on multi-GPU settings. I can help run some tests on my end, can you give me a 1-2 days so we can fully validate the fix and merge? Thanks for your contribution to DeepSpeed.

[ebarkhordar]

Thanks @tohtana, and thanks @PKUWZP for offering to run the multi-GPU validation. That plan works for me, please take the time you need. I don't have a multi-GPU box, so I verified the fix single-process only (the unit test simulates world_size=2 with padded numels); a real multi-rank run on the all_reduce fetch path is exactly the check it should have before merging, so I would rather you confirm it than merge on my word. Happy to follow up on anything the runs turn up.

On the red checks, so they do not muddy the picture: the two failing modal-torch-latest legs are the base workflow's pull_request_target checkout refusing fork code on the re-run ("Refusing to check out fork pull request code from a 'pull_request_target' workflow"), not the change. The same commit (409ccd7) passed the full modal CI on the first run.

[ebarkhordar]

Thanks @tohtana and @PKUWZP. No rush at all, take the days you need. I verified the stride fix with a 2-rank test locally (test_zero_allreduce_fetch_params.py), but I do not have real multi-GPU hardware here, so your validation is the right gate. Happy to add cases or adjust if your runs surface anything.