Multi-GPU Training#

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:

Parameter	Description
`--nodes`	Number of cluster nodes. Each node runs one `srun` task.
`--gpus-per-node`	GPUs per node. Controls `nproc_per_node` through `SLURM_GPUS_ON_NODE` — `num_gpus` in `config.yaml` is ignored under Slurm.
`--account`	Your Slurm allocation account.
`--time`	Wall-clock limit.
`DATA_DIR`	Path on the shared filesystem to your graph data directory.
`CONFIG_FILE`	Path to your config YAML (must be accessible from the head node).
`OUTPUT_DIR`	Path on the shared filesystem where artifacts are written.

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

Launch mode	What sets `NPROC`
Docker single-node	`config.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)

Note: 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>