Skip to content

Fix simplified layer norm fp16 overflow - #29045

Closed
SakshamKapoor2911 wants to merge 1 commit into
microsoft:mainfrom
SakshamKapoor2911:fix-simplified-layer-norm-fp16-overflow
Closed

Fix simplified layer norm fp16 overflow#29045
SakshamKapoor2911 wants to merge 1 commit into
microsoft:mainfrom
SakshamKapoor2911:fix-simplified-layer-norm-fp16-overflow

Conversation

@SakshamKapoor2911

Copy link
Copy Markdown

Description

This PR resolves a critical numerical stability bug in the SkipSimplifiedLayerNormalization CUDA kernel when operating under float16 (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:

  1. Precision Upgrade in Kernel Helpers (layer_norm.cuh): Modified the simplified variance reduction loops to accumulate intermediate squared states using float instead of half (via cub::BlockReduce<float, TPB>).
  2. Precision Upgrade in Host/Device Implementations (skip_layer_norm_impl.cu): Updated the shared memory allocation and thread-local data registers to use float for 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 SkipSimplifiedLayerNormalization to 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 NaN or silent all-zero outputs without throwing a runtime exception or crash. This is particularly problematic in production (e.g., openai/privacy-filter deployments), 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.

@tianleiwu

Copy link
Copy Markdown
Contributor

Please provide a regression test that reproduce overflow issue.

I think the main branch has fixed the overflow issue in #28682.

@SakshamKapoor2911

Copy link
Copy Markdown
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.

@SakshamKapoor2911
SakshamKapoor2911 deleted the fix-simplified-layer-norm-fp16-overflow branch June 14, 2026 05:02
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.

[CUDA] SkipSimplifiedLayerNormalization fp16 silently returns all zeros for large-magnitude inputs (variance accumulated in fp16)

2 participants