Skip to content

GeoT optimization 3/4: Add fused batched Muon optimizer #1743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
coreyjadams merged 6 commits into from
Jun 29, 2026
+446 −4
Merged

GeoT optimization 3/4: Add fused batched Muon optimizer #1743

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
47521d9
Add fused batched Muon optimizer
coreyjadams Jun 22, 2026
0d20603
Apply pre-commit
coreyjadams Jun 22, 2026
413cf35
Merge branch 'main' into geoT-opt-muon-opt-fusion
coreyjadams Jun 29, 2026
4ea5813
Update muon: inherit, not reimplement. Catch sharded weights and
coreyjadams Jun 29, 2026
34e0207
Merge branch 'main' into geoT-opt-muon-opt-fusion
coreyjadams Jun 29, 2026
fdc0b5a
Clean up muon PR, purge DTensor checks
coreyjadams Jun 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from torch.amp import autocast

from physicsnemo.mesh import DomainMesh, Mesh
from physicsnemo.optim import CombinedOptimizer
from physicsnemo.optim import CombinedOptimizer, Muon

### Recipe-wide type aliases. Re-exported for use in loss.py, metrics.py,
### output_normalize.py, forward_kwargs.py, collate.py, train.py, infer.py,
Expand Down Expand Up @@ -112,7 +112,7 @@ def build_muon_optimizer(
if muon_params and other_params:
return CombinedOptimizer(
[
torch.optim.Muon(
Muon(
muon_params,
lr=lr,
weight_decay=weight_decay,
Expand All @@ -129,7 +129,7 @@ def build_muon_optimizer(
torch_compile_kwargs=compile_kwargs,
)
elif muon_params:
opt = torch.optim.Muon(
opt = Muon(
muon_params,
lr=lr,
weight_decay=weight_decay,
Expand Down
3 changes: 2 additions & 1 deletion physicsnemo/optim/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,5 +17,6 @@
"""Optimizer utilities for PhysicsNeMo."""

from physicsnemo.optim.combined_optimizer import CombinedOptimizer
from physicsnemo.optim.muon import Muon

__all__ = ["CombinedOptimizer"]
__all__ = ["CombinedOptimizer", "Muon"]
Loading
Loading