diff --git a/.gitignore b/.gitignore index a911299..638ec95 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ experiments unsloth_compiled_cache results bct.txt -.codex \ No newline at end of file +.codex +.claude \ No newline at end of file diff --git a/DGX_SPARK_SETUP.md b/DGX_SPARK_SETUP.md new file mode 100644 index 0000000..8c1a9ac --- /dev/null +++ b/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//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//.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/.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 +``` diff --git a/configs/gemma-3-4b-it.yaml b/configs/gemma-3-4b-it.yaml index 9dca09e..ec280f9 100644 --- a/configs/gemma-3-4b-it.yaml +++ b/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 diff --git a/core/ftpo_trainer.py b/core/ftpo_trainer.py index f5cfe4c..34ce97b 100644 --- a/core/ftpo_trainer.py +++ b/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, :] diff --git a/utils/vllm_manager.py b/utils/vllm_manager.py index 43bc2b7..8427ac2 100644 --- a/utils/vllm_manager.py +++ b/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: