auto-antislop/DGX_SPARK_SETUP.md

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

```python
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`. That should cut
per-step compute roughly 3x, bringing the ~34h estimate down to somewhere around ~11-12h. We left
`finetune_batch_size` at the default (`1`) since increasing it has no effect either way here.

If you need it faster still, the other lever is `finetune_max_train_examples` (fewer total steps,
less data coverage) — or just accept the 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; 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
```