Repository navigation
Commit a5d7261
[Inductor] Fix NaN output in LayerNorm CPU by guarding Welford variance. (pytorch#173989)
Running torch.compile on CPU with float16 precision causes LayerNorm to produce NaN outputs when the input contains Inf values. This happens because the Welford variance calculation in lowering.py hits a numerical collapse ($\infty - \infty$) in the $M2$ accumulator, which then propagates as a NaN.This PR adds an ops.where guard in var_mean_welford_ to catch these non-finite results and reset them to 0.0. This ensures that even if the input overflows the precision range, the LayerNorm output remains numerically stable, matching the behavior of the CUDA backend and Eager mode.
**Fixes**: pytorch#173885
Pull Request resolved: pytorch#173989
Approved by: https://github.com/jansel1 parent 48f9fce commit a5d7261
4 files changed
Lines changed: 49 additions & 2 deletions
File tree
- test/inductor
- torch
- _inductor
- runtime
- csrc/inductor
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
6624 | 6624 | | |
6625 | 6625 | | |
6626 | 6626 | | |
| 6627 | + | |
| 6628 | + | |
| 6629 | + | |
| 6630 | + | |
| 6631 | + | |
| 6632 | + | |
| 6633 | + | |
| 6634 | + | |
| 6635 | + | |
| 6636 | + | |
| 6637 | + | |
| 6638 | + | |
| 6639 | + | |
| 6640 | + | |
| 6641 | + | |
| 6642 | + | |
| 6643 | + | |
| 6644 | + | |
| 6645 | + | |
| 6646 | + | |
| 6647 | + | |
| 6648 | + | |
| 6649 | + | |
| 6650 | + | |
| 6651 | + | |
| 6652 | + | |
| 6653 | + | |
| 6654 | + | |
| 6655 | + | |
6627 | 6656 | | |
6628 | 6657 | | |
6629 | 6658 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1257 | 1257 | | |
1258 | 1258 | | |
1259 | 1259 | | |
1260 | | - | |
| 1260 | + | |
| 1261 | + | |
| 1262 | + | |
| 1263 | + | |
| 1264 | + | |
| 1265 | + | |
| 1266 | + | |
| 1267 | + | |
1261 | 1268 | | |
1262 | 1269 | | |
1263 | 1270 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
262 | 262 | | |
263 | 263 | | |
264 | 264 | | |
265 | | - | |
| 265 | + | |
| 266 | + | |
| 267 | + | |
266 | 268 | | |
267 | 269 | | |
268 | 270 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
192 | 192 | | |
193 | 193 | | |
194 | 194 | | |
| 195 | + | |
| 196 | + | |
195 | 197 | | |
| 198 | + | |
| 199 | + | |
| 200 | + | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
196 | 205 | | |
197 | 206 | | |
198 | 207 | | |
| |||
0 commit comments