Skip to navigation
Training

Multi-GPU Training

View as Markdown

Single Node, Multiple GPUs (most common setup)

models:
  - kind: GNN_XGBoost
    gpu: multi
    num_gpus: 4
docker run --rm -it --gpus all \
  --shm-size=10g --ulimit memlock=-1 --ulimit stack=67108864 \
  -v /your/data:/data:ro \
  -v /your/output:/workspace/output \
  -v /path/to/my_config.yaml:/workspace/config.yaml:ro \
  nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0

The container reads num_gpus from the config and passes --nproc_per_node=4 to torchrun. Each GPU runs one worker process. PyTorch DDP synchronises gradients across all workers after every mini-batch.

Single GPU

models:
  - kind: GNN_XGBoost
    gpu: single
    # num_gpus is omitted (defaults to 1)

Or with CLI:

torchrun --standalone --nproc_per_node=1 /workspace/mnmg/train.py --config my_config.yaml

How Distributed Training Works

GNN phase:

  • WholeGraph partitions the graph across all GPUs. Each GPU owns a shard of node feature tensors.
  • During neighbourhood sampling, if a sampled neighbour is on a different GPU, WholeGraph fetches its features transparently over NVLink or PCIe.
  • The GNN model is wrapped with PyTorch DistributedDataParallel. Gradients are all-reduced using NCCL after each mini-batch.
  • Only rank 0 saves the final model weights.

XGBoost phase:

  • XGBoost uses its built-in RabitTracker for tree building. Each process holds a shard of the embedding data.
  • All GPUs collaborate on each tree (splits are computed across all shards simultaneously).
  • Rank 0 saves the final booster to output_dir/xgboost_fraud.json.

Multi-Node Multi-GPU (MNMG) — Slurm

For clusters managed by Slurm, the training container can be launched with srun across multiple nodes. Save the script below as e.g. slurm_train_ep.sh and submit with sbatch slurm_train_ep.sh.

#!/bin/bash
# ══════════════════════════════════════════════════════════════════════════════
# USER-EDITABLE SECTION — adjust these before submitting
# ══════════════════════════════════════════════════════════════════════════════
#SBATCH --job-name=fraud-mnmg
#SBATCH --nodes=4
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=16
#SBATCH --mem=128G
#SBATCH --time=00:15:00
#SBATCH --partition=batch
#SBATCH --account=<your-slurm-account>
#SBATCH --output=logs/fraud-mnmg-%j.out
#SBATCH --error=logs/fraud-mnmg-%j.err
# ══════════════════════════════════════════════════════════════════════════════
# DO NOT EDIT BELOW THIS LINE
# ══════════════════════════════════════════════════════════════════════════════
#SBATCH --ntasks-per-node=1
set -euo pipefail
mkdir -p logs
CONTAINER_IMAGE="nvcr.io/nvidia/cugraph/financial-fraud-training:3.0.0"
DATA_DIR="/path/to/your/graph/data"
CONFIG_FILE="$(pwd)/config.yaml"
OUTPUT_DIR="$(pwd)/output_dir"
PARTITION_DIR="$(pwd)/partition_dir"
mkdir -p "$OUTPUT_DIR"
mkdir -p "$PARTITION_DIR"
echo "Job        : $SLURM_JOB_ID"
echo "Nodes      : $SLURM_JOB_NUM_NODES ($SLURM_JOB_NODELIST)"
echo "Container  : $CONTAINER_IMAGE"
echo "Data       : $DATA_DIR"
echo "Config     : $CONFIG_FILE"
echo "Output     : $OUTPUT_DIR"
# The entrypoint.sh inside the container reads SLURM_* env vars to derive
# NNODES, NODE_RANK, MASTER_ADDR, and nproc_per_node automatically.
# No manual torchrun flags are needed.
srun --ntasks-per-node=1 \
    --container-image="$CONTAINER_IMAGE" \
    --container-mounts="$DATA_DIR:/data:ro,$CONFIG_FILE:/workspace/config.yaml:ro,$OUTPUT_DIR:/workspace/output,$PARTITION_DIR:/tmp/fraud_partitions" \
    --container-env=SLURM_JOB_ID,SLURM_NNODES,SLURM_JOB_NUM_NODES,SLURM_NODEID,SLURM_GPUS_ON_NODE,SLURM_GPUS_PER_NODE,SLURM_LAUNCH_NODE_IPADDR,SLURM_SRUN_COMM_HOST \
    /workspace/entrypoint.sh
echo "Training complete."

Key parameters to configure:

ParameterDescription
--nodesNumber of cluster nodes. Each node runs one srun task.
--gpus-per-nodeGPUs per node. Controls nproc_per_node through SLURM_GPUS_ON_NODE — num_gpus in config.yaml is ignored under Slurm.
--accountYour Slurm allocation account.
--timeWall-clock limit.
DATA_DIRPath on the shared filesystem to your graph data directory.
CONFIG_FILEPath to your config YAML (must be accessible from the head node).
OUTPUT_DIRPath on the shared filesystem where artifacts are written.

How NPROC (workers per node) is resolved — first match wins:

Launch modeWhat sets NPROC
Docker single-nodeconfig.yaml → num_gpus
Slurm (--gpus-per-node=N)SLURM_GPUS_ON_NODE (set automatically by Slurm)
Any mode — explicit override-e NPROC=N (highest priority)
num_gpus in config.yaml is ignored under Slurm. The entrypoint reads GPU count from SLURM_GPUS_ON_NODE (set automatically from --gpus-per-node). For single-node Docker runs, num_gpus is used instead.

Submit training job:

sbatch slurm_train_ep.sh
# Check status
squeue -j <job_id>