submit_inkling_generate.sbatch

#!/bin/bash
#SBATCH --job-name=inkling-generate
#SBATCH --partition=hopper-prod
#SBATCH --qos=normal
#SBATCH --nodes=4
#SBATCH --gpus-per-node=8
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=88
#SBATCH --mem=0
#SBATCH --exclusive
#SBATCH --time=01:00:00
#SBATCH --output=/fsx/merve/logs/inkling-generate-%j.out

# Usage:
#   sbatch submit_inkling_generate.sbatch
#   PROMPT="What is in this image?" IMAGE=/path/cat.png sbatch submit_inkling_generate.sbatch
#   PROMPT="Transcribe this." AUDIO=/path/clip_16k.wav sbatch submit_inkling_generate.sbatch

set -euo pipefail

MODEL_PATH=${MODEL_PATH:-/fsx/pedro/tm/models/tml-model-share}
PROCESSOR_PATH=${PROCESSOR_PATH:-${MODEL_PATH}}
PROMPT=${PROMPT:-"Explain tensor parallelism in simple terms."}
IMAGE=${IMAGE:-}
AUDIO=${AUDIO:-}
MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-512}
THINKING_EFFORT=${THINKING_EFFORT:-}
TRANSFORMERS_SRC=${TRANSFORMERS_SRC:-/fsx/merve/transformers-tm/new-model-addition-tm}
SCRIPT=${SCRIPT:-${TRANSFORMERS_SRC}/scripts/generate_inkling.py}

NNODES=${SLURM_NNODES:-4}
GPUS_PER_NODE=${GPUS_PER_NODE:-8}
MASTER_PORT=${MASTER_PORT:-29500}

module load cuda/12.8
source /fsx/merve/cluster/bin/activate

export FI_PROVIDER=efa
export HF_HOME=${HF_HOME:-/fsx/merve/hf_cache}
export PYTHONPATH="${TRANSFORMERS_SRC}/src:${PYTHONPATH:-}"
export PYTHONUNBUFFERED=1

HEAD_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
HEAD_IP=$(srun --nodes=1 --ntasks=1 -w "$HEAD_NODE" hostname -I | awk '{print $1}')

echo "Nodes      = $SLURM_JOB_NODELIST"
echo "Head node  = $HEAD_NODE ($HEAD_IP)"
echo "Model      = $MODEL_PATH"
echo "Processor  = $PROCESSOR_PATH"
echo "Prompt     = $PROMPT"
echo "Image(s)   = ${IMAGE:-<none>}"
echo "Audio(s)   = ${AUDIO:-<none>}"

SCRIPT_ARGS=(--model-path "$MODEL_PATH" --processor-path "$PROCESSOR_PATH"
             --prompt "$PROMPT" --max-new-tokens "$MAX_NEW_TOKENS")
for img in $IMAGE; do SCRIPT_ARGS+=(--image "$img"); done
for aud in $AUDIO; do SCRIPT_ARGS+=(--audio "$aud"); done
[[ -n "$THINKING_EFFORT" ]] && SCRIPT_ARGS+=(--thinking-effort "$THINKING_EFFORT")

srun --ntasks-per-node=1 torchrun \
  --nnodes="$NNODES" --nproc-per-node="$GPUS_PER_NODE" \
  --rdzv-backend=c10d --rdzv-endpoint="${HEAD_IP}:${MASTER_PORT}" --rdzv-id="$SLURM_JOB_ID" \
  "$SCRIPT" "${SCRIPT_ARGS[@]}"