Repository navigation
[WebGPU] Keep LayerNormalization arithmetic in FP32 - #32896
Merged
Jiajia Qin (qjia7) merged 1 commit intoSep 29, 2026
Merged
Conversation
Narrowing the mean to FP16 before subtraction can amplify a small FP32 rounding error into a full FP16 step. On D3D12, round-to-zero conversion causes LayerNorm_WebGPU.Float16 to exceed its existing tolerance for small normalization dimensions. Normalize full LayerNormalization in FP32 in both the ordinary and split shader paths, converting to the output type after normalization. Preserve same-type SimplifiedLayerNormalization arithmetic to maintain agreement with the fused RMSNorm kernels. The existing LayerNorm_WebGPU tests cover the affected behavior; input and output storage types and test tolerances remain unchanged.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Yang Gu (gyagp)
marked this pull request as ready for review
September 29, 2026 05:35
Contributor
Author
|
Jiajia Qin (@qjia7) PTAL |
Contributor
Author
|
/run bots |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟢 Approval recommended
The focused change is correct and covered by existing WebGPU tests across relevant shader paths.
Review effort: Balanced
Findings: None
What changed in this PR
Keeps WebGPU LayerNormalization arithmetic in FP32 while preserving existing SimplifiedLayerNormalization behavior. No actionable issues found.
Changes:
- Selects FP32 normalization for standard and mixed-type normalization.
- Preserves same-type RMSNorm arithmetic.
| File | Description |
|---|---|
onnxruntime/core/providers/webgpu/nn/layer_norm.cc |
Updates normalization arithmetic type selection. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
|
The request was received. Approve workflow runs. Thanks. |
Jiajia Qin (qjia7)
approved these changes
Sep 29, 2026
Jiajia Qin (qjia7)
enabled auto-merge (squash)
September 29, 2026 05:44
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
Keep full LayerNormalization arithmetic in FP32 until the normalized value is converted to the output type. This prevents FP16 mean conversion from amplifying small FP32 rounding errors before subtraction, and applies to both ordinary and split normalization shaders.
Preserve the existing same-type SimplifiedLayerNormalization arithmetic used by fused RMSNorm kernels. Mixed-type normalization already uses FP32. Input/output storage types and test tolerances are unchanged.
Motivation and Context
LayerNorm_WebGPU.Float16fails on a Radeon 880M with the D3D12 backend (driver32.0.22024.3004) for full normalization at shapes{2,3}and{2,6}, with and without bias.For
{2,3}, the FP32 mean is-1.7499998807907104. Narrowing it toward zero gives-1.7490234375in FP16. Subtracting this from the middle input,-1.75, and normalizing produces-0.004779815673828125instead of zero, exceeding the existing0.002tolerance. The CPU oracle normalizes before converting to the output type.WGSL permits either adjacent representable value when narrowing to a floating-point type, so the shader cannot assume round-to-nearest will recover the exact mean. A standalone D3D12 probe reproduces the conversion behavior on both AMD and WARP; enabling IEEE strictness does not resolve it.
The existing
LayerNorm_WebGPU.Float16test provides regression coverage. Its related tests also cover mixed types, FP32, scalar/vec2/vec4, split normalization, bias and statistics outputs. No additional tolerance adjustment or duplicate test is needed.Validation
onnxruntime_provider_testwith WebGPU enabled (Ninja/MSVC), on the Radeon 880M machine above.--gtest_filter=LayerNorm_WebGPU.*: all 5 tests pass, with no skips. The 4 numerical tests exercise 72 subcases with CPU fallback disabled, including 24 SimplifiedLayerNormalization subcases.--gtest_filter=LayerNormTest.*:CudaKernelTest.*LayerNorm*: all 54 broader tests pass, with no skips, usingORT_TEST_RANDOM_SEED_VALUE=2345. These tests include CPU and WebGPU coverage as supported by each case.LayerNorm_WebGPU.Float16failure was reproduced before the change. An independent D3D12 arithmetic probe also improves from 10 failing comparisons to zero across 36 outputs at the original tolerance.git diff --checkpasses.