Joey Grasty
4edb2135be
- utils/vllm_manager.py: drop --disable-log-requests, removed in vLLM 0.26.0 - core/ftpo_trainer.py: pass token_type_ids to the 3 model forward calls in compute_loss -- transformers 5.5.0's Gemma3 requires it during training for causal-mask construction (Gemma3 is natively multimodal) - configs/gemma-3-4b-it.yaml: lower vllm_gpu_memory_utilization 0.85->0.5, since the DGX Spark's 121GB is unified CPU/GPU memory and the default starved the OS, causing swap thrashing - DGX_SPARK_SETUP.md: full writeup of the above plus the parts that don't live in this repo (two-conda-env split to resolve a vllm/unsloth transformers version conflict, flash-attn source build flags, torch/CUDA version matching, ~/.triton/cache permissions) The antislop-vllm submodule also needed a one-line fix (removing an invalid reference_compile kwarg in utils/refusal_detector.py that was silently disabling refusal filtering) -- documented in DGX_SPARK_SETUP.md rather than committed as a submodule pointer change, since we don't have push access to upstream's antislop-vllm repo. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
222 lines
10 KiB
Markdown
222 lines
10 KiB
Markdown
# Running auto-antislop on a DGX Spark
|
||
|
||
Notes from getting [auto-antislop](https://github.com/sam-paech/auto-antislop) working end-to-end
|
||
(generation → slop analysis → FTPO fine-tuning) on an NVIDIA DGX Spark. The upstream quickstart
|
||
assumes a conventional x86_64 box with a discrete-VRAM GPU; the Spark's aarch64 CPU, Blackwell
|
||
(GB10) GPU, and **unified** CPU/GPU memory break several of its assumptions. This doc covers what
|
||
to change and why.
|
||
|
||
## Hardware/platform recap
|
||
|
||
- CPU/GPU: aarch64, NVIDIA GB10 (Blackwell, compute capability sm_121)
|
||
- Memory: 121GB **unified** — CPU and GPU draw from the same pool, unlike a normal discrete GPU
|
||
- Driver: reports CUDA 13.0 as the max supported runtime; system `nvcc` is also 13.0
|
||
|
||
The unified memory is the single biggest thing to keep in mind — anywhere upstream defaults assume
|
||
"GPU memory" is separate from "the OS's memory," that assumption is wrong here.
|
||
|
||
## 1. Two conda environments, not one
|
||
|
||
`vllm` 0.26.0 hard-pins `transformers>=5.5.3`. `unsloth` (needed for FTPO fine-tuning) caps it at
|
||
`<=5.5.0`. There is no single transformers version that satisfies both.
|
||
|
||
The fix: since `main.py` launches vLLM as a **subprocess** (`vllm serve ...` via `PATH` lookup, not
|
||
an in-process import — see `utils/vllm_manager.py`), it doesn't actually need to share an
|
||
interpreter with unsloth/transformers at all. Give it its own env:
|
||
|
||
```bash
|
||
conda create -n antislop python=3.11 -y
|
||
conda activate antislop
|
||
pip install "torch==2.11.0" # see §2 re: which index
|
||
pip install -r requirements.txt # minus flash-attn, see §3
|
||
# transformers ends up needing to be pinned to <=5.5.0 for unsloth — see below
|
||
|
||
conda create -n antislop-vllm-serve python=3.11 -y
|
||
conda activate antislop-vllm-serve
|
||
pip install vllm # pulls its own compatible torch/transformers
|
||
```
|
||
|
||
Then expose *only* the `vllm` binary from the second env on `PATH`, without shadowing `python`:
|
||
|
||
```bash
|
||
mkdir -p ~/.local/bin
|
||
cat > ~/.local/bin/vllm << 'EOF'
|
||
#!/bin/bash
|
||
exec /home/<user>/miniconda3/envs/antislop-vllm-serve/bin/vllm "$@"
|
||
EOF
|
||
chmod +x ~/.local/bin/vllm
|
||
```
|
||
|
||
Make sure `~/.local/bin` is on `PATH`. From inside the `antislop` env, `which vllm` should resolve
|
||
to the wrapper while `which python` stays in `antislop`.
|
||
|
||
In the `antislop` env, after `pip install -r requirements.txt` pulls in `unsloth`/`trl`/etc., pin
|
||
transformers back down:
|
||
|
||
```bash
|
||
pip install "transformers<=5.5.0,>=4.51.3"
|
||
```
|
||
|
||
## 2. Match torch's CUDA build to the *system* CUDA toolkit
|
||
|
||
`pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu128` installs a CUDA
|
||
12.8 build. The Spark's system `nvcc` (used later to build flash-attn from source) is CUDA 13.0.
|
||
Building any CUDA extension against a mismatched torch build fails immediately with:
|
||
|
||
```
|
||
RuntimeError: The detected CUDA version (13.0) mismatches the version that was used to compile
|
||
PyTorch (12.8). Please make sure to use the same CUDA versions.
|
||
```
|
||
|
||
Fix: install torch from **plain PyPI** (no `--index-url` override) — for aarch64 it resolves to a
|
||
cu130 build that matches the system toolkit:
|
||
|
||
```bash
|
||
pip install "torch==2.11.0" "torchvision==0.26.0" --force-reinstall
|
||
```
|
||
|
||
Also watch for `torchvision` being pulled by a downstream dependency (e.g. installing `vllm` into
|
||
the wrong env) with a different CUDA tag than `torch` — same mismatch error, same fix (reinstall
|
||
matching versions from the same index).
|
||
|
||
## 3. Build flash-attn from source, restricted to one arch, low parallelism
|
||
|
||
There's no aarch64 wheel for `flash-attn` on PyPI — it always builds from source here.
|
||
|
||
- **Restrict architectures.** flash-attn's default build targets `sm_80;90;100;120` — four separate
|
||
kernel variants. GB10 is `sm_121`, which runs fine against `sm_120` binaries (Blackwell
|
||
family-compatible), so there's no reason to build the other three:
|
||
```bash
|
||
FLASH_ATTN_CUDA_ARCHS=120 MAX_JOBS=6 pip install -U --no-build-isolation \
|
||
"git+https://github.com/Dao-AILab/flash-attention.git@v2.8.3.post1#egg=flash_attn"
|
||
```
|
||
- **Lower `MAX_JOBS`.** The default (one job per core — 20 here) OOM-killed the build: too many
|
||
parallel `nvcc` processes each compiling flash-attn's heavy templated CUDA kernels. `MAX_JOBS=6`
|
||
with the single-arch restriction above completed cleanly in a few minutes.
|
||
- Install `wheel ninja packaging cmake` first (`pip install -U wheel ninja packaging cmake`).
|
||
|
||
## 4. `~/.triton/cache` may be root-owned
|
||
|
||
If Triton (used by vLLM's compilation backend) was ever invoked as root on the machine before
|
||
(e.g. during initial OS/image setup), `~/.triton/cache` and its subdirectories can end up owned by
|
||
`root`, which breaks any user-level process trying to JIT-compile through it:
|
||
|
||
```
|
||
PermissionError: [Errno 13] Permission denied: '/home/<user>/.triton/cache/.../tmp.pid_...'
|
||
```
|
||
|
||
Fix (needs an interactive sudo prompt, so run it yourself, not from an agent/script):
|
||
|
||
```bash
|
||
sudo chown -R $USER:$USER ~/.triton
|
||
```
|
||
|
||
## 5. Lower `vllm_gpu_memory_utilization` — this is not a normal GPU
|
||
|
||
vLLM eagerly reserves `gpu_memory_utilization × total_memory` at startup for its KV cache pool, to
|
||
avoid reallocating mid-serve. On a normal discrete GPU this is a free lunch — it only competes with
|
||
other GPU processes. On the Spark, "GPU memory" **is** system RAM, so the default `0.85` reserves
|
||
~103GB out of 121GB total, starving the OS and desktop. We saw this manifest as active swap
|
||
thrashing (`vmstat` showing nonzero `si`/`so` continuously) at 118/121GB used with 1.3GB free.
|
||
|
||
In `configs/<your-config>.yaml`:
|
||
|
||
```yaml
|
||
vllm_gpu_memory_utilization: 0.5 # ~60GB reserved -- generous for a 4B model, leaves the OS room
|
||
```
|
||
|
||
Adjust upward if you know the box is otherwise idle; 0.5 was comfortably enough for gemma-3-4b-it
|
||
serving 50 concurrent generation threads.
|
||
|
||
## 6. Code patches needed for current library versions
|
||
|
||
The repo (as of the commit we tested against) predates some of the exact library versions that
|
||
`pip install` resolves to today. Three small patches were needed — none are DGX-Spark-specific,
|
||
they'd bite on any platform once these versions are current on PyPI:
|
||
|
||
**`utils/vllm_manager.py`** — vLLM 0.26.0 removed the `--disable-log-requests` flag (replaced by an
|
||
opt-in `--enable-log-requests`, off by default already). Delete the line:
|
||
```python
|
||
"--disable-log-requests", # Cleaner logs during generation
|
||
```
|
||
|
||
**`antislop-vllm/utils/refusal_detector.py`** — passes a `reference_compile=False` kwarg to
|
||
`AutoModelForSequenceClassification.from_pretrained(...)` that current transformers doesn't
|
||
recognize. It fails *silently* (caught, logged once, falls back to a no-op sentinel for the rest of
|
||
the run) — refusal filtering just quietly does nothing unless you check the log for
|
||
`[RefusalDetector ERROR]`. Remove the `reference_compile=False,` line.
|
||
|
||
**`core/ftpo_trainer.py`** — transformers 5.5.0's Gemma3 forward pass now *requires*
|
||
`token_type_ids` during training (used to build the causal mask; Gemma3 is natively multimodal and
|
||
needs to know which tokens are image vs. text). The FTPO trainer's `compute_loss` calls the model
|
||
in three places without it — all three need `token_type_ids=torch.zeros_like(ids)` added (text-only
|
||
training, so all-zero/all-text is correct):
|
||
- the main forward pass (`outputs = model(...)`)
|
||
- the reference-model forward pass, `self.ref_model is None` branch (inside `null_ref_context()`)
|
||
- the reference-model forward pass, `self.ref_model is not None` branch
|
||
|
||
## 7. FTPO fine-tuning is compute-bound, not throughput-bound
|
||
|
||
With `finetune_batch_size: 1` / `gradient_accumulation_steps: 16`, we measured ~750 optimizer steps
|
||
at ~160-185s/step — a ~34 hour run for the full 12,000-example dataset. Bumping
|
||
`finetune_batch_size` to 4 (with `gradient_accumulation_steps` dropped to 4 to keep the same
|
||
effective batch size) made **no meaningful difference** — still ~160-175s/step.
|
||
|
||
Takeaway: this workload is compute-bound (two full forward passes per micro-batch — the model plus
|
||
a reference-model pass for the MSE tether loss term — over sequences up to
|
||
`finetune_max_seq_length: 4000` tokens), not limited by batch-size/scheduling overhead. Increasing
|
||
batch size doesn't reduce total FLOPs for a fixed effective batch size, so it doesn't help here.
|
||
If you need a faster run, the actual levers are:
|
||
- lower `finetune_max_train_examples` (fewer total steps, less data coverage)
|
||
- lower `finetune_max_seq_length` (less compute per step, truncates longer training examples)
|
||
- accept the long runtime and let it run in the background
|
||
|
||
We left `finetune_batch_size` at the default (`1`) since increasing it only costs more memory for
|
||
no speed benefit on this hardware.
|
||
|
||
## Validated results
|
||
|
||
Ran the full pipeline against `unsloth/gemma-3-4b-it` (2 iterations, 1200 prompts each):
|
||
|
||
- Iteration 0 (baseline, no bans): completed in 28m24s, `repetition_per_100k_chars` = 160
|
||
- Iteration 1 (with ban lists from iteration 0's analysis): completed in 1h31m43s (slower — active
|
||
backtracking around bans), `repetition_per_100k_chars` = 56 — a real, measured reduction in slop
|
||
- FTPO training: confirmed working end-to-end (750 steps, 12,000 preference pairs) after the
|
||
patches in §6; not run to completion due to the ~34h runtime (§7)
|
||
|
||
## Quick-reference: full env setup
|
||
|
||
```bash
|
||
git clone --recurse-submodules https://github.com/sam-paech/auto-antislop.git
|
||
cd auto-antislop
|
||
|
||
conda create -n antislop python=3.11 -y
|
||
conda activate antislop
|
||
pip install "torch==2.11.0" "torchvision==0.26.0" # plain PyPI, matches system CUDA 13.0
|
||
grep -v '^flash-attn$' requirements.txt > /tmp/reqs_no_fa.txt
|
||
pip install -r /tmp/reqs_no_fa.txt
|
||
pip install "transformers<=5.5.0,>=4.51.3" # re-pin down for unsloth
|
||
|
||
pip install -U wheel ninja packaging cmake
|
||
FLASH_ATTN_CUDA_ARCHS=120 MAX_JOBS=6 pip install -U --no-build-isolation \
|
||
"git+https://github.com/Dao-AILab/flash-attention.git@v2.8.3.post1#egg=flash_attn"
|
||
|
||
conda create -n antislop-vllm-serve python=3.11 -y
|
||
conda activate antislop-vllm-serve
|
||
pip install vllm
|
||
|
||
mkdir -p ~/.local/bin
|
||
printf '#!/bin/bash\nexec %s/miniconda3/envs/antislop-vllm-serve/bin/vllm "$@"\n' "$HOME" \
|
||
> ~/.local/bin/vllm
|
||
chmod +x ~/.local/bin/vllm
|
||
# ensure ~/.local/bin is on PATH
|
||
|
||
sudo chown -R $USER:$USER ~/.triton # only if it's root-owned
|
||
|
||
# apply the 3 code patches from §6, then edit vllm_gpu_memory_utilization down to ~0.5
|
||
# in whichever configs/*.yaml you're running
|
||
|
||
conda activate antislop
|
||
python main.py -c configs/gemma-3-4b-it.yaml
|
||
```
|