Files
auto-antislop/DGX_SPARK_SETUP.md
Joey Grasty 4edb2135be Add DGX Spark (aarch64/Blackwell) setup fixes and documentation
- 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>
2026-08-08 07:39:37 -05:00

10 KiB
Raw Blame History

Running auto-antislop on a DGX Spark

Notes from getting 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:

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:

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:

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:

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

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:

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:

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

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