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>

.gitignore

@@ -6,3 +6,4 @@ unsloth_compiled_cache
 results
 bct.txt
 .codex
+.claude

DGX_SPARK_SETUP.md

@@ -0,0 +1,221 @@
+# 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
+```

configs/gemma-3-4b-it.yaml

@@ -22,7 +22,7 @@ vllm_model_id: null # Model served by vLLM (if unset, will use model_id)
 vllm_port: 8000
 vllm_hf_token: null # Optional: Your Hugging Face token if model is gated
 vllm_cuda_visible_devices: "0"  # set to e.g. "0,1,2,3" for multiple gpus
-vllm_gpu_memory_utilization: 0.85 # leave some room for the refusal classifier if you are using it (about 3gb)
+vllm_gpu_memory_utilization: 0.5 # DGX Spark has unified CPU/GPU memory (121GB shared) -- 0.85 starved the OS and caused swap thrashing. 0.5 (~60GB) is generous for a 4B model + KV cache.
 vllm_max_model_len: 4500
 vllm_dtype: "bfloat16"
 # Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus

core/ftpo_trainer.py

@@ -219,6 +219,7 @@ class FTPOTrainer(DPOTrainer):
             ids,
             attention_mask=attn,
             position_ids=pos_full,
+            token_type_ids=torch.zeros_like(ids),
             use_cache=False,
             return_dict=True,
         )
@@ -259,11 +260,13 @@ class FTPOTrainer(DPOTrainer):
                     with self.null_ref_context():
                         ref_logits_last = model(
                             ids, attention_mask=attn, position_ids=pos_full,
+                            token_type_ids=torch.zeros_like(ids),
                             use_cache=False, return_dict=True,
                         ).logits[:, -1, :]
                 else:
                     ref_logits_last = self.ref_model(
                         ids, attention_mask=attn, position_ids=pos_full,
+                        token_type_ids=torch.zeros_like(ids),
                         use_cache=False, return_dict=True,
                     ).logits[:, -1, :]
 

utils/vllm_manager.py

@@ -77,7 +77,6 @@ def start_vllm_server(
         "--gpu-memory-utilization", str(gpu_memory_utilization),
         "--max-model-len", str(max_model_len),
         "--dtype", dtype,
-        "--disable-log-requests", # Cleaner logs during generation
         "--uvicorn-log-level", uvicorn_log_level.lower(),
     ]
     if hf_token: