File size: 2,946 Bytes
fdc6474 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | #!/usr/bin/env bash
# GLM-5.2 hybrid NVFP4+AQLM on 4x RTX PRO 6000 Blackwell (SM120, 96GB).
# Usage: ./run_rtx6000.sh [1m|500k|250k] (default: 1m)
#
# One-time setup (fresh box) — see SETUP steps below the config block.
set -euo pipefail
VARIANT="${1:-1m}"
HF_BASE="jarrelscy/GLM-5.2-NVFP4-AQLM-hybrid" # weights + code/ bundle
VLLM_FORK="https://github.com/jarrelscy/vllm-glm52-sm120" # full fork, patches committed
VLLM_BRANCH="glm52-sm120"
WORK="${WORK:-$HOME/glm52}"
case "$VARIANT" in
1m) HF_REPO="$HF_BASE"; MAXLEN=1048576 ;; # 30% hot experts, 292 GB
500k) HF_REPO="$HF_BASE-500k"; MAXLEN=524288 ;; # 48% hot experts, 334 GB
250k) HF_REPO="$HF_BASE-250k"; MAXLEN=262144 ;; # 57% hot experts, 354 GB
*) echo "usage: $0 [1m|500k|250k]"; exit 1 ;;
esac
MODEL_DIR="$WORK/models/$VARIANT"
# ---------- one-time setup ----------
if [ ! -d "$WORK/vllm/.venv" ]; then
mkdir -p "$WORK" && cd "$WORK"
command -v uv >/dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
# vLLM changes now live in the fork (no patches to apply)
git clone --depth 1 -b "$VLLM_BRANCH" "$VLLM_FORK" vllm && cd vllm
uv venv --python 3.12 .venv && source .venv/bin/activate
uv pip install vllm --torch-backend=auto # bootstrap deps fast
VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
uv pip install ninja "huggingface_hub[cli]"
# flashinfer git main: SM120 sparse-MLA needs kv_scale_format=arbitrary_fp32
uv pip install "flashinfer-python @ git+https://github.com/flashinfer-ai/flashinfer.git"
# JIT toolchain: pin nvcc/crt/nvvm/cccl to torch's CUDA minor, lib64 link
TORCH_CU=$(python -c "import torch;print(torch.version.cuda.split('.')[0]+'.'+torch.version.cuda.split('.')[1])")
uv pip install "nvidia-cuda-nvcc==${TORCH_CU}.*" "nvidia-cuda-crt==${TORCH_CU}.*" \
"nvidia-nvvm==${TORCH_CU}.*" "nvidia-cuda-cccl==${TORCH_CU}.*" || true
CU13="$WORK/vllm/.venv/lib/python3.12/site-packages/nvidia/cu13"
ln -sfn lib "$CU13/lib64" 2>/dev/null || true
for f in "$CU13"/lib/lib*.so.*; do b=$(basename "$f"); s=${b%%.so.*};
[ -e "$CU13/lib/$s.so" ] || ln -s "$b" "$CU13/lib/$s.so"; done
fi
# ---------- fetch weights ----------
if [ ! -f "$MODEL_DIR/config.json" ]; then
source "$WORK/vllm/.venv/bin/activate"
hf download "$HF_REPO" --exclude "code/*" --local-dir "$MODEL_DIR"
fi
# ---------- serve ----------
cd "$WORK/vllm" && source .venv/bin/activate
export CUDA_HOME="$WORK/vllm/.venv/lib/python3.12/site-packages/nvidia/cu13"
export VLLM_PP_LAYER_PARTITION="21,19,19,19"
export NCCL_MAX_NCHANNELS=4 NCCL_BUFFSIZE=1048576
export VLLM_SPARSE_INDEXER_MAX_LOGITS_MB=256
exec vllm serve "$MODEL_DIR" \
--pipeline-parallel-size 4 \
--gpu-memory-utilization 0.95 \
--kv-cache-dtype fp8_ds_mla \
--max-model-len "$MAXLEN" \
--max-num-seqs 2 \
--max-num-batched-tokens 2048 \
--enforce-eager \
--port 8000
|