DeepSpeed

4e3acb1e - Fix ZeRO-3 all_reduce param fetch stride for padded parameters (#8158)

Fix ZeRO-3 all_reduce param fetch stride for padded parameters (#8158) ## 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: ```python 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. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Co-authored-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>