[tab] slurm/jobs/mihir/juelich/crowd-pilot/1-wave-anta-santa/qwen_32b_lora_bs_32_ctx_parallel.sh (shellscript)
#!/bin/bash
#
# Ray-free SFT Training Script
#
    
export NCCL_TIMEOUT=3600
export NCCL_ASYNC_ERROR_HANDLING=1
export TORCH_NCCL_BLOCKING_WAIT=1

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
# FIXME(f.srambical): this is hardcoded for now
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)"

RUN_ID=${RUN_ID:-$SLURM_JOBID}
RUN_NAME="qwen32B_no_lora_bs_32_${RUN_ID}"
WAVE_NAME="ANTA_CLAUS"
LOAD_PATH="/p/scratch/envcomp/mihir/huggingface/Qwen3-32B"
SAVE_PATH="/p/scratch/envcomp/mihir/huggingface/shared_data/${RUN_ID}/checkpoints"

CKPT_ARGS=(
   --hf-checkpoint /p/scratch/envcomp/mihir/huggingface/Qwen3-32B
   --load ${LOAD_PATH}
   --ref-load /p/scratch/envcomp/mihir/huggingface/Qwen3-32B
   --save ${SAVE_PATH}
   --save-interval 1000
)

SFT_ARGS=(
   --rollout-function-path miles.rollout.sft_rollout.generate_rollout
   --prompt-data /p/scratch/envcomp/mihir/huggingface/nemo_hf_part_jsonl_4k_tokens.parquet
   --val-prompt-data /p/scratch/envcomp/mihir/huggingface/nemo_hf_part_jsonl_4k_tokens_validation.parquet
   --val-interval 500
   --val-steps 100
   --input-key messages
   --apply-chat-template
   --rollout-shuffle
   --num-rollout 5000
   --rollout-batch-size 32
   --global-batch-size 32
   --context-parallel-size 8

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

LORA_ARGS=(
    --use-lora
    --lora-rank 8
    --lora-alpha 16
    --lora-dropout 0.0
    --lora-target-modules q_proj v_proj
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-4
   --lr-decay-style WSD
   --lr-wsd-decay-style linear
   --lr-warmup-iters 250
   --lr-decay-iters 5000
   --lr-wsd-decay-iters 1000
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)


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}"
)

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

PERF_ARGS=(
   --micro-batch-size 1
)

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/qwen3-600M-fsdp-1116-noref/dump_details
)
echo ${WANDB_ARGS[@]}

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[@]}