Repository navigation
FNO: run spectral convolutions in float32 under autocast - #2150
GeorgeLiang3 wants to merge 2 commits into
Conversation
CODEOWNERS review mapCurrent for commit ⏳ @ktangsali — 1 file(s)
⏳ @loliverhennigh — 1 file(s)
No CODEOWNER
Comment |
|
[Medium impact] The changes affect spectral layers shared by multiple model workflows. The PR appears safe to merge, with a non-blocking improvement needed in the training test. Findings
SummaryThis PR lets FNO use mixed precision while keeping spectral convolutions in float32.
GeorgeLiang3 explicitly identifies default autocast in Reviews (1) · Last reviewed commit: "fix(fno): run spectral convolutions in f..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| scaler.scale(loss).backward() | ||
| scaler.step(optimizer) | ||
| scaler.update() | ||
| assert torch.isfinite(loss) | ||
| assert all(torch.isfinite(p).all() for p in model.parameters()) |
There was a problem hiding this comment.
test_fno_autocast_training_step can pass without updating any weights. For float16, GradScaler.step skips the update when gradients contain infinities or NaNs, leaving the initial weights finite. Checking only the loss and weights therefore misses a broken backward pass. Check that gradients are finite after unscaling and that at least one weight changes.
There was a problem hiding this comment.
The test now unscales the gradients before scaler.step and checks that every trainable parameter has a finite gradient. After the step it checks that every trainable parameter changed. It only checks parameters with requires_grad, because Conv2dFCLayer sets requires_grad=False on its bias. I confirmed the updated test fails when a gradient is forced to inf or NaN, and when a parameter receives no gradient or an all-zero gradient; the previous version passed in all of these cases under float16.
|
Hi @ktangsali @loliverhennigh , could you review this please? thanks |
FNO was marked amp=False, so validate_amp skipped it. Under float16 autocast, SpectralConv1d-4d failed: cuFFT accepts half precision only for power-of-two sizes, and the complex-half mode multiplication is not implemented. torch 2.13 also rejects bfloat16 FFTs on CUDA. SpectralConv1d-4d now upcast float16 and bfloat16 inputs to float32 and run the FFT, mode multiplication and inverse FFT with autocast disabled, when autocast is active for the input's device. Layers around them keep running under autocast. Without autocast the forward is unchanged. FNO is marked AMP-capable. Also: - build the output spectrum in the input spectrum's complex dtype, so SpectralConv1d returns float64 for float64 inputs - reset_parameters keeps the weights' device and dtype Adds test/nn/module/test_spectral_conv.py with direct tests for the spectral convolution layers. Signed-off-by: GeorgeLiang3 <liangzhouji@gmail.com>
…g test GradScaler skips the optimizer step when gradients are not finite, so a broken backward pass under float16 could leave the weights untouched and still pass. The test now unscales and checks the gradients of every trainable parameter before the step, and checks that each one changed after it. Signed-off-by: GeorgeLiang3 <liangzhouji@gmail.com>
31cca4f to
92f5082
Compare
|
@melo-gonzo this moves the float32 spectral-convolution guard from xDeepONet (#1738) into |
PhysicsNeMo Pull Request
Description
Closes #2149.
FNOis markedamp=Falsein itsMetaData, sovalidate_ampreturns early and the FNO AMP test checks nothing. Under float16 autocast on CUDA,SpectralConv1d-4dfail. With a non-power-of-two grid, cuFFT raises "cuFFT only supports dimensions whose sizes are powers of two when computing in half precision". With a power-of-two grid, the complex-half mode multiplication raises"baddbmm_cuda" not implemented for 'ComplexHalf'. On torch 2.13, the minimum version, CUDA FFTs also reject bfloat16 (Unsupported dtype BFloat16), so bfloat16 autocast fails too.This PR adds a small decorator on
SpectralConv1d-4d.forward. When autocast is active for the input's device, it upcasts float16 and bfloat16 inputs to float32 and runs the FFT, mode multiplication and inverse FFT with autocast disabled. Layers around the spectral convolutions keep running under autocast. Without autocast,forwardis unchanged, so full-precision outputs and the existing golden files are unchanged. This is the same approach xDeepONet uses inSpatialBranch._spectral(#1737, #1738), placed in the shared layers as suggested in the design note of #1737.FNOis now marked AMP-capable. Note thatStaticCaptureTrainingandStaticCaptureEvaluateNoGraddefault touse_amp=True, so they now enable autocast for FNO. The in-repo FNO examples passuse_amp=Falseand are unaffected.Two small fixes in the same layers:
torch.cfloat, soSpectralConv1dreturns float64 for float64 inputs. 2D-4D already did.reset_parametersusestorch.rand_like, so it keeps the weights' device and dtype. Before, calling it after.to("cuda")or.double()moved the weights back to CPU float32. Default initialization is unchanged under a fixed seed.New tests are in
test/nn/module/test_spectral_conv.py. Until now onlySpectralConv4dhad a direct test, intest_nd_conv_layers.py. The new tests cover:2 * modes == sizeand corners that overlap on each axis;GradScalerfor float16;torch.compile(fullgraph=True), with and without autocast;reset_parameterskeeping device and dtype.Against
main, 26 of the new tests fail on CUDA: reduced precision, the FNO autocast step, compile with autocast, 1D float64, andreset_parameters.Validation on torch 2.13.0+cu130 and 2.14.1+cu130:
test/nn/module,test/models/fno,afno,sfnoandexperimental/models/xdeeponetpass: 1116 passed, 150 skipped (optional dependencies and ONNX).test/models/fnopass.pre-commit run --all-filespasses.fullgraphforward and backward for FNO under float16 and bfloat16 autocast run without graph breaks. CI's compile check uses a no-op backend.Checklist
Dependencies
None.