physicsnemo

PhysicsNeMo

Overview

Open-source Python framework for building, training, and fine-tuning physics-informed AI models on NVIDIA GPUs. PhysicsNeMo provides optimized model architectures (FNO, AFNO, GNNs, diffusion models), scalable distributed training, and physics-constrained learning — enabling real-time AI surrogates for scientific simulation across weather, CFD, molecular dynamics, and engineering domains.

Quick Pattern

Incorrect — manual PyTorch training without physics optimization:

model = MyModel().cuda()
for batch in dataloader:
    pred = model(batch)
    loss = F.mse_loss(pred, target)
    loss.backward()

Correct — PhysicsNeMo with optimized training and CUDA graphs:

import physicsnemo
from physicsnemo.datapipes.benchmarks.darcy import Darcy2D
from physicsnemo.metrics.general.mse import mse
from physicsnemo.models.fno.fno import FNO

model = FNO(
    in_channels=1, out_channels=1,
    dimension=2, latent_channels=32,
    num_fno_layers=4, num_fno_modes=12,
    padding=5,
).to("cuda")

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
dataloader = Darcy2D(resolution=256, batch_size=64)

for batch in dataloader:
    pred = model(batch["permeability"])
    loss = mse(pred, batch["darcy"])
    loss.backward()
    optimizer.step()

Quick Command

# Install PhysicsNeMo
pip install nvidia-physicsnemo

# Install with all optional dependencies
pip install nvidia-physicsnemo[all]

# Install PhysicsNeMo-Sym (physics-informed constraints)
pip install Cython
pip install nvidia-physicsnemo-sym --no-build-isolation

# Run with Docker container
docker pull nvcr.io/nvidia/physicsnemo/physicsnemo:latest
docker run --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 \
  --runtime nvidia -v ${PWD}:/workspace \
  -it --rm nvcr.io/nvidia/physicsnemo/physicsnemo:latest bash

# Distributed training
torchrun --nproc_per_node=4 train.py

# Clone training recipes
git clone https://github.com/NVIDIA/physicsnemo.git