Skip to content

fix: prevent division by zero in SAM gradient normalization - #1743

Open
wahid18-maqs wants to merge 1 commit into
google-deepmind:mainfrom
wahid18-maqs:fix/sam-normalize-zero-norm
Open

fix: prevent division by zero in SAM gradient normalization#1743
wahid18-maqs wants to merge 1 commit into
google-deepmind:mainfrom
wahid18-maqs:fix/sam-normalize-zero-norm

Conversation

@wahid18-maqs

Copy link
Copy Markdown

Summary

This PR improves the numerical stability of normalize() in the SAM optimizer by preventing division by zero when the gradient norm is zero.

Changes

  • Add an eps parameter to the gradient normalization step.
  • Normalize gradients using g_norm + eps instead of g_norm.
  • Add a test to verify that zero gradients do not produce NaN values.

This change preserves the existing behavior for non-zero gradients while avoiding NaN propagation in zero-gradient scenarios, such as at convergence or after aggressive gradient clipping.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant