[tab] scripts/run-qwen3-0.6B-lora-rl.sh (shellscript)
#!/bin/bash
#
# RL Training Script with LoRA (testing LoRA checkpoint loading in RL pipeline)
#
# This script tests RL training with a LoRA-finetuned model by:
# 1. Loading a LoRA checkpoint from SFT training
# 2. Running RL training with LoRA enabled
# 3. Verifying that SGLang loads the LoRA adapter correctly

# for rerun the task
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python




set -ex

# will prevent ray from buffering stdout/stderr
export PYTHONUNBUFFERED=1
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

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

# --- 1. DYNAMIC HOST IP DETECTION ---
# Don't hardcode IP. Get the actual IP of the current node.
export HEAD_NODE_IP=$(hostname -I | awk '{print $1}')
echo "Detected Head Node IP: ${HEAD_NODE_IP}"

# --- 2. PROXY CONFIGURATION ---
# Ensure local traffic doesn't go through a corporate proxy
export no_proxy="${HEAD_NODE_IP},localhost,127.0.0.1,0.0.0.0"
export NO_PROXY="${HEAD_NODE_IP},localhost,127.0.0.1,0.0.0.0"

# --- 3. DEBUGGING & STABILITY ENV VARS ---
# Force NCCL/Distributed into a robust mode to prevent initialization hangs
export NCCL_DEBUG=INFO
export TORCH_DISTRIBUTED_DEBUG=INFO



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

# Use the same checkpoint path as run-sft-torchrun.sh
RUN_ID="run_20251222_220632"
LOAD_PATH="/fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/${RUN_ID}/checkpoints"
SAVE_PATH="/fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/${RUN_ID}/checkpoints"

CKPT_ARGS=(
   --hf-checkpoint /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
   --load ${LOAD_PATH}
   --ref-load /fast/project/HFMI_SynergyUnit/tab_model/huggingface/Qwen3-0.6B
   --save ${SAVE_PATH}
   --save-interval 50
)

ROLLOUT_ARGS=(
   --prompt-data /fast/project/HFMI_SynergyUnit/tab_model/huggingface/nemo_hf_part_jsonl_4k_tokens.parquet
   --input-key messages
   --apply-chat-template
   --rollout-shuffle
   --rm-type random
   --num-rollout 20
   --rollout-batch-size 8
   --n-samples-per-prompt 4
   --rollout-max-response-len 512
   --rollout-temperature 0.8
   --global-batch-size 32
)

GRPO_ARGS=(
   --use-kl-loss
   --advantage-estimator grpo
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --kl-coef 0.00
   --entropy-coef 0.00
   --eps-clip 0.2
   --eps-clip-high 0.28
)

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-6
   --lr-decay-style constant
   --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 qwen3-0.6b-lora-rl-test
)

SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 1
   --sglang-mem-fraction-static 0.75
   --sglang-decode-log-interval 100
   --sglang-chunked-prefill-size 4096
   --sglang-attention-backend fa3
)

TRAIN_BACKEND_ARGS=(
   --train-backend fsdp
   --update-weight-buffer-size 536870912
   --attn-implementation flash_attention_3
   --train-env-vars '{"PYTORCH_CUDA_ALLOC_CONF":"expandable_segments:True"}'
)

PERF_ARGS=(
   --use-dynamic-batch-size
   --max-tokens-per-gpu 9216
)

MISC_ARGS=(
   --actor-num-nodes 1
   --actor-num-gpus-per-node 2
   --colocate
   --dump-details /fast/project/HFMI_SynergyUnit/tab_model/huggingface/shared_data/qwen3-0.6b-lora-rl-test/dump_details
)

# launch the master node of ray in container - 8 GPUs for training
export MASTER_ADDR=${MASTER_ADDR:-"${HEAD_NODE_IP}"}
python3 -m ray.scripts.scripts start --head \
    --node-ip-address=${HEAD_NODE_IP} \
    --num-gpus 2 \
    --num-cpus 4 \
    --memory=214748364800 \
    --disable-usage-stats \
    --dashboard-host=0.0.0.0 \
    --dashboard-port=8265 \
    --port=6379

echo "Ray started. Waiting for Dashboard to be ready..."

# --- 4. WAIT FOR DASHBOARD (FIX FOR 504 ERROR) ---
# Loop until the dashboard port accepts connections
for i in {1..30}; do
    if curl -s "http://${HEAD_NODE_IP}:8265" > /dev/null; then
        echo "Dashboard is up!"
        break
    fi
    echo "Waiting for Ray Dashboard..."
    sleep 2
done
# Add a small safety buffer
sleep 5

# Build runtime env
RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"PYTHONPATH\": \"${SCRIPT_DIR}\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
  }
}"

python3 -m ray.scripts.scripts job submit --address="http://${HEAD_NODE_IP}:8265" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
   -- python3 train.py \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${LORA_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${GRPO_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${SGLANG_ARGS[@]} \
   ${TRAIN_BACKEND_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${MISC_ARGS[@]}