Fix FSDP2 gradient norm/max reporting and clipping - #2900
Merged
Conversation
FSDP2 gradients are DTensors, so rank-local maxima can under-report the logical gradient, empty local shards must still join collectives, and reporting must not receive device tensors. The training loop also measured a max on every accumulation microstep even when norm clipping would immediately recompute and overwrite it. - In Trainer._max_grad_value(), measure non-empty local shards, contribute negative infinity for empty DTensor shards, stage CPU scalars on the mesh device, and all-reduce MAX once across the mesh. - At synchronized optimizer steps, reuse the norm returned by clipping in norm mode and run the separate max path only when clipping does not already provide the reported value. - Preserve existing FSDP1, FSDP2, non-FSDP, optimizer-exclusion, and FSDP value-clipping policies while avoiding duplicate gradient traversal and collectives. - Normalize norm and max values through one scalar path for logs, progress, webhooks, and regularisation metrics.
Restore the clone_norm_value keyword used by existing callers and keep value-mode grad_absmax entries as the original metric object. Continue converting norm metrics to scalars only for progress payloads that explicitly request cloning. This keeps the FSDP2 mesh-wide max reduction and synchronized clipping changes while preserving the pre-existing logging API and tensor identity behavior exercised by the upstream trainer suite.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adjusts gradient metric computation and reporting during training—especially under FSDP2—so that gradient max/norm measurements are consistent across ranks, avoid unsafe empty/CPU-offload edge cases, and reduce redundant gradient traversals when clipping is enabled.
Changes:
- Updates gradient metric reporting to gate metric computation/clipping on
accelerator.sync_gradients, and to reuseclip_grad_norm_()’s returned pre-clipping norm in norm mode. - Reworks
Trainer._max_grad_value()to correctly compute an infinity-norm / absolute-max style metric for DTensor gradients, including mesh-wideall_reduce(MAX)and safe handling of empty local shards. - Normalizes metric scalar conversion for structured progress payload construction while preserving existing accelerator logging behavior.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+567
to
571
| grad_value = self._normalize_metric_value(self.grad_norm) | ||
| if grad_value is None: | ||
| return | ||
| target_logs[f"{prefix}grad_norm"] = grad_value | ||
| elif ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fix FSDP2 gradient norm/max reporting and clipping
Summary
gradient; empty local shards and CPU-offloaded gradients were unsafe cases.
when norm clipping would recompute and overwrite it at the optimizer step.
paths and could receive inconsistent metric types.
across the FSDP2 mesh, reuse the norm returned by clipping, and preserve the
existing logging API while normalizing structured progress payloads.
Changes
Trainer._max_grad_value(): compute the infinity norm from non-empty localgradients, contribute
-inffor an empty DTensor shard, stage CPU scalars onthe mesh device, and run one mesh-group
all_reduce(MAX).accelerator.sync_gradientsis true; useclip_grad_norm_()'s return value fornorm mode instead of running a separate max traversal first.
torch.nn.utils.clip_grad_norm_(),FSDP1 on the wrapped model method, and non-FSDP clipping through Accelerate.
Existing optimizer exclusions and the FSDP value-clipping limitation remain
unchanged.
_update_grad_metrics()and progress payloads: retain tensor objects for theexisting accelerator log path, restore the
clone_norm_valuecompatibilityargument, and normalize norm values only for structured progress/webhook
serialization.
clip_grad_norm_()Result
is used.
norm is reported directly.
-inf; CPU-offloaded scalarvalues are staged on the mesh device before reduction.
collectives.
0.8, rank 1 has3.2, and rank 2 hasan empty shard, their contributions are
0.8,3.2, and-inf; all ranksreceive and report
3.2afterall_reduce(MAX).metric contracts, norm/value clipping, empty DTensor shards, CPU-offload
staging, and accumulation gating. A real two-GPU FSDP2 run also verified the
mesh-wide max and post-clipping global norm.
Environment
3.13.142.13.0+cu1301.14.00.39.05.13.1590.48.01