Add coalesced full gradient access

[Marchematics]

Add safe_get_full_model_grads to collect model gradients with coalesced ZeRO-1/2 communication.

This replaces one collective per parameter with one coalesced collective per data-parallel group.

Fixes #7644

Tests: targeted gradient API test, pre-commit.

[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.