Fix ZeRO++ secondary shard copy for small params (#8210)
## Summary
This PR fixes a ZeRO++ edge case in _partition_param_sec() when a small
parameter does not overlap with its computed secondary shard.
In the reported case from deepspeedai/DeepSpeed#6659, a parameter of
shape [32] is DP-aligned to 2048, and with zero_hpz_partition_size=16
the secondary shard size becomes 128. For some secondary-group ranks,
secondary_start is already beyond param.ds_numel, so sec_numel becomes
0. The old code still executes:
one_dim_param.narrow(0, secondary_start, sec_numel)
PyTorch raises IndexError for this case even when sec_numel == 0 if the
start index is out of range.
This change fixes the issue by:
- skipping the secondary copy when sec_numel == 0
- zero-filling the secondary shard buffer first so uncovered padding
remains deterministic and does not leak uninitialized values into later
coalesced quantization
This PR also adds focused regression tests covering:
- the small-parameter, no-overlap secondary shard case
- zeroed padding for partially covered secondary shards
Fixes deepspeedai/DeepSpeed#6659
## Testing
python3 -m pytest -q tests/unit/runtime/zero/test_zeropp.py -k
'small_param_secondary_shard_without_overlap or
secondary_shard_padding_is_zeroed'
Observed locally:
- current master reproduces the IndexError on the
deepspeedai/DeepSpeed#6659 geometry
- this PR avoids the out-of-range narrow() call
- 2 passed
---------
Signed-off-by: zengyong <2595650269@qq.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
