Repository navigation
Conversation
The FNO encoders built their coordinate grid with torch.linspace, and Dynamo specializes linspace's step count. A model compiled with dynamic=True therefore recompiled for every new input resolution (4 sizes -> 4 graphs) until it hit the recompile limit. Build the grid with torch.arange(n) / max(n - 1, 1) instead, which keeps the size symbolic. Values match linspace to within 1 float32 ULP and the saved reference outputs still pass. Closes NVIDIA#2207 Signed-off-by: aarochu <aaronchu022@gmail.com>
CODEOWNERS review mapCurrent for commit ⏳ @loliverhennigh — 1 file(s)
No CODEOWNER
Comment |
|
[Medium impact] The change reaches FNO runs across one through four dimensions. The PR appears safe to merge, with a non-blocking request to test grid values directly. Findings
SummaryReplaces
Acknowledged behavior: aarochu explicitly accepts small float32 rounding differences from the replacement formula. Reviews (1) · Last reviewed commit: "Keep FNO grid size dynamic under torch.c..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| for size in (6, 7, 8): | ||
| invar = torch.randn(2, 2, *([size] * dimension), device=device) | ||
| torch.testing.assert_close(compiled(invar), model(invar)) |
There was a problem hiding this comment.
Both calls in this comparison use _unit_grid, and the test only tries sizes 6, 7, and 8. An incorrect grid could therefore pass, including a broken [0.] result for size 1. Add a focused test against torch.linspace for sizes 1, 2, and several larger sizes, with a tight float32 tolerance and exact endpoint checks.
PhysicsNeMo Pull Request
Description
Closes #2207.
If you compile FNO with
torch.compile(dynamic=True)and feed it different grid sizes, it recompiles every time. That's the opposite of what you want from a resolution-independent model. The cause is the coordinate grid in the encoders'_meshgrid: it usestorch.linspace(0, 1, size), and Dynamo specializes on linspace's step count. WithTORCH_LOGS=dynamicyou can see it guardingEq(s53, 16)right at that line, and thentensor 'x' size mismatch at index 2on the next size.The fix adds a small
_unit_gridhelper that builds the same grid astorch.arange(size) / max(size - 1, 1)(usingtorch.sym_maxso the size stays symbolic). All four encoders (1D–4D) now use it. Nothing else changes:linspaceto within 1 float32 ULP (I checked every size from 1 to 512). The endpoints are exactly 0 and 1, andsize=1still gives[0.]..pthreference outputs still pass, as do the JIT/AMP/CUDA-graph optim tests and the xFNO tests.Before/after for the repro from the issue (CPU, torch 2.13, sizes 16/24/32/40 for 1D/2D, smaller for 3D/4D):
I ran it with the
eager,aot_eagerandinductorbackends. Compiled outputs match eager in every case. The issue mentions 3D hitting an InductorValueRangeErroron CUDA. I don't have a GPU on this machine, so I couldn't check that path, but 3D compiles to a single graph with Inductor on CPU. If it still trips on CUDA, that looks like an Inductor bug rather than something on our side.I also added
test_fno_compile_dynamic_resolution. It compiles withdynamic=Trueon the eager backend, runs three sizes, and asserts that there's one graph and the outputs match eager. It runs on CPU too, unlike the existing fullgraph check, which is GPU-only and never changes size. On main it fails for all four dimensions.Checklist
Dependencies
None.