[NVIDIA] Update the algorithm to compute fp8 scales

[kaixih]

This pull request implements updates to the algorithm for computing new FP8 scales, in line with the design outlined in the following documentation: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/jax.html#transformer_engine.jax.update_fp8_metas.

cc. @wenscarl @reedwm @lukaszlew @levskaya

[levskaya]

The change looks OK, but mypy is now yielding a type error - would you mind double-checking if that's finding something real, or is it spurious?

[codecov-commenter]

Codecov Report

Merging #3441 (3a0e87a) into main (5557649) will decrease coverage by 0.03%.
Report is 2 commits behind head on main.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #3441      +/-   ##
==========================================
- Coverage   83.68%   83.65%   -0.03%     
==========================================
  Files          56       56              
  Lines        6808     6797      -11     
==========================================
- Hits         5697     5686      -11     
  Misses       1111     1111              

Files	Coverage Δ
flax/linen/__init__.py	100.00% <100.00%> (ø)
flax/linen/fp8_ops.py	100.00% <100.00%> (ø)
flax/training/train_state.py	100.00% <100.00%> (ø)

[kaixih]

@levskaya thank you for bringing this to my attention. I've removed the function return value annotation that was causing the previous mypy check to trigger. I don't believe it's necessary because this custom operation for FP8 is not configured to the layers, such as DenseGeneral, by default. There won't be any other type checking tracing back to this operation.

[kaixih]

@levskaya Any updates? I see this PR has stayed in "pull ready" status for a couple of days.

[levskaya]

@kaixih - super sorry for the delay, not sure why it didn't go through earlier! It should have merged last night.