Skip to content

Conversation

james77777778
Copy link
Contributor

@james77777778 james77777778 commented Jun 30, 2025

  • Added tests.
  • Fixed numeric stability issues in half-precision training by using float32 for computations (which is also the default in Flax)
  • Updated the fn signature of the Operation.

Ref: https://github.com/google/flax/blob/main/flax/linen/normalization.py

@codecov-commenter
Copy link

codecov-commenter commented Jun 30, 2025

Codecov Report

Attention: Patch coverage is 90.69767% with 4 lines in your changes missing coverage. Please review.

Project coverage is 82.73%. Comparing base (713172a) to head (b49e51d).
Report is 4 commits behind head on master.

Files with missing lines Patch % Lines
keras/src/ops/nn.py 90.69% 1 Missing and 3 partials ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##           master   #21438   +/-   ##
=======================================
  Coverage   82.72%   82.73%           
=======================================
  Files         565      565           
  Lines       55247    55259   +12     
  Branches     8613     8617    +4     
=======================================
+ Hits        45702    45717   +15     
- Misses       7433     7437    +4     
+ Partials     2112     2105    -7     
Flag Coverage Δ
keras 82.54% <90.69%> (+0.01%) ⬆️
keras-jax 63.35% <86.04%> (+0.01%) ⬆️
keras-numpy 58.54% <86.04%> (+0.02%) ⬆️
keras-openvino 33.78% <58.13%> (+0.01%) ⬆️
keras-tensorflow 63.77% <86.04%> (+0.02%) ⬆️
keras-torch 63.38% <83.72%> (+0.02%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@james77777778 james77777778 changed the title Update rms_normalization and layer_normalization to ensure consistent precision in computations and tests. Update rms_normalization and layer_normalization to ensure consistent precision in computations. Jun 30, 2025
- Add tests.
- Fix the numeric stability issue in half precision training.
- Update the signature of the Operation.
@gbaned gbaned requested a review from amitsrivastava78 July 2, 2025 18:14
@gbaned gbaned added this to PR Queue Jul 2, 2025
@github-project-automation github-project-automation bot moved this to Assigned Reviewer in PR Queue Jul 2, 2025
@james77777778
Copy link
Contributor Author

@fchollet
I think we should prioritize this PR because without it mixed-precision training will be unstable or broken.

Here I need to explicitly set dtype="float32" for the CLIP model to overcome the dtype issues:
keras-team/keras-hub#2316

@james77777778 james77777778 changed the title Update rms_normalization and layer_normalization to ensure consistent precision in computations. Fix rms_normalization and layer_normalization to ensure consistent precision in computations. Jul 3, 2025
Copy link
Collaborator

@fchollet fchollet left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thank you for the fixes.

@google-ml-butler google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Jul 3, 2025
@github-project-automation github-project-automation bot moved this from Assigned Reviewer to Approved by Reviewer in PR Queue Jul 3, 2025
@fchollet fchollet merged commit 6a54c36 into keras-team:master Jul 3, 2025
8 checks passed
@google-ml-butler google-ml-butler bot removed awaiting review ready to pull Ready to be merged into the codebase labels Jul 3, 2025
@github-project-automation github-project-automation bot moved this from Approved by Reviewer to Merged in PR Queue Jul 3, 2025
@james77777778 james77777778 deleted the fix-rms-layer-norm branch July 4, 2025 01:10
Mohamed-Ashraf273 pushed a commit to Mohamed-Ashraf273/keras that referenced this pull request Jul 8, 2025
- Add tests.
- Fix the numeric stability issue in half precision training.
- Update the signature of the Operation.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
Status: Merged
Development

Successfully merging this pull request may close these issues.

4 participants