From 05ae2edbd7318581893122383717ebca80d18e97 Mon Sep 17 00:00:00 2001 From: GeorgeLiang3 Date: Tue, 20 Oct 2026 23:20:11 -0600 Subject: [PATCH 1/2] fix(fno): run spectral convolutions in float32 under autocast 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 --- CHANGELOG.md | 12 + physicsnemo/models/fno/fno.py | 2 +- physicsnemo/nn/module/spectral_layers.py | 89 ++++++-- test/nn/module/test_spectral_conv.py | 265 +++++++++++++++++++++++ 4 files changed, 347 insertions(+), 21 deletions(-) create mode 100644 test/nn/module/test_spectral_conv.py diff --git a/CHANGELOG.md b/CHANGELOG.md index dc4e708b8d..62d06363a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -133,6 +133,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `DistributedGraph` and GraphCast's distributed mode) no longer loops over destination nodes in Python. A 1-million-node grid now takes well under a second instead of half a minute or more. +- `FNO` supports automatic mixed precision. `SpectralConv1d`-`4d` upcast + float16 and bfloat16 inputs to float32 and run their FFTs and mode + multiplication with autocast disabled, so float16 autocast no longer fails + with a cuFFT or `ComplexHalf` error, and bfloat16 autocast works on PyTorch + versions whose FFT does not accept bfloat16. Layers outside the spectral + convolutions still run under autocast. `FNO` is now marked as AMP-capable, + so `StaticCaptureTraining` and `StaticCaptureEvaluateNoGrad` enable autocast + for it by default; pass `use_amp=False` to keep full precision. - `ClusterTree` builds tighter trees by default (`split="morton"`), so Barnes-Hut plans need 2-4x fewer interactions at the same `theta`. `split="midpoint"` gives the previous trees. @@ -422,6 +430,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 DoMINO's `sample_points_on_mesh` now works on meshes with over 2^24 faces. - `MeshDataset.close()` no longer raises `AttributeError` with a `MeshReader` or `DomainMeshReader`. Both readers now have a `close()` method. +- `SpectralConv1d`-`4d` build the output spectrum in the input spectrum's + complex dtype, so `SpectralConv1d` no longer returns float32 for float64 + inputs, and `reset_parameters` keeps the weights' device and dtype instead + of moving them to CPU float32. - `GLOBE` accepts `boundary_meshes` as a TensorDict, as the DrivAerML example passes them. Before, every forward pass of that example failed. - Under `torch.compile`, GLOBE no longer recompiles on its second training diff --git a/physicsnemo/models/fno/fno.py b/physicsnemo/models/fno/fno.py index aed088e5d8..4e4639c5ee 100644 --- a/physicsnemo/models/fno/fno.py +++ b/physicsnemo/models/fno/fno.py @@ -46,7 +46,7 @@ class MetaData(ModelMetaData): # Optimization jit: bool = False cuda_graphs: bool = True - amp: bool = False + amp: bool = True # Inference onnx_cpu: bool = False onnx_gpu: bool = False diff --git a/physicsnemo/nn/module/spectral_layers.py b/physicsnemo/nn/module/spectral_layers.py index 31dec5ee2d..2ff24b57f4 100644 --- a/physicsnemo/nn/module/spectral_layers.py +++ b/physicsnemo/nn/module/spectral_layers.py @@ -14,7 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Tuple +import functools +from typing import Callable, List, Tuple import numpy as np import torch @@ -24,6 +25,50 @@ from torch import Tensor +def _fp32_under_autocast(forward: Callable) -> Callable: + r"""Run a spectral convolution ``forward`` in float32 when autocast is active. + + A float16 input yields a complex-half spectrum, which cuFFT supports only + for power-of-two signal sizes and which the mode multiplication does not + support at all. A bfloat16 input is not accepted by every FFT backend, and + where it is, the transform is not guaranteed to be computed in float32. So + when autocast is active for the input's device, float16 and bfloat16 + inputs are upcast to float32 and the FFT, mode multiplication and inverse + FFT run with autocast disabled. Layers around the spectral convolution + keep running under autocast. Without autocast, the wrapped ``forward`` is + called unchanged. + + .. code-block:: python + + def forward(self, x: torch.Tensor) -> torch.Tensor: + ... + + Parameters + ---------- + forward : Callable[[torch.nn.Module, torch.Tensor], torch.Tensor] + ``forward`` method of a spectral convolution. See code block above. + + Returns + ------- + Callable[[torch.nn.Module, torch.Tensor], torch.Tensor] + Wrapped ``forward`` with the same signature. + """ + + @functools.wraps(forward) + def wrapper(self, x: Tensor) -> Tensor: + device_type = x.device.type + if torch.amp.is_autocast_available(device_type) and torch.is_autocast_enabled( + device_type + ): + if x.dtype in (torch.float16, torch.bfloat16): + x = x.float() + with torch.autocast(device_type=device_type, enabled=False): + return forward(self, x) + return forward(self, x) + + return wrapper + + class SpectralConv1d(nn.Module): """1D Fourier layer. It does FFT, linear transform, and Inverse FFT. @@ -75,6 +120,7 @@ def compl_mul1d( cweights = torch.view_as_complex(weights) return torch.einsum("bix,iox->box", input, cweights) + @_fp32_under_autocast def forward( self, x: Float[Tensor, "batch in_channels x"] ) -> Float[Tensor, "batch out_channels x"]: @@ -87,7 +133,7 @@ def forward( self.out_channels, x.size(-1) // 2 + 1, device=x.device, - dtype=torch.cfloat, + dtype=x_ft.dtype, ) # (batch, out_channels, x//2+1) complex out_ft[:, :, : self.modes1] = self.compl_mul1d( x_ft[:, :, : self.modes1], @@ -100,7 +146,7 @@ def forward( def reset_parameters(self): """Reset spectral weights with distribution scale*U(0,1)""" - self.weights1.data = self.scale * torch.rand(self.weights1.data.shape) + self.weights1.data = self.scale * torch.rand_like(self.weights1) class SpectralConv2d(nn.Module): @@ -161,6 +207,7 @@ def compl_mul2d( cweights = torch.view_as_complex(weights) return torch.einsum("bixy,ioxy->boxy", input, cweights) + @_fp32_under_autocast def forward( self, x: Float[Tensor, "batch in_channels h w"] ) -> Float[Tensor, "batch out_channels h w"]: @@ -169,7 +216,7 @@ def forward( # Initialize output in frequency space out_ft = torch.zeros( - x.size(0), self.out_channels, h, w, dtype=torch.cfloat, device=x.device + x.size(0), self.out_channels, h, w, dtype=x_ft.dtype, device=x.device ) # (batch, out_channels, h, w) complex # Accumulate Fourier modes. Use .contiguous() on sliced complex tensors and @@ -196,8 +243,8 @@ def forward( def reset_parameters(self): """Reset spectral weights with distribution scale*U(0,1)""" - self.weights1.data = self.scale * torch.rand(self.weights1.data.shape) - self.weights2.data = self.scale * torch.rand(self.weights2.data.shape) + self.weights1.data = self.scale * torch.rand_like(self.weights1) + self.weights2.data = self.scale * torch.rand_like(self.weights2) class SpectralConv3d(nn.Module): @@ -278,6 +325,7 @@ def compl_mul3d( cweights = torch.view_as_complex(weights) return torch.einsum("bixyz,ioxyz->boxyz", input, cweights) + @_fp32_under_autocast def forward( self, x: Float[Tensor, "batch in_channels d1 d2 d3"] ) -> Float[Tensor, "batch out_channels d1 d2 d3"]: @@ -293,7 +341,7 @@ def forward( d1, d2, d3, - dtype=torch.cfloat, + dtype=x_ft.dtype, device=x.device, ) # (batch, out_channels, d1, d2, d3) complex @@ -340,10 +388,10 @@ def forward( def reset_parameters(self): """Reset spectral weights with distribution scale*U(0,1)""" - self.weights1.data = self.scale * torch.rand(self.weights1.data.shape) - self.weights2.data = self.scale * torch.rand(self.weights2.data.shape) - self.weights3.data = self.scale * torch.rand(self.weights3.data.shape) - self.weights4.data = self.scale * torch.rand(self.weights4.data.shape) + self.weights1.data = self.scale * torch.rand_like(self.weights1) + self.weights2.data = self.scale * torch.rand_like(self.weights2) + self.weights3.data = self.scale * torch.rand_like(self.weights3) + self.weights4.data = self.scale * torch.rand_like(self.weights4) class SpectralConv4d(nn.Module): @@ -501,6 +549,7 @@ def compl_mul4d( cweights = torch.view_as_complex(weights) return torch.einsum("bixyzt,ioxyzt->boxyzt", input, cweights) + @_fp32_under_autocast def forward( self, x: Float[Tensor, "batch in_channels d1 d2 d3 d4"] ) -> Float[Tensor, "batch out_channels d1 d2 d3 d4"]: @@ -522,7 +571,7 @@ def forward( d2, d3, d4, - dtype=torch.cfloat, + dtype=x_ft.dtype, device=x.device, ) # (batch, out_channels, d1, d2, d3, d4) complex @@ -622,14 +671,14 @@ def forward( def reset_parameters(self): """Reset spectral weights with distribution scale*U(0,1)""" - self.weights1.data = self.scale * torch.rand(self.weights1.data.shape) - self.weights2.data = self.scale * torch.rand(self.weights2.data.shape) - self.weights3.data = self.scale * torch.rand(self.weights3.data.shape) - self.weights4.data = self.scale * torch.rand(self.weights4.data.shape) - self.weights5.data = self.scale * torch.rand(self.weights5.data.shape) - self.weights6.data = self.scale * torch.rand(self.weights6.data.shape) - self.weights7.data = self.scale * torch.rand(self.weights7.data.shape) - self.weights8.data = self.scale * torch.rand(self.weights8.data.shape) + self.weights1.data = self.scale * torch.rand_like(self.weights1) + self.weights2.data = self.scale * torch.rand_like(self.weights2) + self.weights3.data = self.scale * torch.rand_like(self.weights3) + self.weights4.data = self.scale * torch.rand_like(self.weights4) + self.weights5.data = self.scale * torch.rand_like(self.weights5) + self.weights6.data = self.scale * torch.rand_like(self.weights6) + self.weights7.data = self.scale * torch.rand_like(self.weights7) + self.weights8.data = self.scale * torch.rand_like(self.weights8) # ========================================== diff --git a/test/nn/module/test_spectral_conv.py b/test/nn/module/test_spectral_conv.py new file mode 100644 index 0000000000..11a1b654e8 --- /dev/null +++ b/test/nn/module/test_spectral_conv.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch + +from physicsnemo.models.fno import FNO +from physicsnemo.nn.module.spectral_layers import ( + SpectralConv1d, + SpectralConv2d, + SpectralConv3d, + SpectralConv4d, +) + +IN_CHANNELS = 3 +OUT_CHANNELS = 4 +CONV_CLASSES = { + 1: SpectralConv1d, + 2: SpectralConv2d, + 3: SpectralConv3d, + 4: SpectralConv4d, +} +CORNERS = { + 1: [()], + 2: [(False,), (True,)], + 3: [(False, False), (True, False), (False, True), (True, True)], + 4: [ + (False, False, False), + (True, False, False), + (False, True, False), + (False, False, True), + (True, True, False), + (True, False, True), + (False, True, True), + (True, True, True), + ], +} + + +def _make_conv(dimension: int, modes=None) -> torch.nn.Module: + modes = modes if modes is not None else [2] * dimension + return CONV_CLASSES[dimension](IN_CHANNELS, OUT_CHANNELS, *modes) + + +def _make_input(dimension: int, size: int, device, dtype=torch.float32): + return torch.randn( + 2, IN_CHANNELS, *([size] * dimension), device=device, dtype=dtype + ) + + +AMP_DTYPES = [torch.float16, torch.bfloat16] + + +def _skip_unsupported_amp(device, amp_dtype): + # Only float16 is skipped on CPU, where these tests use bfloat16 autocast + if not str(device).startswith("cuda") and amp_dtype == torch.float16: + pytest.skip("CPU autocast is tested with bfloat16 only") + + +def _reference_forward(conv: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: + # Spectral convolution written with slice accumulation into a full-size + # spectrum, independent of the layer implementation. + dimension = x.ndim - 2 + modes = [getattr(conv, f"modes{i + 1}") for i in range(dimension)] + dims = list(range(-dimension, 0)) + x_ft = torch.fft.rfftn(x, dim=dims) + out_ft = torch.zeros( + x.shape[0], + conv.out_channels, + *x_ft.shape[2:], + dtype=x_ft.dtype, + device=x.device, + ) + letters = "xyzt"[:dimension] + for k, corner in enumerate(CORNERS[dimension]): + index = [slice(None), slice(None)] + for high, m in zip(corner, modes[:-1]): + index.append(slice(-m, None) if high else slice(None, m)) + index.append(slice(None, modes[-1])) + index = tuple(index) + weights = torch.view_as_complex(getattr(conv, f"weights{k + 1}")) + block = torch.einsum( + f"bi{letters},io{letters}->bo{letters}", x_ft[index], weights + ) + out_ft[index] = out_ft[index] + block + return torch.fft.irfftn(out_ft, s=x.shape[2:], dim=dims) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +def test_spectral_conv_shape(dimension, device): + conv = _make_conv(dimension).to(device) + x = _make_input(dimension, 6, device) + y = conv(x) + assert y.shape == (2, OUT_CHANNELS, *([6] * dimension)) + assert y.dtype == torch.float32 + + +@pytest.mark.parametrize( + "sizes, modes", + [ + ((16,), (5,)), + ((15, 12), (4, 5)), + ((8, 9), (4, 3)), # 2 * modes == size, no gap between corners + ((10, 9), (6, 3)), # corners overlap along the first axis + ((9, 8, 10), (3, 2, 4)), + ((8, 8, 8), (5, 3, 3)), # corners overlap along the first axis + ((8, 6, 8), (3, 4, 3)), # corners overlap along the second axis + ((6, 7, 6, 8), (2, 3, 2, 3)), + ((6, 6, 6, 6), (2, 2, 4, 2)), # corners overlap along the third axis + ], +) +def test_spectral_conv_matches_reference(sizes, modes, device): + dimension = len(sizes) + conv = _make_conv(dimension, modes).to(device) + ref = _make_conv(dimension, modes).to(device) + ref.load_state_dict(conv.state_dict()) + x = torch.randn(2, IN_CHANNELS, *sizes, device=device, requires_grad=True) + x_ref = x.detach().clone().requires_grad_(True) + + y = conv(x) + y_ref = _reference_forward(ref, x_ref) + torch.testing.assert_close(y, y_ref) + + grad = torch.randn_like(y) + y.backward(grad) + y_ref.backward(grad) + torch.testing.assert_close(x.grad, x_ref.grad) + for p, p_ref in zip(conv.parameters(), ref.parameters()): + torch.testing.assert_close(p.grad, p_ref.grad) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +@pytest.mark.parametrize("size", [6, 8]) +@pytest.mark.parametrize("amp_dtype", AMP_DTYPES) +def test_spectral_conv_autocast_reduced_precision_input( + dimension, size, amp_dtype, device +): + # Under autocast, the layer before a spectral convolution produces a + # reduced-precision tensor. The spectral convolution must run on it in + # float32 and match the float32 result exactly. Size 6 is not a power of + # two, which cuFFT rejects in half precision; at size 8, a half-precision + # FFT runs but the mode multiplication does not, and a bfloat16 FFT may + # not be computed in float32. + _skip_unsupported_amp(device, amp_dtype) + device_type = torch.device(device).type + conv = _make_conv(dimension).to(device) + x = _make_input(dimension, size, device).to(amp_dtype).requires_grad_(True) + with torch.autocast(device_type, dtype=amp_dtype): + y = conv(x) + assert y.dtype == torch.float32 + torch.testing.assert_close(y, conv(x.detach().float()), rtol=0, atol=0) + + y.square().mean().backward() + assert torch.isfinite(x.grad.float()).all() + assert all(torch.isfinite(p.grad).all() for p in conv.parameters()) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +@pytest.mark.parametrize("amp_dtype", AMP_DTYPES) +def test_spectral_conv_autocast_float32_input(dimension, amp_dtype, device): + # Regression guard: a float32 input under autocast gives the same output + # as without autocast. + _skip_unsupported_amp(device, amp_dtype) + device_type = torch.device(device).type + conv = _make_conv(dimension).to(device) + x = _make_input(dimension, 8, device) + reference = conv(x) + with torch.autocast(device_type, dtype=amp_dtype): + y = conv(x) + torch.testing.assert_close(y, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize("dimension", [2, 3]) +@pytest.mark.parametrize("amp_dtype", AMP_DTYPES) +def test_fno_autocast_training_step(dimension, amp_dtype, device): + _skip_unsupported_amp(device, amp_dtype) + device_type = torch.device(device).type + model = FNO( + in_channels=2, + out_channels=1, + dimension=dimension, + latent_channels=8, + num_fno_layers=2, + num_fno_modes=3, + padding=2, + ).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + scaler = torch.amp.GradScaler(device_type, enabled=amp_dtype == torch.float16) + x = torch.randn(2, 2, *([10] * dimension), device=device) + with torch.autocast(device_type, dtype=amp_dtype): + loss = model(x).float().square().mean() + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + assert torch.isfinite(loss) + assert all(torch.isfinite(p).all() for p in model.parameters()) + + +@pytest.mark.parametrize("dimension", [2, 3]) +@pytest.mark.parametrize("use_autocast", [False, True]) +def test_spectral_conv_compile(dimension, use_autocast, device): + device_type = torch.device(device).type + # CPU autocast is tested with bfloat16 only + amp_dtype = torch.float16 if device_type == "cuda" else torch.bfloat16 + conv = _make_conv(dimension).to(device) + compiled = torch.compile(conv, fullgraph=True, backend="aot_eager") + x = _make_input(dimension, 8, device) + if use_autocast: + x = x.to(amp_dtype) + x = x.requires_grad_(True) + x_compiled = x.detach().clone().requires_grad_(True) + + with torch.autocast(device_type, dtype=amp_dtype, enabled=use_autocast): + y = conv(x) + y_compiled = compiled(x_compiled) + torch.testing.assert_close(y_compiled, y) + y.sum().backward() + y_compiled.sum().backward() + torch.testing.assert_close(x_compiled.grad, x.grad) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +def test_spectral_conv_meta_device(dimension): + conv = _make_conv(dimension).to("meta") + y = conv(_make_input(dimension, 6, "meta")) + assert y.shape == (2, OUT_CHANNELS, *([6] * dimension)) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +def test_spectral_conv_float64(dimension, device): + conv = _make_conv(dimension).to(device) + x = _make_input(dimension, 6, device) + reference = conv(x) + + conv = conv.double() + y = conv(x.double()) + assert y.dtype == torch.float64 + torch.testing.assert_close(y.float(), reference, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("dimension", [1, 2, 3, 4]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_spectral_conv_reset_parameters_keeps_device_and_dtype( + dimension, dtype, device +): + conv = _make_conv(dimension).to(device=device, dtype=dtype) + conv.reset_parameters() + for param in conv.parameters(): + assert param.device == torch.device(device) + assert param.dtype == dtype + assert param.requires_grad From 92f5082c8b36d617c4056bf0b866eb12d94155bf Mon Sep 17 00:00:00 2001 From: GeorgeLiang3 Date: Wed, 21 Oct 2026 10:07:45 -0600 Subject: [PATCH 2/2] test(fno): check gradients and weight updates in the autocast training 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 --- test/nn/module/test_spectral_conv.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/test/nn/module/test_spectral_conv.py b/test/nn/module/test_spectral_conv.py index 11a1b654e8..e1861d3779 100644 --- a/test/nn/module/test_spectral_conv.py +++ b/test/nn/module/test_spectral_conv.py @@ -201,13 +201,21 @@ def test_fno_autocast_training_step(dimension, amp_dtype, device): optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scaler = torch.amp.GradScaler(device_type, enabled=amp_dtype == torch.float16) x = torch.randn(2, 2, *([10] * dimension), device=device) + params = [p for p in model.parameters() if p.requires_grad] + initial = [p.detach().clone() for p in params] with torch.autocast(device_type, dtype=amp_dtype): loss = model(x).float().square().mean() + assert torch.isfinite(loss) scaler.scale(loss).backward() + + # GradScaler skips the optimizer step when gradients are not finite, so + # check the unscaled gradients and that the step updated every weight + scaler.unscale_(optimizer) + assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in params) scaler.step(optimizer) scaler.update() - assert torch.isfinite(loss) - assert all(torch.isfinite(p).all() for p in model.parameters()) + assert all(torch.isfinite(p).all() for p in params) + assert all(not torch.equal(p, p0) for p, p0 in zip(params, initial)) @pytest.mark.parametrize("dimension", [2, 3])