Fix simplified layer norm fp16 overflow - #29045
Closed
SakshamKapoor2911 wants to merge 1 commit into
Closed
Conversation
Contributor
|
Please provide a regression test that reproduce overflow issue. I think the main branch has fixed the overflow issue in #28682. |
Author
|
Ah! I see that this was already caught and fixed in #28682. I was reproducing this on the older v1.24.2 release and didn't realize it had recently been fixed upstream on main. Closing this PR as a duplicate. |
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.
Description
This PR resolves a critical numerical stability bug in the
SkipSimplifiedLayerNormalizationCUDA kernel when operating underfloat16(FP16) precision.To prevent variance overflow during FP16 reduction, this change upgrades the precision of the intermediate reduction variables to
float(FP32) within the non-strict simplified layer normalization paths:layer_norm.cuh): Modified the simplified variance reduction loops to accumulate intermediate squared states usingfloatinstead ofhalf(viacub::BlockReduce<float, TPB>).skip_layer_norm_impl.cu): Updated the shared memory allocation and thread-local data registers to usefloatfor tracking the sum-of-squares.This maintains FP16 inputs/outputs and retains high performance while guaranteeing numerical stability during reduction, achieving parity with the CPU EP and the strict-mode CUDA kernels.
Motivation and Context
Fixes #29034
When executing transformer architectures in FP16 mode, input activations with large magnitudes (e.g., maximum absolute values > 7600, common in deep layer activations, scaling factors, or specific outlier features) can cause the variance calculation inside
SkipSimplifiedLayerNormalizationto overflow silently.Because the maximum representable value for FP16 is$65504$ , accumulating squared activations for a hidden dimension (e.g., $hidden_size = 768$ or $1024$ ) easily exceeds this limit. In the existing non-strict kernels, this results in an overflow to infinity ($\infty$ ), producing
NaNor silent all-zero outputs without throwing a runtime exception or crash. This is particularly problematic in production (e.g.,openai/privacy-filterdeployments), where the PII detection can silently fail by outputting zero logits.By upgrading the intermediate accumulation to FP32, the numerical bounds are dramatically increased, fully eliminating the overflow risk for these models.