Fix ZeRO++ secondary shard copy for small params

[ZyEng-Art]

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 #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 #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 [BUG] ZeRO++ sharding small parameter raise IndexError #6659 geometry
• this PR avoids the out-of-range narrow() call
• 2 passed

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: fb6daa6efc

ℹ️ 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]

Zero-fill skipped secondary shards

In HPZ runs with zero_quantized_weights enabled, partitioned_param_coordinator.py:580 passes quantize=True and _all_gather_coalesced quantizes a torch.cat of every ds_secondary_tensor (partition_parameters.py:1406-1409). When sec_numel == 0, this branch skips the only write into the newly allocated torch.empty secondary tensor, so a small no-overlap parameter can feed arbitrary padding values into the shared quantization group for neighboring parameters instead of deterministic zeros. Please initialize the skipped shard before it can be gathered or quantized.

Useful? React with 👍 / 👎.

[goransh-buh]

This looks like a solid, focused change.

[sfc-gh-truwase]

PyTorch raises IndexError for this case even when sec_numel == 0.

@ZyEng-Art thanks for the PR. This looks like a valid issue, but I am curious why it has not previously showed up. Can you please share details of this repro or a stack trace? Thanks!

[ZyEng-Art]

@sfc-gh-truwase Thanks for the review. This seems to have stayed relatively rare because it only affects the ZeRO++ secondary-shard path (when zero_hpz_partition_size > 1), and it further requires a small parameter together with a
large DP partition world size so that the computed secondary shard has no overlap with the real tensor.

I reproduced it locally using the same geometry as the original bug report in issue #6659: a parameter of shape [32], dp_world_size=2048, and zero_hpz_partition_size=16 (so
num_ranks_in_param_group=16). In that setup, the parameter is DP-aligned to 2048, so secondary_partition_size = 2048 / 16 = 128. For secondary-group ranks beyond the first one, secondary_start is already beyond
param.ds_numel; for example, with rank_in_group=15, secondary_start=1920, while the real tensor only has 32 elements. That makes sec_numel=0, but the old code still executes one_dim_param.narrow(0, secondary_start,
sec_numel).

PyTorch still raises IndexError for a zero-length narrow() when the start index is out of range. The stack trace on current master is:

  Traceback (most recent call last):
    File ".../tmp_repro_issue_6659_shape32.py", line 36, in <module>
      Init._partition_param_sec(dummy_init, param)
    File ".../deepspeed/utils/nvtx.py", line 33, in wrapped_fn
      ret_val = func(*args, **kwargs)
    File ".../deepspeed/runtime/zero/partition_parameters.py", line 1900, in _partition_param_sec
      sec_numel).copy_(one_dim_param.narrow(0, secondary_start, sec_numel))
  IndexError: start out of range (expected to be in range of [-32, 32], but got 1920)

So this appears to be a narrow edge case rather than a new regression: ZeRO++ must be enabled, the parameter must be small enough, and the DP partition world size must be large enough that some secondary ranks have no
overlap with the real tensor.

I can also add this repro detail to the PR description if that would be helpful.

[sfc-gh-truwase]

@ZyEng-Art thanks for the explanation. This is a great contribution. Thanks so much!