NVIDIA PhysicsNeMo: Deep Learning Framework for Physics-ML Models

Introduction

NVIDIA PhysicsNeMo is an open-source deep learning framework from NVIDIA, specifically engineered for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods. This powerful framework enables researchers and engineers to develop AI models that combine physics knowledge with data, facilitating real-time predictions across various scientific and engineering domains. Built on PyTorch, PhysicsNeMo offers a scalable, GPU-optimized stack for exploring neural operators, Graph Neural Networks (GNNs), transformers, Physics-Informed Neural Networks (PINNs), and hybrid approaches. The project is currently undergoing an update to v2.0, promising easier installation and integration.

PhysicsNeMo provides modular Python components, including:

physicsnemo.models: A collection of optimized model architectures like Neural Operators, GNNs, and Diffusion models.
physicsnemo.datapipes: Scalable data pipelines for scientific data structures.
physicsnemo.distributed: Utilities for parallel training on NVIDIA GPUs.
physicsnemo.sym: Symbolic PDE residual computation for physics-informed losses.

Installation

PhysicsNeMo offers flexible installation options to suit different environments, including pip, uv, NVIDIA container images, or building from source.

Via pip

To install the latest version from PyPI:

pip install nvidia-physicsnemo

For GPU-accelerated packages and a CUDA-matched PyTorch build, specify the CUDA backend (e.g., cu13 for CUDA 13 or cu12 for CUDA 12) along with optional feature extras:

# CUDA 13 backend with nn-extras
pip install "nvidia-physicsnemo[cu13,nn-extras]"

# CUDA 12 backend with nn-extras
pip install "nvidia-physicsnemo[cu12,nn-extras]"

Using NVIDIA Container

The PhysicsNeMo Docker image is available from the NVIDIA Container Registry. Pull the latest tag:

docker pull nvcr.io/nvidia/physicsnemo/physicsnemo:25.06

You can then run the container and clone the repository to access examples:

docker run --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --runtime nvidia \
--rm -it nvcr.io/nvidia/physicsnemo/physicsnemo:25.06 bash
git clone https://github.com/NVIDIA/physicsnemo.git
cd physicsnemo/examples/cfd/darcy_fno/
pip install warp-lang # install NVIDIA Warp to run the Darcy example
python train_fno_darcy.py

Examples

Getting started with PhysicsNeMo is straightforward. Here's a "Hello World" example demonstrating a simple fully connected model:

import torch
from physicsnemo.models.mlp.fully_connected import FullyConnected

model = FullyConnected(in_features=32, out_features=64)
input = torch.randn(128, 32)
output = model(input)
print(output.shape)
# Expected output: torch.Size([128, 64])

For distributed training, PhysicsNeMo integrates seamlessly with torch.distributed:

import torch
from torch.nn.parallel import DistributedDataParallel
from physicsnemo.distributed import DistributedManager
from physicsnemo.models.mlp.fully_connected import FullyConnected

def main():
    DistributedManager.initialize()
    dist = DistributedManager()

    arch = FullyConnected(in_features=32, out_features=64).to(dist.device)

    if dist.distributed:
        ddps = torch.cuda.Stream()
        with torch.cuda.stream(ddps):
            arch = DistributedDataParallel(
                arch,
                device_ids=[dist.local_rank],
                output_device=dist.device,
                broadcast_buffers=dist.broadcast_buffers,
                find_unused_parameters=dist.find_unused_parameters,
            )
        torch.cuda.current_stream().wait_stream(ddps)

    # Set up the optimizer
    optimizer = torch.optim.Adam(
        arch.parameters(),
        lr=0.001,
    )

    def training_step(invar, target):
        pred = arch(invar)
        loss = torch.sum(torch.pow(pred - target, 2))
        loss.backward()
        optimizer.step()
        return loss

    # Sample training loop
    for i in range(20):
        # Random inputs and targets for simplicity
        input = torch.randn(128, 32, device=dist.device)
        target = torch.randn(128, 64, device=dist.device)

        # Training step
        loss = training_step(input, target)

if __name__ == "__main__":
    main()

PhysicsNeMo also includes a symbolic PDE module for defining equations:

from physicsnemo.sym.eq.pdes.navier_stokes import NavierStokes
ns = NavierStokes(nu=0.01, rho=1, dim=2)
ns.pprint()
# Expected output:
# continuity: u__x + v__y
# momentum_x: u*u__x + v*u__y + p__x + u__t - 0.01*u__x__x - 0.01*u__y__y
# momentum_y: u*v__x + v*v__y + p__y + v__t - 0.01*v__x__x - 0.01*v__y__y

For a comprehensive collection of examples and reference samples, refer to the PhysicsNeMo examples directory.

Why Use PhysicsNeMo?

PhysicsNeMo is designed to accelerate scientific machine learning workflows with several key advantages:

Scalable GPU-Optimized Training Library: Maximizes the power of NVIDIA GPUs with distributed computing utilities for efficient scaling from single to multi-node GPU clusters. It includes advanced optimization utilities, tailor-made datapipes, and symbolic PDE utilities for enhanced training speed.
A Suite of Physics-Informed ML Models: Offers a rich library of state-of-the-art models optimized for Physics-ML applications, such as Neural Operators (FNOs, DeepONet), Graph Neural Networks (MeshGraphNet), Diffusion Models, and Physics-Informed Neural Networks (PINNs). These models significantly reduce development time for high-fidelity simulations.
Seamless PyTorch Integration: Built on top of PyTorch, it provides a familiar and user-friendly experience, allowing users to leverage the extensive PyTorch ecosystem while benefiting from PhysicsNeMo's specialized capabilities.
Easy Customization and Extension: Highly extensible with Pythonic APIs for defining new physics models, geometries, and constraints. Features like ONNX support, robust logging, and efficient checkpointing further enhance its adaptability.
AI4Science Library: Serves as a complementary tool to PyTorch for SciML and AI4Science applications, providing a deep learning research platform with optimal performance on NVIDIA GPUs. It also includes domain-specific packages like PhysicsNeMo CFD, PhysicsNeMo Curator, and Earth-2 Studio.
Key Benefits: PhysicsNeMo enables SciML benchmarking and validation, offers ease of using generalized SciML recipes with heterogeneous datasets, and provides out-of-the-box performance and scalability across multi-GPU and multi-node GPU setups.