[tab] slurm/jobs/mihir/juelich/crowd-pilot/2-wave-broken-resolutions/run_template.sh (shellscript)
#!/bin/bash
#
# Ray-free SFT Training Script - Template Version
# Configurable via environment variables:
#   - MODEL_SIZE: 1.7b, 4b, or 8b
#   - BATCH_SIZE: 8, 16, or 32
#   - USE_LORA: true or false
#

set -euo pipefail

# ============================================
# Validate required environment variables
# ============================================
if [[ -z "${MODEL_SIZE:-}" ]]; then
    echo "ERROR: MODEL_SIZE not set. Must be one of: 1.7b, 4b, 8b"
    exit 1
fi

if [[ -z "${BATCH_SIZE:-}" ]]; then
    echo "ERROR: BATCH_SIZE not set. Must be one of: 8, 16, 32"
    exit 1
fi

if [[ -z "${USE_LORA:-}" ]]; then
    echo "ERROR: USE_LORA not set. Must be: true or false"
    exit 1
fi

# ============================================
# Environment setup
# ============================================
export NCCL_TIMEOUT=3600
export NCCL_ASYNC_ERROR_HANDLING=1
export TORCH_NCCL_BLOCKING_WAIT=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# ============================================
# Cluster configuration
# ============================================
GPUS_PER_NODE=4
NUM_NODES=${SLURM_JOB_NUM_NODES}
NODE_RANK=${SLURM_NODEID}
MASTER_ADDR=${MASTER_ADDR:-$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)}

NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"

# ============================================
# Model-specific configuration
# ============================================
case "${MODEL_SIZE}" in
    1.7b|1.7B)
        MODEL_NAME="Qwen3-1.7B"
        MODEL_SIZE_LABEL="1.7B"
        ;;
    4b|4B)
        MODEL_NAME="Qwen3-4B"
        MODEL_SIZE_LABEL="4B"
        ;;
    8b|8B)
        MODEL_NAME="Qwen3-8B"
        MODEL_SIZE_LABEL="8B"
        ;;
    *)
        echo "ERROR: Invalid MODEL_SIZE '${MODEL_SIZE}'. Must be one of: 1.7b, 4b, 8b"
        exit 1
        ;;
esac

# ============================================
# LoRA-specific configuration
# ============================================
if [[ "${USE_LORA}" == "true" ]]; then
    LORA_LABEL="lora"
    LR="1e-4"
else
    LORA_LABEL="no_lora"
    LR="1e-5"
fi

# ============================================
# Run identification
# ============================================
RUN_ID=${RUN_ID:-$SLURM_JOBID}
RUN_NAME="qwen${MODEL_SIZE_LABEL}_${LORA_LABEL}_bs_${BATCH_SIZE}_${RUN_ID}"
WAVE_NAME="BROKEN_RESOLUTIONS"

# ============================================
# Paths
# ============================================
MODEL_PATH="/p/scratch/envcomp/mihir/huggingface/${MODEL_NAME}"
LOAD_PATH="${MODEL_PATH}"
SAVE_PATH="/p/scratch/envcomp/mihir/huggingface/shared_data/${RUN_ID}/checkpoints"

# ============================================
# Dynamic Hyperparameters based on Batch Size
# ============================================
VAL_INTERVAL=$((16000 / BATCH_SIZE))
NUM_ROLLOUT=$((160000 / BATCH_SIZE))
LR_WARMUP_ITERS=$((8000 / BATCH_SIZE))
LR_DECAY_ITERS=$((160000 / BATCH_SIZE))
LR_WSD_DECAY_ITERS=$((32000 / BATCH_SIZE))

echo "============================================"
echo "Run Configuration:"
echo "  MODEL_SIZE:  ${MODEL_SIZE}"
echo "  MODEL_NAME:  ${MODEL_NAME}"
echo "  BATCH_SIZE:  ${BATCH_SIZE}"
echo "  USE_LORA:    ${USE_LORA}"
echo "  CP_SIZE:     ${CONTEXT_PARALLEL_SIZE:-1}"
echo "  VAL_INT:     ${VAL_INTERVAL}"
echo "  NUM_ROLL:    ${NUM_ROLLOUT}"
echo "  LR_WARMUP:   ${LR_WARMUP_ITERS}"
echo "  RUN_NAME:    ${RUN_NAME}"
echo "  RUN_ID:      ${RUN_ID}"
echo "  WAVE_NAME:   ${WAVE_NAME}"
echo "  MODEL_PATH:  ${MODEL_PATH}"
echo "  SAVE_PATH:   ${SAVE_PATH}"
echo "============================================"

# ============================================
# Checkpoint arguments
# ============================================
CKPT_ARGS=(
   --hf-checkpoint ${MODEL_PATH}
   --load ${LOAD_PATH}
   --ref-load ${MODEL_PATH}
   --save ${SAVE_PATH}
   --save-interval 1000
)

# ============================================
# SFT arguments
# ============================================
DATASET_DIR=/p/scratch/envcomp/mihir/huggingface/dataset_05_01_2026/qwen/miles_hf_part_jsonl_4k_tokens
SFT_ARGS=(
   --rollout-function-path miles.rollout.sft_rollout.generate_rollout
   --prompt-data $DATASET_DIR/training.jsonl
   --val-prompt-data $DATASET_DIR/validation.jsonl
   --val-interval ${VAL_INTERVAL}
   --val-steps 100
   --input-key messages
   --apply-chat-template
   --rollout-shuffle
   --num-rollout ${NUM_ROLLOUT}
   --rollout-batch-size ${BATCH_SIZE}
   --global-batch-size ${BATCH_SIZE}
   --context-parallel-size ${CONTEXT_PARALLEL_SIZE:-1}

   --loss-type sft_loss
   --calculate-per-token-loss
   --disable-compute-advantages-and-returns
)

# ============================================
# LoRA arguments (conditionally added)
# ============================================
LORA_ARGS=()
if [[ "${USE_LORA}" == "true" ]]; then
    LORA_ARGS=(
        --use-lora
        --lora-rank 8
        --lora-alpha 16
        --lora-dropout 0.0
        --lora-target-modules q_proj v_proj
    )
fi

# ============================================
# Optimizer arguments
# ============================================
OPTIMIZER_ARGS=(
   --optimizer adam
   --lr ${LR}
   --lr-decay-style WSD
   --lr-wsd-decay-style linear
   --lr-warmup-iters ${LR_WARMUP_ITERS}
   --lr-decay-iters ${LR_DECAY_ITERS}
   --lr-wsd-decay-iters ${LR_WSD_DECAY_ITERS}
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)

# ============================================
# Wandb arguments
# ============================================
WANDB_ARGS=(
   --use-wandb
   --wandb-project crowd-pilot-miles
   --wandb-team instant-uv
   --wandb-group $RUN_NAME
   --wandb-run-id $RUN_ID
   --wandb-dir "/p/scratch/envcomp/logs/${WAVE_NAME}"
)

# ============================================
# Training backend arguments
# ============================================
TRAIN_BACKEND_ARGS=(
   --train-backend fsdp
   --update-weight-buffer-size 536870912
   --attn-implementation flash_attention_2
)

# ============================================
# Performance arguments
# ============================================
PERF_ARGS=(
   --micro-batch-size 1
)

# ============================================
# Miscellaneous arguments
# ============================================
MISC_ARGS=(
   --rollout-max-context-len 4100
   --rollout-max-prompt-len 4096
   --rollout-max-response-len 4100
   --dump-details /p/scratch/envcomp/mihir/huggingface/shared_data/${RUN_ID}/dump_details
)

echo "WANDB_ARGS: ${WANDB_ARGS[@]}"

# ============================================
# Launch training
# ============================================
torchrun \
    --nproc_per_node=${GPUS_PER_NODE} \
    --nnodes=${NUM_NODES} \
    --node_rank=${NODE_RANK} \
    --master_addr=${MASTER_ADDR} \
    --master_port=${MASTER_PORT:-29500} \
    train_sft.py \
    ${CKPT_ARGS[@]} \
    ${SFT_ARGS[@]} \
    ${LORA_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
    ${WANDB_ARGS[@]} \
    ${TRAIN_BACKEND_ARGS[@]} \
    ${PERF_ARGS[@]} \
    ${MISC_ARGS[@]}