feat: Switch from LayerNorm to RMSNorm

recipes/aero_cfd/scripts/train_ahmedml.py

@@ -1,6 +1,7 @@
 #  Copyright © 2026 Emmi AI GmbH. All rights reserved.
 
 from aero_cfd.presets import AhmedMLPreset
+
 from noether.core.distributed.utils import accelerator_to_device
 from noether.training.runners import HydraRunner
 

recipes/aero_cfd/scripts/train_drivaerml.py

@@ -1,6 +1,7 @@
 #  Copyright © 2026 Emmi AI GmbH. All rights reserved.
 
 from aero_cfd.presets import DrivAerMLPreset
+
 from noether.core.distributed.utils import accelerator_to_device
 from noether.training.runners import HydraRunner
 

recipes/aero_cfd/scripts/train_drivaernet.py

@@ -1,6 +1,7 @@
 #  Copyright © 2026 Emmi AI GmbH. All rights reserved.
 
 from aero_cfd.presets import DrivAerNetPreset
+
 from noether.core.distributed.utils import accelerator_to_device
 from noether.training.runners import HydraRunner
 

recipes/aero_cfd/scripts/train_emmi_wing.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from aero_cfd.presets import EmmiWingPreset
+
 from noether.core.distributed.utils import accelerator_to_device
 from noether.data.datasets.cfd.emmi_wing.dataset_hf import EmmiWingHFDataset
 from noether.training.runners import HydraRunner

recipes/aero_cfd/scripts/train_shapenet_car.py

@@ -1,6 +1,7 @@
 #  Copyright © 2026 Emmi AI GmbH. All rights reserved.
 
 from aero_cfd.presets import ShapeNetCarPreset
+
 from noether.core.distributed.utils import accelerator_to_device
 from noether.training.runners import HydraRunner
 

recipes/aero_cfd/showcase/utils/forces.py

@@ -8,7 +8,6 @@
 from pathlib import Path
 
 import torch
-
 from aero_cfd.utils.drag_lift import FlowConditions, compute_force_coefficients
 
 

src/noether/modeling/models/ab_upt.py

@@ -63,7 +63,7 @@ class ReadoutLayer(nn.Module):
 
     def __init__(self, decoder_config: LinearProjectionConfig, hidden_dim: int, condition_dim: int | None = None):
         super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_dim, elementwise_affine=condition_dim is None, eps=1e-6)
+        self.norm_final = nn.RMSNorm(hidden_dim, elementwise_affine=condition_dim is None, eps=1e-6)
         self.linear = LinearProjection(config=decoder_config)
         self.modulation = None
         if condition_dim is not None:

src/noether/modeling/models/aerodynamics.py

@@ -150,7 +150,7 @@ def __init__(self, model_config: AeroTransformerConfig, **kwargs):
 
         self.backbone = Transformer(config=model_config)
 
-        self.norm = nn.LayerNorm(hidden_dim, eps=1e-6)
+        self.norm = nn.RMSNorm(hidden_dim, eps=1e-6)
         self.out = LinearProjection(
             config=LinearProjectionConfig(
                 input_dim=hidden_dim,
@@ -258,7 +258,7 @@ def __init__(self, model_config: AeroTransolverConfig, **kwargs):
 
         self.backbone = Transformer(config=model_config)
 
-        self.norm = nn.LayerNorm(hidden_dim, eps=1e-6)
+        self.norm = nn.RMSNorm(hidden_dim, eps=1e-6)
         self.out = LinearProjection(
             config=LinearProjectionConfig(
                 input_dim=hidden_dim,

src/noether/modeling/models/upt.py

@@ -39,7 +39,7 @@ def __init__(
 
         self.decoder = DeepPerceiverDecoder(config=config.decoder_config)  # type: ignore[arg-type]
 
-        self.norm = nn.LayerNorm(
+        self.norm = nn.RMSNorm(
             config.decoder_config.perceiver_block_config.hidden_dim,
             eps=config.decoder_config.perceiver_block_config.eps,
         )

src/noether/modeling/modules/blocks/perceiver.py

@@ -43,11 +43,9 @@ def __init__(
                     "If modulation is enabled, modulation_linear_projection_config must be provided. Likely condition_dim is not set."
                 )
 
-        self.norm1q = torch.nn.LayerNorm(
-            config.hidden_dim, elementwise_affine=elementwise_affine, bias=config.bias, eps=config.eps
-        )
-        self.norm1kv = torch.nn.LayerNorm(
-            config.kv_dim or config.hidden_dim, elementwise_affine=elementwise_affine, bias=config.bias, eps=config.eps
+        self.norm1q = nn.RMSNorm(config.hidden_dim, eps=config.eps, elementwise_affine=elementwise_affine)
+        self.norm1kv = nn.RMSNorm(
+            config.kv_dim or config.hidden_dim, eps=config.eps, elementwise_affine=elementwise_affine
         )
 
         self.attn = PerceiverAttention(config=config.perceiver_attention_config)  # type: ignore[arg-type]
@@ -56,9 +54,7 @@ def __init__(
 
         self.drop_path1 = UnquantizedDropPath(config=config.drop_path_config)  # type: ignore[arg-type]
 
-        self.norm2 = torch.nn.LayerNorm(
-            config.hidden_dim, elementwise_affine=elementwise_affine, bias=config.bias, eps=config.eps
-        )
+        self.norm2 = nn.RMSNorm(config.hidden_dim, eps=config.eps, elementwise_affine=elementwise_affine)
 
         self.mlp = UpActDownMlp(config=config.up_act_down_mlp_config)  # type: ignore[arg-type]
         self.ls2 = LayerScale(config=config.layerscale_config)  # type: ignore[arg-type]

src/noether/modeling/modules/blocks/transformer.py

@@ -40,9 +40,7 @@ def __init__(
             self.modulation = LinearProjection(config=config.modulation_linear_projection_config)  # type: ignore[arg-type]
             elementwise_affine = False
 
-        self.norm1 = torch.nn.LayerNorm(
-            config.hidden_dim, elementwise_affine=elementwise_affine, bias=config.bias, eps=config.eps
-        )
+        self.norm1 = nn.RMSNorm(config.hidden_dim, eps=config.eps, elementwise_affine=elementwise_affine)
 
         try:
             if callable(config.attention_constructor):
@@ -65,10 +63,7 @@ def __init__(
         self.drop_path1 = UnquantizedDropPath(
             config=config.drop_path_config  # type: ignore[arg-type]
         )
-
-        self.norm2 = torch.nn.LayerNorm(
-            config.hidden_dim, elementwise_affine=elementwise_affine, bias=config.bias, eps=config.eps
-        )
+        self.norm2 = nn.RMSNorm(config.hidden_dim, eps=config.eps, elementwise_affine=elementwise_affine)
         self.mlp = UpActDownMlp(config=config.up_act_down_mlp_config)  # type: ignore[arg-type]
         self.ls2 = LayerScale(config=config.layerscale_config)  # type: ignore[arg-type]
         self.drop_path2 = UnquantizedDropPath(config=config.drop_path_config)  # type: ignore[arg-type]