[Marchematics]

Add safe_get_full_model_grads for model-wide gradient access after backward and before step. ZeRO-1/2 uses one coalesced reduction per data-parallel group when supported by the backend, with per-parameter reductions as the fallback. ZeRO-3 uses the existing accessor. AutoTP/AutoEP shards remain local.

The coalesced path shares fragment assembly with safe_get_full_grad, including the upstream BF16 fixes: contribute only the owned fragment and prefer FP32 accumulation over the low-precision staging gradient. Parameters owned by another rank still contribute zero buffers to the collective.

Fixes #7644. The open-PR search for #7644 returns this PR.

Synced with upstream 0eb3a8bcc and resolved the conflicts. .github matches upstream exactly; the PR diff contains the utility export, implementation, API documentation, and tensor-fragment tests under tests/unit/v1/zero.

Validation on one Apple M4 host, macOS 15.7.3, CPU/Gloo with two worker processes, Python 3.12.12, PyTorch 2.14.1:

• DS_ACCELERATOR=cpu LOCAL_SIZE=2 DS_UNITTEST_TIMEOUT=90 python -m pytest tests/unit/v1/zero/test_zero_tensor_fragment.py -k 'safe_get_full_model_grads or TestFullModelGrads or full_gradient_reconstruction_uses_only_the_owned_fragment or test_safe_get_full_grad_preserves_fp32_accumulated_average' --torch_ver=2.14.1 -q — 17 passed. Covers ZeRO-1/2 FP32 against an ordinary PyTorch model, BF16 with FP32 accumulation on ZeRO-1/2/3, ZeRO-2 CPU optimizer offload and noncontiguous gradients, remote-owned parameters, fallback, and separate groups.
• The two BF16 regression cases failed before the compatibility fix and passed afterward.
• pre-commit run --files deepspeed/utils/__init__.py deepspeed/utils/tensor_fragment.py docs/code-docs/source/zero3.rst tests/unit/v1/zero/test_zero_tensor_fragment.py — passed.
• python ci/tests_fetcher.py --base upstream/master — selects 88 test files, including tests/unit/v1/zero/test_zero_tensor_fragment.py.

GPU, FP16, NVMe-offload, AutoTP/AutoEP, throughput, peak-memory, and convergence measurements remain. This machine has no CUDA accelerator or driver.

Updated with AI assistance. The author has confirmed the DCO statement for eb0c37c; maintainer handling of its existing DCO check is pending.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2ad7b07915

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[sfc-gh-truwase]

@Marchematics thanks for the PR. Please address the agent feedback.

[Marchematics]

Addressed the agent feedback in 5e314ed: fixed unmapped ZeRO parameters in the coalesced collective, added the regression test, and documented the public API. All commits are signed off.

[pengdurice]

Thanks for the PR, how does this work with when autoTP or autoEP is enabled.

[Marchematics]

The current implementation handles ZeRO data-parallel groups only. AutoTP and AutoEP use additional tensor/expert-parallel sharding, so this helper does not yet reconstruct their full gradients. I can add dedicated AutoTP/AutoEP coverage and reduction handling if you want this API to support those paths.

[Marchematics]

Thanks for the PR, how does this work with when autoTP or autoEP is enabled.

Would it make sense to merge this PR with the scope explicitly limited to ZeRO data-parallel groups, and handle AutoTP/AutoEP support in a follow-up PR?
I can document that limitation clearly in the API and tests. I’m also happy to continue working on the AutoTP/AutoEP extension separately so this PR doesn’t get blocked on the broader parallelism support.

[sfc-gh-truwase]

@Marchematics reduced scope is fine.

[Marchematics]

Added the reduced-scope documentation and regression coverage in ab0f4dd. The helper reconstructs ZeRO data-parallel gradients; AutoTP/AutoEP shards remain local. ZeRO-3 keeps the existing per-parameter accessor.

All 7 targeted tests passed on Apple M4 CPU/Gloo with PyTorch 2.14.1, including two-process ZeRO-1/2 FP32 comparisons against an ordinary PyTorch model. Pre-commit passed on the changed files.

The earlier eb0c37c commit deletes 30 workflows and leaves merge conflicts with master. Restoring those workflows is awaiting authorization. GPU, mixed-precision, offload, and ZeRO-3 validation remains.

[Marchematics]

Synced with upstream 0eb3a8b in ac0c8a6 and resolved the conflicts. .github now matches upstream exactly. The PR diff is back to four gradient-related files, including the tests in tests/unit/v1/zero.

The coalesced helper now shares the upstream BF16 fragment assembly and preserves FP32 accumulation. The two regression cases failed before the fix; all 17 targeted tests now pass on Apple M4 CPU/Gloo with PyTorch 2.14.1, including two-process ZeRO-1/2/3 BF16 accumulation and ZeRO-2 CPU offload. Pre-commit passed.

@sfc-gh-truwase Jiahao Zhang marchematics@gmail.com, the author of eb0c37c, has confirmed authorship, the right to contribute under the project license, compliance with DCO 1.1, and agreement to the public, lasting record of the contribution and sign-off. Please review this author confirmation and handle the existing DCO check through the maintainer process. The commit history is preserved.

The PR workflows for formatting, CPU tests, Python install smoke, CUDA op compilation, and DCO validation currently show ACTION_REQUIRED; GitHub is awaiting maintainer approval before running them. The Modal test-selection job passed, and its GPU jobs were skipped for this PR event.