onnxruntime
e75f6c89 - Enable LayerNormFusion for the WebGPU EP to fix fp16 inference correctness (#32294)

Commit
14 days ago
Enable LayerNormFusion for the WebGPU EP to fix fp16 inference correctness (#32294) ### Description The Level2 `LayerNormFusion` transformer isn't registered for `kWebGpuExecutionProvider`, so decomposed LayerNorm subgraphs never get fused on WebGPU. This is a **correctness issue**, not just a missed optimization. For models exported before opset 17, `LayerNormalization` doesn't exist as a standard op, so the Level1 fusion bails out (`layer_norm_fusion.cc` L234). Level2 is the only remaining chance, and it skips WebGPU because of the provider allowlist. The chain then runs elementwise in fp16, where `Pow(x - mean, 2)` overflows in the FFN-output LayerNorms of a typical BERT encoder. No exception, no NaN, just silently wrong values. This looks like an oversight: - `GeluFusion` on the line right above already uses `cpu_acl_cuda_dml_webgpu_eps`. - `SkipLayerNormFusion` and `BiasSkipLayerNormFusion` already include WebGPU. - The WebGPU kernel itself handles fp16 fine (`layer-norm.ts` casts to f32 internally). Note this only affects exports below opset 17: from opset 17 the Level1 instance fuses before EP partitioning, where no provider allowlist applies, and the Level2 instance bails out instead. That is likely why it has gone unnoticed. Refs #31626 — full analysis and measurements are in [this comment](https://github.com/microsoft/onnxruntime/issues/31626#issuecomment-5438068663). ### Change `LayerNormFusion`: `cpu_acl_cuda_dml_eps` → `cpu_acl_cuda_dml_webgpu_eps` (constant already exists). The fusion now also preserves scalar epsilon initializers in supported tensor types, including FLOAT16, instead of replacing non-FLOAT values with the default epsilon. `SimplifiedLayerNormFusion` on the next line is left as-is. It's not needed for this pattern (`LayerNormFusion` matches first), and it has its own `is_gpu_ep` check that only treats CUDA as a GPU, which would want separate consideration. ### Results `Xenova/bge-small-en-v1.5` (opset 11, 25 decomposed LayerNorms), onnxruntime-web 1.29.0, Apple M1 (Metal-3). Reference = fp32 weights on CPU EP; mean pooling + L2 normalisation. | graph | EP | cosine vs fp32 reference | |---|---|---| | official `model_fp16.onnx` (decomposed) | webgpu | 0.63 – 0.83 | | same graph, chains hand-replaced with `LayerNormalization` | webgpu | **1.0000** | | same graph, fused to `SimplifiedLayerNormalization` by CPU EP optimizer | webgpu | **1.0000** | | official `model.onnx` (fp32) | webgpu | **1.0000** | | official `model_fp16.onnx` | wasm (CPU) | **1.0000** | (Cosine to four decimals, not bit-exact.) The fp32 graph run through the CPU optimizer at `ORT_ENABLE_EXTENDED` fuses all 25 chains. The fp16 graph has the same topology (same op counts, no `Cast` nodes), and `tensor(float16)` is in `supported_data_types` (L17). The only fp16-specific gates in the file (L561, L820) are CPU-only. So once WebGPU is on the allowlist, these chains should fuse. The regression fixture uses a non-default FLOAT16 epsilon (`1e-4`) and asserts that the fused `LayerNormalization` node preserves its value.
Author
Parents
Loading