prain: README; example config with the measured eGPU defaults (miss auto <=2 tokens, 16 CPU threads, no pinned RAM tier)
This commit is contained in:
@@ -0,0 +1,45 @@
|
|||||||
|
# prain
|
||||||
|
|
||||||
|
Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places
|
||||||
|
each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk.
|
||||||
|
Missing experts are copied to the GPU or computed on the CPU (`[miss]`). Eviction, heat, prefetch
|
||||||
|
and miss policies are configurable, and `prain tune` picks them from a routing trace.
|
||||||
|
Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 GB RX 6600 XT.
|
||||||
|
|
||||||
|
## Build (GPU host, user space)
|
||||||
|
|
||||||
|
scripts/build.sh wrappers sdk fetch torch torchvision vllm prain
|
||||||
|
|
||||||
|
- `sdk`: venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
|
||||||
|
- `torch`, `torchvision`: from source with clang and `-march=native` (hours)
|
||||||
|
- `vllm`: v0.32.0 from source with gfx1032 added and `patches/vllm-*.patch` applied
|
||||||
|
- `prain`: editable install and the test suite
|
||||||
|
|
||||||
|
Check the stack: `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`
|
||||||
|
|
||||||
|
## Serve
|
||||||
|
|
||||||
|
. scripts/env.sh
|
||||||
|
prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --served-model-name gemma4 \
|
||||||
|
--language-model-only --max-model-len 4096 --max-num-batched-tokens 1024 \
|
||||||
|
--enable-prefix-caching --attention-backend TRITON_ATTN
|
||||||
|
|
||||||
|
The first start writes the pack (~3 min). Chat with `prain tui`, or `prain web --config FILE` for
|
||||||
|
several users with drafts.
|
||||||
|
|
||||||
|
## Tune for a machine
|
||||||
|
|
||||||
|
python scripts/measure.py PACK_DIR/layer029.w13 --threads 16 > cost.json # server stopped
|
||||||
|
prain sim TRACE CONFIG MODEL_DIR/config.json --cost cost.json
|
||||||
|
prain tune TRACE MODEL_DIR/config.json --cost cost.json --expert-vram-gib 1.67 --ram-gib 0 \
|
||||||
|
--pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
|
||||||
|
python scripts/bench.py http://127.0.0.1:8000
|
||||||
|
|
||||||
|
The simulator treats every disk-tier expert as a disk read; the page cache makes real runs faster.
|
||||||
|
|
||||||
|
## vLLM patches
|
||||||
|
|
||||||
|
- `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"
|
||||||
|
checks took it for Blackwell (B200 attention tiles: 128 KiB of LDS on a 64 KiB GPU).
|
||||||
|
- `vllm-0.32-gemma4-fit-8gb`: Gemma 4 shares `embed_tokens` as `lm_head` instead of allocating a
|
||||||
|
second 1.4 GiB copy while loading, and sizes rope caches by `max_model_len` (saves ~0.4 GiB).
|
||||||
@@ -1,6 +1,9 @@
|
|||||||
# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
|
# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
|
||||||
# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
|
# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
|
||||||
# VRAM: 8 GiB - 4.45 GiB dense weights - 0.75 GiB KV cache - ~0.9 GiB runtime = ~1.7 GiB experts.
|
# VRAM, measured: 6.22 GiB weights (4.46 dense + 5 slots x 30 layers) + 0.75 GiB KV cache: 7.9 of
|
||||||
|
# 8 GiB in use. Serve with --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than RDNA2 has).
|
||||||
|
# RAM tier: pinning is very slow behind a strict IOMMU (Thunderbolt eGPU: minutes per GiB), so this
|
||||||
|
# file keeps no RAM tier; the kernel page cache holds the hot part of the pack instead.
|
||||||
#
|
#
|
||||||
# Tiers: vram = pinned in a VRAM slot for the whole run
|
# Tiers: vram = pinned in a VRAM slot for the whole run
|
||||||
# ram = pinned host memory, copied into a cache slot when a token needs it
|
# ram = pinned host memory, copied into a cache slot when a token needs it
|
||||||
@@ -38,10 +41,11 @@ budget = 2 # experts per layer and forward
|
|||||||
min_prob = 0.05 # skip guesses below this router probability
|
min_prob = 0.05 # skip guesses below this router probability
|
||||||
|
|
||||||
[miss] # an expert that is not in VRAM
|
[miss] # an expert that is not in VRAM
|
||||||
mode = "copy" # copy | cpu (always compute it on the CPU) | auto (CPU if it has
|
mode = "auto" # copy | cpu (always compute it on the CPU) | auto (CPU if it has
|
||||||
cpu_max_tokens = 1 # at most cpu_max_tokens tokens, else copy)
|
cpu_max_tokens = 2 # at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
|
||||||
cpu_threads = 4 # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
|
# eGPU host: a copy costs ~4.6 ms + disk read, the CPU ~2.5 ms per token
|
||||||
admit = 1 # cpu/auto: copy this many CPU-run experts into VRAM afterwards
|
cpu_threads = 16 # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
|
||||||
|
admit = 0 # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
|
||||||
|
|
||||||
[serve] # parallel requests (`prain serve` passes these to vLLM)
|
[serve] # parallel requests (`prain serve` passes these to vLLM)
|
||||||
max_requests = 2 # requests vLLM runs at once (--max-num-seqs); warm-ups count too
|
max_requests = 2 # requests vLLM runs at once (--max-num-seqs); warm-ups count too
|
||||||
|
|||||||
Reference in New Issue
Block a user