Running Slurm batch jobs

​

How to run a batch job

1. Connect to a login node.
2. Create a batch script and name it, for example, my_ml_job.sh. Here is a basic example of a job script that runs your training application written in Python:
   #!/bin/bash
   
   # Directives that define the job's settings
   #SBATCH --job-name=my_ml_job
   #SBATCH --output=%x_%j.out      # <job-name>_<job-id>.out
   #SBATCH --error=%x_%j.err       # <job-name>_<job-id>.err
   #SBATCH --time=01:00:00         # time limit = 1 hour
   #SBATCH --exclusive             # allocate all CPUs on nodes
   #SBATCH --gpus-per-node=8       # allocate 8 GPUs per node
   #SBATCH --ntasks-per-node=8     # launch 8 tasks in each job step by default
   
   # Commands that the batch script runs
   export CHECKPOINT_PATH="/mnt/data/gpt3/checkpoints"
   export DATA_PATH="/mnt/data/c4_data"
   export DTYPE="fp16"
   
   source .venv/bin/activate
   
   # Launch a job step on 4 nodes with 8 tasks per node, each pinned to 16 CPUs.
   srun --cpus-per-task=16 python train.py
   
   For more details about writing batch scripts, see Job configuration and Examples.
3. Prepare the environment on the login node.
   In Soperator clusters, all nodes share a root filesystem, so files and dependencies that you set up on the login node automatically appear on other nodes.
   1. Make sure that the files that the script uses are on the login node: create them from scratch, upload them from another machine or copy their contents to files on the login node. The example above requires train.py to be in the same directory as the batch script.
   2. Install the script dependencies.

      For example, if your workload uses Python packages listed in requirements.txt, create a Python virtual environment and install the packages into it:

      python -m venv .venv
      source .venv/bin/activate
      pip install -r requirements.txt
      

      You can make working with dependencies easier by turning your workload into a containerized job. For more details, see Running jobs in containers in Soperator clusters.
   3. Make sure all files have suitable permissions.
4. Define the required environment variables. Usually, you can put your environment variables right in the batch script, but there are cases when it’s more convenient to define them inside your login shell.

   For example, if your workload uses Weights & Biases (W&B), define an environment variable for your W&B API key:

   export WANDB_API_KEY=<W&B_API_key>
   

5. Submit the script to Slurm with sbatch, providing additional settings if needed.

   For example, if you named your batch script my_ml_job.sh and you want to run the job on 4 worker nodes, run the following command:

   sbatch --nodes=4 my_ml_job.sh
   

   The output contains the job ID, which you can use to monitor the job’s status:
   Submitted batch job 610
   

cat output_<job_ID>.txt

tail -f output_<job_ID>.txt

​

Job configuration

​

sbatch settings

1. Command parameters: sbatch --time="02:00:00" The sbatch command has parameters that work and behave like regular parameters of a Linux command. Use command parameters when you want to use different values for different job runs or to overwrite settings from other sources.
2. Input environment variables: export SBATCH_TIMELIMIT="01:00:00" in login shell, /etc/profile or ~/profile Most (but not all) command parameters have corresponding input environment variables with the same meaning. Some of these environment variables have different names from their command parameters.

   For example, the --time command parameter corresponds to the SBATCH_TIMELIMIT input environment variable.

   For the list of all environment variables, see the Slurm documentation. You can define them on different levels, from highest to lowest priority:
   1. Session variables: environment variables that you export in your login shell. They disappear after you reconnect to the cluster. Use session variables when you want to launch a few jobs with the same setting without keeping it forever.
   2. User variables: environment variables from ~/profile. The login shell exports them every time you connect to the cluster; after that, session variables can override them. Use user variables to set your personal default values for settings.
   3. Cluster variables: environment variables from /etc/profile. The login shell exports them every time any user connects to the cluster; after that, user variables and then session variables can override them. Use cluster variables to set default values for all users.
3. #SBATCH directives: #SBATCH --time="01:30:00" in a batch script #SBATCH directives are special comments at the beginning of a batch script. Use #SBATCH directives to define settings that should apply to most runs of this particular job.
   All #SBATCH directives must appear at the beginning of the script. As soon as sbatch encounters the first line that does not start with # or consist of whitespace (spaces, tab characters, etc.), it interprets the rest of the file as commands and silently ignores misplaced #SBATCH directives.
4. Slurm defaults: time="00:30:00" in ~/.slurm/defaults Slurm defaults are key-value pairs in the ~/.slurm/defaults. These values apply only to you and are not applicable to other users. Use Slurm defaults when you need a low-priority default value, or a default value for a command parameter that does not have a corresponding input environment variable.
   ~/.slurm/defaults defines settings for both sbatch and srun. Do not use them for settings that have the same name but different meaning for these commands; for example, --exclusive (sbatch, srun).

​

Common sbatch settings

​

Commands

In the example above, srun python train.py makes python train.py run in parallel. If it was just python train.py, the Python script would run on one worker node.

• SLURM_JOB_NODELIST is the list of worker nodes allocated to the job; SLURM_NNODES is the number of the worker nodes.
• SLURM_SUBMIT_DIR is the directory where you executed sbatch.
• SLURM_NODEID is the node ID for each worker node.

​

Examples

​

Training with Hugging Face Accelerate

#!/bin/bash

#SBATCH --job-name=llm_training
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --exclusive
#SBATCH --mem=0

# Get the hostname of worker node where this sbatch script runs,
# SLURMD_NODENAME is the Slurm output environment variable for sbatch.
MAIN_PROCESS_ADDR=$SLURMD_NODENAME
MAIN_PROCESS_PORT=12345

srun \
  --cpus-per-task=64 \
  --hint=nomultithread \
  bash -c 'accelerate launch \
    # SLURM_STEP_NUM_NODES is the number of worker nodes allocated for the step (2)
    --num_machines $SLURM_STEP_NUM_NODES \
    # SLURM_NODEID is the ID of the current worker node (0 or 1)
    --machine_rank $SLURM_NODEID \
    --main_process_ip $MAIN_PROCESS_ADDR \
    --main_process_port $MAIN_PROCESS_PORT \
    --num_processes $(($SLURM_STEP_NUM_NODES * $SLURM_GPUS_ON_NODE)) \
    train.py'

​

Fine-tuning with PyTorch (torchrun)

• sbatch.sh:
  #!/bin/bash
  
  #SBATCH --job-name=llama-finetune
  #SBATCH --nodes=3
  #SBATCH --output=O-%x_%j.txt
  #SBATCH --error=E-%x_%j.txt
  #SBATCH --gres=gpu:8
  #SBATCH --cpus-per-task=120
  # Loads all environment variables into the job
  #SBATCH --export=ALL
  
  srun srun.sh
  
• srun.sh:
  #!/bin/bash
  
  export GPUS_PER_NODE=8
  # Get the hostname of the first node from the list of the job's worker nodes
  # (SLURM_JOB_NODELIST, set by Slurm)
  HOST_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
  MAIN_PROCESS_PORT=12345
  
  echo "SLURM_NNODES=$SLURM_NNODES"
  echo "SLURM_NODEID=$SLURM_NODEID"
  echo "HOST_ADDR=$HOST_ADDR"
  echo "MAIN_PROCESS_PORT=$MAIN_PROCESS_PORT"
  
  source .venv/bin/activate
  
  torchrun --nnodes $SLURM_NNODES \
    --nproc_per_node $GPUS_PER_NODE \
    --master_addr $HOST_ADDR \
    --master_port $MAIN_PROCESS_PORT \
    --node_rank=$SLURM_NODEID \
    finetuning.py \
    --model_name ./llama-3-8b \
    --output_dir saved_peft_model \
    --use_peft \
    --peft_method lora \
    --enable_fsdp \
    --use_fast_kernels \
    --use_wandb \
    --dataset alpaca_dataset
  
  For details about torchrun, see the PyTorch documentation and the source code of torch.distributed.run on GitHub.

​

Containerized jobs

srun --container-image="nvcr.io#nvidia/tensorflow:23.02-tf1-py3" \
  python -c "import tensorflow as tf; print (tf.__version__)"