Inference with multiple GPUs.

[zhangqi94]

In the eval_configs.py file, the data_loader is encountering an issue where the length of edge_index (and other related quantities) varies for each structure in the batch. This causes an issue when attempting to use torch.nn.DataParallel for inference. Any suggestions on how to resolve this?

[ilyes319]

Hey @zhangqi94 ,

Can you post a model detail problem you have?

[zhangqi94]

Thank you! The program runs normally on a single GPU, but an error occurs when using multiple GPUs. I think the issue might be related to the dataloader? Here is the log output during the error:
log.txt

The code was modified based on eval_configs.py.

import torch
import torch.multiprocessing as mp
torch.set_default_dtype(torch.float64)

from mace import data
from mace.tools import torch_geometric, torch_tools, utils
from mace.cli.convert_e3nn_cueq import run as run_e3nn_to_cueq
from copy import deepcopy

import ase
import argparse
import numpy as np


import os
import torch
import torch.distributed as dist

print("="*100)
world_size = torch.cuda.device_count()
print("torch device count:", world_size)
print("="*100)


####################################################################################################

mace_modelpath = "macemodel/mace_iceX_l1x64r4.0.model"
mace_batchsize = 4
batch = 32

atoms_list = []
for i in range(batch):
    i = i + 1
    stru_path = f"data_vasp/ice10_n16/{i:06d}/POSCAR"
    atoms = ase.io.read(stru_path)
    atoms_list.append(atoms.copy())
    
####################################################################################################
# Parse model arguments
mace_args = argparse.Namespace(default_type = "float64",
                                model = mace_modelpath,
                                device = f"cuda",
                                compute_stress = True,
                                )

# Set the default data type for PyTorch and initialize the computation device
torch_tools.set_default_dtype(mace_args.default_type)
device = torch_tools.init_device(mace_args.device)


# Load the pretrained model and move it to the specified device
model = torch.load(f=mace_args.model, map_location=mace_args.device)
model = run_e3nn_to_cueq(deepcopy(model), device=device)
# model = model.to(mace_args.device)
for param in model.parameters():
    param.requires_grad = False

#######################################################
atomic_numbers = model.atomic_numbers
r_max = model.r_max
try:
    heads = model.heads
except AttributeError:
    heads = None
    
#######################################################
print("="*100)
print("Start inference on various devices")
print("="*100)

if world_size > 1:
    model = torch.nn.DataParallel(model)
model = model.to(mace_args.device)

#######################################################
# Convert ASE Atoms objects to model-compatible configurations
configs = [data.config_from_atoms(atoms) for atoms in atoms_list]
z_table = utils.AtomicNumberTable([int(z) for z in atomic_numbers])


data_set = [data.AtomicData.from_config(
            config, z_table=z_table, cutoff=float(r_max), heads=heads
            )
            for config in configs
            ]
# Create a data loader for batched processing of input data
data_loader = torch_geometric.dataloader.DataLoader(
    dataset = data_set,
    batch_size = mace_batchsize,
    shuffle = False,
    drop_last = False,
)


print("="*100)
print(data_loader)
for batch in data_loader:
    print(batch)
print("="*100)

############################

# Initialize lists to store predictions
energies_list = []
stresses_list = []
forces_collection = []

# Process each batch
for batch in data_loader:
    batch = batch.to(device)
    batch_dict = batch.to_dict()
    output = model(batch_dict, compute_stress=mace_args.compute_stress,)

    energies_list.append(torch_tools.to_numpy(output["energy"]))
    stresses_list.append(torch_tools.to_numpy(output["stress"]))
    forces = np.split(
        torch_tools.to_numpy(output["forces"]),
        indices_or_sections=batch.ptr[1:],
        axis=0,
    )
    forces_collection.append(forces[:-1])  # drop last as its empty

# Concatenate and convert predictions to numpy arrays
energies = np.concatenate(energies_list, axis=0).astype(np.float64)
stresses = np.concatenate(stresses_list, axis=0).astype(np.float64)

# Flatten forces and ensure consistency with input
forces_list = [forces for forces_list in forces_collection for forces in forces_list]
forces = np.array(forces_list, dtype=np.float64)


print("energies:", energies.shape, energies)
print("forces:", forces.shape)
print("stresses:", stresses.shape)

[RokasEl]

I don't think nn.DataParallel works correctly for graph data. In fact, looking at the source code shows that the data will be split apart roughly evenly across the first dimension. That is, some configurations will probably split and sent to different GPUs hence your error.

A better way is to just split your data in two and run the individual halves on the two GPUs you have.