Multi-GPU training

metatrain supports training a model with several GPUs, which can accelerate the training, especially when the training dataset is large / there are many training epochs. This feature is enabled by the torch.distributed module, and thus can do multiprocess parallelism across several nodes.

In multi-GPU training, every batch of samples is split into smaller mini-batches and the computation is run for each of the smaller mini-batches in parallel on different GPUs. The different gradients obtained on each device are then summed. This approach allows the user to reduce the time it takes to train models.

To know if the model supports multi-GPU training, please check Available Architectures and see if the default hyperparameters have the distributed option.

Input file

You do not need to change anything in the .yaml file for the training: distributed training is enabled automatically whenever the job runs under more than one SLURM task. For example, the options.yaml file of the Training a model from scratch example can be used as it is:

seed: 42
device: cuda

architecture:
  name: pet
  training:
    batch_size: 50
    num_epochs: 100

training_set:
  systems:
    read_from: ethanol_reduced_100.xyz
    length_unit: angstrom
  targets:
    energy:
      key: energy
      unit: eV
      forces: true

test_set: 0.0
validation_set: 0.5

Slurm script

Below is an example Slurm script for submitting the job. Please be aware that the actual configurations vary from clusters to clusters, so you have to modify it. Different scheduler will require similar options. metatrain will automatically use all the GPUs that you have asked for. You should make a single GPU visible for each process (setting --gpus-per-node equal to the number of GPUs, or setting --gpus-per-task=1, depending on your cluster configuration).

#!/bin/bash
#SBATCH --nodes 1
#SBATCH --ntasks 2  # must equal to the number of GPUs
#SBATCH --ntasks-per-node 2
#SBATCH --gpus-per-node 2  # use 2 GPUs
#SBATCH --cpus-per-task 8
#SBATCH --exclusive
#SBATCH --partition=h100  # adapt this to your cluster
#SBATCH --time=1:00:00

# load modules and/or virtual environments and/or containers here

srun mtt train options.yaml

Performance

If the multi-GPU training runs successfully, you should see this in the training log:

[2025-10-08 11:34:22][INFO] - Distributed environment set up with MASTER_ADDR=kh080,
MASTER_PORT=39591, WORLD_SIZE=2, RANK=0, LOCAL_RANK=0
[2025-10-08 11:34:23][INFO] - Training on 2 devices with dtype torch.float32

This 100-epoch training takes 23 seconds.

[2025-10-08 11:34:22][INFO] - Starting training from scratch
...
[2025-10-08 11:34:45][INFO] - Training finished!

Now let’s switch off the multi-GPU training by submitting the same job with a single task (--ntasks 1). The training takes 69 seconds.

[2025-10-08 11:37:38][INFO] - Setting up model
...
[2025-10-08 11:38:47][INFO] - Training finished!

Multi-GPU fine-tuning

You can use multi-GPU for fine-tuning too, in the same way: launch the job with one SLURM task per GPU. For information about fine-tuning, please refer to the Fine-tuning a pre-trained model example.