slurm.sub

#!/bin/bash
# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# =============================================================================
# NeMo AutoModel — Reference SLURM submission script
# =============================================================================
#
# Copy this file, adapt it to your cluster, and submit with sbatch:
#
#   cp slurm.sub my_cluster.sub
#   vim my_cluster.sub          # edit CONFIG, SBATCH directives, container, etc.
#   sbatch my_cluster.sub
#
# The script launches torchrun across all allocated nodes.  Each worker runs
# `automodel` which detects the torchrun environment and executes the recipe
# in-process.
# =============================================================================

# ---------------------------------------------------------------------------
# SLURM directives — change these for your cluster
# ---------------------------------------------------------------------------
#SBATCH -A <your_account>           # account
#SBATCH -p batch                    # partition
#SBATCH -t 01:00:00                 # wall time limit, hr:min:sec
#SBATCH -N 8                        # number of nodes
#SBATCH --gpus-per-node=8
#SBATCH --ntasks-per-node=1
#SBATCH -J automodel-job            # job name
#SBATCH --output=slurm_jobs/%x_%j.out
#SBATCH --error=slurm_jobs/%x_%j.err
#SBATCH --dependency=singleton
#SBATCH --signal=B:USR1@300

# ---------------------------------------------------------------------------
# Recipe config — point this at your YAML
# ---------------------------------------------------------------------------
CONFIG=examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml

echo "Running on hosts: $(echo $(scontrol show hostname))"

# ---------------------------------------------------------------------------
# Container image and mounts — edit paths for your environment
# ---------------------------------------------------------------------------
CONT=/lustre/fsw/portfolios/<team>/users/$USER/automodel.sqsh
CONT_NAME=automodel-training
CONT_MOUNT="\
/lustre/fsw/portfolios/<team>/users/$USER/:/$USER/,\
/home/$USER/.ssh:/root/.ssh,\
/home/$USER/.gitconfig:/root/.gitconfig,\
/home/$USER/Automodel:/opt/Automodel"

# ---------------------------------------------------------------------------
# Multi-node environment
# ---------------------------------------------------------------------------
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=${MASTER_PORT:-13742}
export NUM_GPUS=${SLURM_GPUS_PER_NODE:-8}
export WORLD_SIZE=$(($NUM_GPUS * $SLURM_NNODES))

# ---------------------------------------------------------------------------
# Secrets — set these or export them in your shell before sbatch
# ---------------------------------------------------------------------------
export WANDB_API_KEY=${WANDB_API_KEY:-}
export HF_TOKEN=${HF_TOKEN:-}

# ---------------------------------------------------------------------------
# Preemption signal forwarding
# ---------------------------------------------------------------------------
# Forward SIGUSR1 to the training worker processes on this node only.
# Deliberately NOT signalled: srun (would propagate to torchrun on every
# node) and torchrun itself (terminates on unhandled USR1, and slurmstepd
# then SIGKILLs the whole step cgroup).  Signalling the local workers is
# sufficient: DistributedSignalHandler all_gathers the flag, so every rank
# checkpoints at the next step boundary once any single rank sees the signal.
#
# The trap is installed before the launch so a signal arriving during
# startup cannot hit the default USR1 disposition and kill this script.
forward_preemption_signal() {
    echo "[$(date)] SIGUSR1 from SLURM; forwarding to local training workers"
    for pid in $(pgrep -u "$(id -un)" -f "nemo_automodel.cli.app"); do
        # Skip the torchrun launcher (its cmdline also contains the module name).
        if ! tr '\0' ' ' < "/proc/${pid}/cmdline" 2>/dev/null | grep -q torchrun; then
            kill -USR1 "${pid}" 2>/dev/null || true
        fi
    done
}
trap forward_preemption_signal USR1

# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
srun \
    --container-name="${CONT_NAME}" \
    --container-image="${CONT}" \
    --container-mounts="${CONT_MOUNT}" \
    --container-entrypoint \
    --no-container-mount-home \
    --export=ALL \
    bash -c "\
        cd /opt/Automodel && \
        torchrun \
            --nproc_per_node=\${SLURM_GPUS_PER_NODE:-8} \
            --nnodes=\${SLURM_NNODES:-1} \
            --rdzv_backend=c10d \
            --rdzv_endpoint=\${MASTER_ADDR}:\${MASTER_PORT} \
            -m nemo_automodel.cli.app ${CONFIG}" &
SRUN_PID=$!


# First wait returns early (128+SIGUSR1) if the trap fires; the second wait
# collects srun's real exit status after the graceful shutdown completes.
wait ${SRUN_PID}
trap '' USR1
wait ${SRUN_PID}
EXIT_CODE=$?
echo "[$(date)] srun exited with status ${EXIT_CODE}"
exit ${EXIT_CODE}