auto-antislop/DGX_SPARK_SETUP.md

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 was slow because of fixed-length padding, not raw compute

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.

Initial hypothesis was that this was inherent — genuinely compute-bound (two full forward passes per micro-batch, over long sequences), not limited by batch-size/scheduling overhead. Measuring the actual training data disproved that. ftpo_trainer.py's collator pads every batch to a fixed finetune_max_seq_length (4000 tokens) regardless of content:

max_len  = self.args.max_length          # always 4000, never pad-to-longest-in-batch
prompt_ids = torch.full((batch_sz, max_len), pad_id, dtype=torch.long)

We tokenized all 12,000 training contexts with the real tokenizer to see how much of that 4000 was actually needed:

	tokens
mean	529.9
median	509
p90 / p99	953 / 1080
max across all 12,000 examples	1126

Not one example reaches even a third of the 4000-token padding target; the mean uses 13.2% of it. Every forward pass — both the main model and the reference-model pass — was processing ~4000 tokens of mostly padding, roughly 4-7x more than the actual content needs. This also explains why the batch-size bump did nothing: total padded-token compute is invariant to how the effective batch of 16 gets split into micro-batches, so reshuffling batch/accum never touched the real cost. This is a collator-design issue, not a hardware ceiling — it would waste the same proportion on any GPU.

Fix: lower finetune_max_seq_length to comfortably cover the real distribution, e.g. 1280 (covers p99 with headroom, nothing in the dataset gets truncated) instead of 4000. We left finetune_batch_size at the default (1) since increasing it has no effect either way here.

Measured, not just predicted: re-ran training with finetune_max_seq_length: 1280 for a 30-minute validation window (46 steps, timing fully steady by the end — no drift):

	max_seq_length=4000	max_seq_length=1280
steady-state step time	~160-185s/step	~40s/step
memory used during training	~82GB	~25GB
speedup	—	~4.3x
extrapolated full run (750 steps)	~34h	~8.3h

Better than the ~3x predicted from the token-count ratio alone — the memory savings from shorter sequences apparently helped beyond just the raw compute reduction. No errors across the validation run.

If you need it faster still, the other lever is finetune_max_train_examples (fewer total steps, less data coverage) — or just accept the ~8h runtime and let it run in the background.

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. At the original finetune_max_seq_length: 4000, steady-state was ~160-185s/step (~34h for the full run). After the fix in §7 (finetune_max_seq_length: 1280), measured ~40s/step over a 30-minute validation run — a confirmed ~4.3x speedup, ~8.3h extrapolated for the full 750 steps. Not run to full completion.

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