Add an opt-in ZeRO-1/2 gradient norm fast path

[yh0903]

Summary

• add zero_optimization.compute_grad_norm, defaulting to true
• allow ZeRO Stage 1/2 GPU optimizers to skip an unused global gradient norm while preserving overflow detection
• skip the gradient multiply when clipping is disabled and the effective loss scale is exactly 1
• return None from get_global_grad_norm() when norm computation is explicitly disabled
• fail fast for gradient clipping, optimizer offload, ZenFlow, unsupported ZeRO stages, checkpoint-restored clipping, and the dedicated ZeRO-1 BF16 optimizer with FP32 gradient accumulation

Safety contract

Finite/overflow checking is unchanged. This does not remove the required non-finite scan; it only removes work with no mathematical consumer.

The default configuration and numerical behavior are unchanged. Identity unscaling is skipped automatically when clipping is disabled and the effective loss scale is exactly 1. Gradient-norm computation remains enabled by default because its cached result is externally observable through get_global_grad_norm().

Performance

Qwen3-30B-A3B, 48 layers, EP16 on 2x8 H100, sequence length 1024, BF16, ZeRO-1, activation checkpointing enabled:

Metric	Baseline	Fast path	Change
Median step	987.92 ms	966.67 ms	-2.2%
Optimizer update	84.93 ms	59.82 ms	-29.6%
Global norm	20.30 ms	0 ms	removed
Identity unscale	5.09 ms	0.03 ms	removed

The paired block used a benchmark-only control that restores the pre-change multiply-by-one path. Peak allocated and reserved memory were unchanged. Two additional rotated blocks isolating the opt-in norm change were both positive (+3.3% and +0.6% E2E); their optimizer-phase saving was stable at about 20 ms, while full-step variance remained larger.

Testing Done

• rebased onto current master
• H100 targeted suite: 20 passed, covering Stage 1/2 update parity, overflow skip behavior, get_global_grad_norm(), clipping/offload/ZenFlow guards, checkpoint clipping, identity/non-identity scaling, and the ZeRO-1 BF16 optimizer with FP32 gradient accumulation
• changed-file pre-commit hooks passed; flake8 5.0.4 was run separately under Python 3.13 because its pyflakes dependency is incompatible with Python 3.14
• current-head multi-GPU H100 suite: 20 passed, 0 failed, 0 skipped
• full repository CI still requires maintainer approval to run on the fork PR

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 349e8c91fe

ℹ️ 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".

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-03T08:37:30.299076Z	349e8c9	Draft marked ready

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[tohtana]

Hi @yh0903,
Thank you for the improvement! Sorry for the delay of review.
This looks good to me.