Files
auto-antislop/DGX_SPARK_SETUP.md
Joey Grasty 3d05fa42bf Complete FTPO training run; document orchestration bug fix and thermal shutdown
Ran the full FTPO fine-tune to completion: 4h12m wall-clock, early-stopped at
step 380/750 by the configured chosen_win=0.85 threshold, train_loss 4.0->2.19.
Output is a merged 16-bit model plus LoRA adapter.

Also fixes an UnboundLocalError in core/orchestration.py when resuming a run
with generation disabled (--generation-step-enabled false), which is needed to
re-run just the finetune stage against an already-generated dataset.

Documents an unrelated platform incident: an uncapped multi-hour finetune
triggered a thermal shutdown mid-run with no checkpoint saved, requiring a
restart from step 0. Capping GPU clocks/power before the retry let it complete
without issue.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0141RdLKiMeXLMXsWyb4U3B5
2026-09-04 20:52:55 -05:00

14 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. Four small patches were needed — none are DGX-Spark-specific, they'd bite on any platform in the same circumstances:

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

core/orchestration.py — resuming a run with generation disabled (--generation-step-enabled false, used to re-run just the finetune stage against an already-generated dataset) crashes:

UnboundLocalError: cannot access local variable 'banned_ngrams_json_path' where it is not associated with a value

banned_ngrams_json_path / banned_slop_phrases_json_path are only ever assigned inside the if generation_enabled: block, but a later resume-logging block references them unconditionally. Guard that block with the same flag:

if start_iter_idx > 0 and generation_enabled:
    if banned_ngrams_json_path.exists(): ...

(There's a second, non-fatal instance of the same root cause — iter0_output_file_for_dpo referenced before assignment when loading stats from a prior completed run — it's caught and logged as a warning rather than raised, so it doesn't block anything, but it's the same bug shape.)

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.

8. Sustained full-GPU load can trigger a thermal shutdown

An unthrottled attempt at the full finetuning run caused the machine to shut down from heat partway through (no OOM/crash signature in the logs — the process and GPU state simply vanished when the box power-cycled). No checkpoint had been written yet, so the run restarted from step 0 — FTPO doesn't checkpoint mid-run by default, so a mid-run shutdown here means losing all progress so far.

Fix: cap GPU clocks/power before starting a sustained multi-hour job (finetuning, not the bursty generation step):

nvidia-smi -pl <lower_watts>          # if supported on this platform
# or check nvidia-smi -q -d CLOCK,POWER for current caps and headroom

With clocks capped, the retry ran the full ~4h12m to completion without incident — GPU temperature held in the 71-81°C range under sustained 90%+ utilization throughout.

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. Both numbers are from a run predating the refusal_detector.py fix in §6, so they reflect n-gram/phrase-ban slop reduction only, not refusal-aware filtering.
  • FTPO training: ran to completion. With the fixes from §6 and §7 (finetune_max_seq_length: 1280), trained on the 12,000-example preference-pair dataset (quota- sampled from 58,369 raw pairs) at finetune_batch_size: 1 / gradient_accumulation_steps: 16. Stopped early by the configured finetune_early_stopping_wins: 0.85 threshold — chosen_win (fraction of chosen completions preferred over rejected) crossed 0.85 at step 380 of a possible 750 (1 epoch). train_loss fell from ~4.0 to 2.19 over the run. Total wall-clock time (training + LoRA save + 16-bit merge): 4h12m32s, steady-state ~40s/step throughout — matching the §7 prediction. Output: LoRA adapter and a merged 16-bit model (8.1GB) under finetuned_model_ftpo_exp01/.

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