File size: 2,010 Bytes
48db85f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | #!/bin/bash
#SBATCH --job-name=mxf_rl_drgrpo
#SBATCH --partition=general
#SBATCH --qos=high
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=8
#SBATCH --mem=256G
#SBATCH --time=48:00:00
#SBATCH --no-requeue
#SBATCH --output=/workspace-vast/celeste/maxact-fast/logs/%x_%j.out
# Stage 6: Dr. GRPO RL (scripts/rl.py) — colocated single task: vLLM TP across all GPUs at
# gpu_memory_utilization=0.35, HF actor+LoRA on cuda:0. `mkdir -p logs` once before first submit.
#
# sbatch scripts/sbatch_rl.sh # RLConfig defaults
# sbatch scripts/sbatch_rl.sh --groups-per-step 64 --sync-every 4 # extra args pass through
#
# Venv (dedicated — vllm pins are hard, see pyproject [vllm] extra). Build once with:
# uv venv /workspace-vast/celeste/envs/mxf-vllm --python 3.12
# uv pip install --python /workspace-vast/celeste/envs/mxf-vllm/bin/python \
# "vllm==0.19.0" "vllm-lens==1.1.0" "transformers==4.57.1" peft wandb --torch-backend=cu128
# Pins are load-bearing: vllm-lens 1.1.0 is built against vLLM 0.19.0 — on vLLM>=0.22 the hook
# crashes then SILENTLY skips injection; cu128 matches the cluster's driver 570 (cu130 wheels
# fail at import). transformers must stay <5 (apply_chat_template API break).
# /workspace-vast/celeste/envs/vllm-lens (nla-experiments) is a known-good fallback with the same pins.
set -euo pipefail
ROOT=/workspace-vast/celeste/maxact-fast
VENV=${VENV:-/workspace-vast/celeste/envs/mxf-vllm}
source "$VENV/bin/activate"
export VLLM_ALLOW_INSECURE_SERIALIZATION=1 # pickle for apply_model(partial) weight sync
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export HF_HOME=/workspace-vast/pretrained_ckpts
export HF_TOKEN_PATH=/workspace-vast/celeste/.cache/huggingface/token
: "${WANDB_API_KEY:?set WANDB_API_KEY in your shell}"
export PYTHONUNBUFFERED=1
export PYTHONPATH=$ROOT/src:${PYTHONPATH:-}
# NEVER set CUDA_VISIBLE_DEVICES here — SLURM does it.
cd "$ROOT"
python scripts/rl.py --tp "${SLURM_GPUS_ON_NODE:-1}" "$@"
|