multi-node (beyond multi-GPU) inference with UMA 1.2

[gauthierlab]

Python version

3.11.15

fairchem-core version

2.19.0

pytorch version

2.8.0

cuda version

12.8

Operating system version

whatever is on Perlmutter; SUSE Linux Enterprise Server 15 SP6

Minimal example

import os, sys, torch
import ray
from ase.io import read, write
from ase.optimize import BFGS
from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnit
from fairchem.core.units.mlip_unit import guess_inference_settings
from fairchem.core import FAIRChemCalculator


MODEL_PATH = "/global/cfs/cdirs/m5215/jgauth32/mlip_models/uma-s-1p2.pt"

inference_settings = guess_inference_settings("turbo")
predictor = ParallelMLIPPredictUnit(
    inference_model_path=MODEL_PATH,
    device="cuda",
    inference_settings=inference_settings,
    num_workers=8,
    num_workers_per_node=4,
)

calc = FAIRChemCalculator(predictor, task_name="omat")

if os.path.exists("qn.traj"):
    atoms = read("qn.traj", index=-1)
else:
    atoms = read("init.traj")

atoms.calc = calc

optimizer = BFGS(atoms, trajectory="qn.traj", logfile="qn.log")
optimizer.run(fmax=0.05)

write("final_relaxed.traj", atoms)

Current behavior

Multi-GPU inference using UMA works using the above script but only when it is fully within a single node. On Perlmutter, this means up to 4 GPUs. If your system needs more than 4 GPUs (the script above tries to use 8, or two nodes) the system hangs when trying to set up the Ray cluster.

Expected Behavior

I would've anticipated this to work the same across multiple nodes, but there seems to be something that goes wrong when orchestrating the Ray cluster. Claude code was trying to do something complicated that involved a check-delay loop to make sure the Ray cluster on each node was properly communicating but I can't imagine that's the intended functionality here.

Relevant files to reproduce this bug

No response

[rayg1234]

hi @gauthierlab , we don't setup a ray cluster automatically when you use >1 node. This is because the way users configure Ray is extremely different based on their network and cluster setup and ParallelMLIPPredictUnit will automatically connect to an existing cluster (in the case of a single node job, it will try to create a local cluster if it doesn't already exist).

So if you want to use multi-node, we expect you to have a RayCluster already setup. We provide some helper utils to setup a Raycluster on top of Slurm (this is what we use internally as well) here https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/launchers/cluster/ray_cluster.py#L370

here some sample code to get you started

    from fairchem.core.launchers.cluster.ray_cluster import RayCluster

    cluster = RayCluster(log_dir=Path(output_dir))
    cluster_reqs = {
        "slurm_account": "ocp",
        "slurm_qos": qos,
        "timeout_min": 10080,
        "mem_gb": 0,
        "nodes": nodes,
        "gpus_per_task": 8,
        "cpus_per_task": 192,
        "tasks_per_node": 1,
    }
    cluster.start_head_and_workers(
        name=job_name,
        requirements=cluster_reqs,
        payload=main_func, # this is the payload function you need to put all your actual code in
        output_path=output_dir,
        workers=workers,
    )

[gauthierlab]

Thanks, Ray. Misko helped me out with that last week and I settled on the following script for running multi-node inference on Perlmutter:

import logging
import os
from pathlib import Path

from fairchem.core.launchers.cluster.ray_cluster import RayCluster
from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnit
from fairchem.core import FAIRChemCalculator
from ase.io import read, write
from ase.optimize import BFGS

from ase.build import bulk

FMAX = 0.05       # eV/Ang convergence criterion

def my_predict_job(output_path: str, workers: int):
    """This runs on the head node AFTER the Ray cluster is up."""
    settings = InferenceSettings(
        tf32=True,
        merge_mole=True,
        compile=False,
        internal_graph_gen_version=2,
        external_graph_gen=False,
    )

    predictor = ParallelMLIPPredictUnit(
        inference_model_path="/global/cfs/cdirs/m5215/jgauth32/mlip_models/uma-s-1p2.pt",
        device="cuda",
        inference_settings=settings,
        num_workers=workers,
        num_workers_per_node=4,
    )
    calc = FAIRChemCalculator(predictor, task_name="omat")

    if os.path.exists("qn.traj"):
        print("Resuming from existing qn.traj ...", flush=True)
        atoms = read("qn.traj", index=-1)
    else:
        print("Starting from init.traj ...", flush=True)
        atoms = read("init.traj")

    atoms.calc = calc

    optimizer = BFGS(atoms, trajectory="qn.traj", logfile="qn.log")
    optimizer.run(fmax=FMAX)

    write("final_relaxed.traj", atoms)

def main():
    workers = 32          # 32 GPUs total
    gpus_per_node = 4     # 4 GPUs per node
    nodes = math.ceil(workers / gpus_per_node)  # = 2 nodes
    output_dir = "/global/cfs/cdirs/m5215/jgauth32/mahsa_uma_test/relaxation/nodes08"

    cluster = RayCluster(log_dir=Path(output_dir))
    cluster.start_head_and_workers(
        name="predict-test",
        requirements={
            "slurm_account": "m5215",
            "slurm_constraint": "gpu",
            "slurm_qos": "debug",
            "slurm_additional_parameters": {
                "output": f"{output_dir}/opt.log",
                "error": f"{output_dir}/err.log",
            },
            "timeout_min": 30,
            "mem_gb": 0,
            "nodes": nodes,
            "gpus_per_node": gpus_per_node,
            "cpus_per_task": 32,
            "tasks_per_node": 1,
        },
        payload=my_predict_job,
        output_path=output_dir,
        workers=workers,
    )
    print("SLURM job submitted. Check logs in:", output_dir)


if __name__ == "__main__":
    main()```

The job launched and seemed to start as intended, though it was hanging after the very first step of the optimization. Have you all encountered this before?

[matr1x-1]

I’m seeing what looks like the same multi-node inference failure path with FairChem + LAMMPS on CSCS Alps/GH200.

Setup:

• 6 nodes, 4 GPUs per node, 24 GPUs total
• ParallelMLIPPredictUnit through the FairChem LAMMPS interface
• Ray cluster reaches all 6 nodes / 24 GPUs
• UMA model loads successfully on all workers
• Failure happens before the first LAMMPS timestep, inside FairChem distributed inference

Failure progression:

1. Initially Ray placed a remote worker on the same GPU as the local FairChem rank 0, causing NCCL duplicate GPU errors.
2. After reserving a GPU for local rank 0, rank 0 and Ray workers used different network backends until I forced the same Libfabric path.
3. With all ranks using Libfabric, the 24-rank run hangs/times out in an NCCL ALLGATHER.
4. If I cap graph parallelism to 8 inside the 24-worker run, the job gets past setup but NCCL/OFI reports:
   NET/OFI Request ... completed with error. RC: 5. Error: 16 (Device or resource busy)

A smaller 8-GPU / 2-node run completed a short LAMMPS test, but the 24-GPU / 6-node configuration consistently fails before MD starts.

This makes it look like the Ray cluster can come up, but FairChem’s multi-node distributed inference is not stable once the worker count crosses some size/
topology threshold. Happy to provide Slurm/NCCL logs if useful.