Running VulnTrain on an HPC cluster#

This page provides a reference configuration for running VulnTrain on a EuroHPC-style GPU cluster managed by SLURM. Treat it as a template and adapt the account, partition, and resource requests to match your local system.

High-level workflow:

Create a shared Conda environment.
Test the installation on the login node.
Submit a SLURM job that launches distributed training.

1. Shared Conda environment#

Run the following once on a login node to create a reusable environment for VulnTrain. If your site provides Conda/Miniconda via environment modules, load the relevant module first.

# 1. Make a folder for shared environments if it doesn't exist
mkdir -p $HOME/conda_envs

# 2. Create a Conda environment with Python 3.11
conda create -y -p $HOME/conda_envs/vulntrain python=3.11

# 3. Activate the environment
conda activate $HOME/conda_envs/vulntrain

# 4. Upgrade pip
pip install --upgrade pip

# 5. Install vulntrain and dependencies
pip install vulntrain datasets transformers accelerate

This creates an environment at $HOME/conda_envs/vulntrain that can be activated from both login and compute nodes.

2. Test the environment on the login node#

Before requesting GPUs, quickly verify that Python and the VulnTrain CLI are available in the environment:

# Activate environment
conda activate $HOME/conda_envs/vulntrain

# Check Python version
python --version

# Check vulntrain CLI
vulntrain-train-severity-classification --help

If these commands run without errors, you are ready to submit a job.

3. Example SLURM scripts#

Save one of the scripts below as run_vulntrain.slurm (or a similar name) in your working directory. Adjust --account, --partition, time, memory, and GPU counts to match your project and cluster policies.

Single-node configuration using `torchrun`#

#!/bin/bash
#SBATCH --job-name=vulntrain
#SBATCH --account=<-your-account-id->
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --ntasks=4
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=10:00:00
#SBATCH --output=logs/vulntrain_%j.out
#SBATCH --error=logs/vulntrain_%j.err
#SBATCH --qos=default

set -e

source $HOME/miniconda3/etc/profile.d/conda.sh
conda activate $HOME/conda_envs/vulntrain

# --------------------------
# Parameters for the trainer
# --------------------------
BASE_MODEL=roberta-base
DATASET_ID=CIRCL/vulnerability-scores
RESULT_REPO_ID=CIRCL/vulnerability-severity-classification-roberta-base
RESULT_SAVE_DIR=$HOME/models/vulntrain_roberta


# --------------------------
# NCCL configuration
# --------------------------
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=1
export NCCL_P2P_LEVEL=NVL

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MASTER_ADDR=$(hostname)

# Optional but recommended
export HF_HOME=${SLURM_TMPDIR:-$HOME}/hf_cache
mkdir -p $HF_HOME

torchrun --nproc_per_node=$SLURM_NTASKS \
         --master_port=29500 \
         $HOME/conda_envs/vulntrain/bin/vulntrain-train-severity-classification \
            --base-model $BASE_MODEL \
            --dataset-id $DATASET_ID \
            --repo-id $RESULT_REPO_ID \
            --model-save-dir $RESULT_SAVE_DIR \
            --no-push \
            --no-cache

Example of multi-node configuration#

{{< callout type=”warning” >}} The multi-node configuration is less reliable and has caused various issues in practice (NCCL timeouts, rank synchronization failures, inconsistent checkpoint saving across nodes). The single-node configuration above works well and is recommended unless you specifically need to scale beyond the GPUs available on a single node. {{< /callout >}}

#!/bin/bash
#SBATCH --job-name=vulntrain
#SBATCH --account=<-your-account-id->
#SBATCH --partition=gpu
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=10:00:00
#SBATCH --output=logs/vulntrain_%j_%N.out
#SBATCH --error=logs/vulntrain_%j_%N.err
#SBATCH --qos=default

set -e

# -------------------------------
# Activate Conda environment
# -------------------------------
source $HOME/miniconda3/etc/profile.d/conda.sh
conda activate $HOME/conda_envs/vulntrain

# --------------------------
# Parameters for the trainer
# --------------------------
BASE_MODEL=roberta-base
DATASET_ID=CIRCL/vulnerability-scores
RESULT_REPO_ID=CIRCL/vulnerability-severity-classification-roberta-base
RESULT_SAVE_DIR=$HOME/models/vulntrain_roberta

# --------------------------
# NCCL configuration
# --------------------------
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=1        # if no InfiniBand, keep it disabled
export NCCL_P2P_LEVEL=NODE      # change from NVL to NODE for multi-node

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK

# Optional but recommended
export HF_HOME=${SLURM_TMPDIR:-$HOME}/hf_cache
mkdir -p $HF_HOME

torchrun --nnodes=$SLURM_JOB_NUM_NODES \
         --nproc_per_node=$SLURM_NTASKS_PER_NODE \
         --node_rank=$SLURM_NODEID \
         --master_addr=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1) \
         --master_port=29500 \
         $HOME/conda_envs/vulntrain/bin/vulntrain-train-severity-classification \
            --base-model $BASE_MODEL \
            --dataset-id $DATASET_ID \
            --repo-id $RESULT_REPO_ID \
            --model-save-dir $RESULT_SAVE_DIR \
            --no-cache \
            --no-push

4. Submit and monitor the job#

Submit the job from the directory where run_vulntrain.slurm is stored:

sbatch run_vulntrain.slurm

To follow progress, use your cluster’s standard tools, for example:

squeue -u $USER to see queued and running jobs.
tail -f vulntrain-<jobid>.out to stream the job’s output log.

Display accounting data and job steps in the Slurm job accounting log or Slurm database:

sacct -j <-jobid-> --format=JobID,JobName,Partition,AllocTRES,State,Elapsed,TotalCPU

Running VulnTrain on an HPC cluster

Contents