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

@@ -4,4 +4,6 @@ experiments
 *.pyc
 unsloth_compiled_cache
 results
-bct.txt
+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
+```

auto_antislop_config.yaml

@@ -259,7 +259,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 # ── FTPO-specific hyper-parameters ─────────────────────────────────────────
@@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

configs/GLM-4-32B-0414.yaml

@@ -265,7 +265,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 
@@ -283,4 +284,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

configs/Mistral-Small-3.2-24B-Instruct-2506.yaml

@@ -265,7 +265,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.7
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 3 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 
@@ -283,4 +284,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

configs/default.yaml

@@ -259,7 +259,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 # ── FTPO-specific hyper-parameters ─────────────────────────────────────────
@@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

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

@@ -258,7 +258,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 
@@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

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

@@ -21,12 +21,12 @@ manage_vllm: true
 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.92 # leave some room for the refusal classifier if you are using it (about 3gb)
-vllm_max_model_len: 4500
+vllm_cuda_visible_devices: "1"  # set to e.g. "0,1,2,3" for multiple gpus
+vllm_gpu_memory_utilization: 0.97 # leave some room for the refusal classifier if you are using it (about 3gb)
+vllm_max_model_len: 1400
 vllm_dtype: "bfloat16"
 # Additional raw CLI arguments for vLLM server, e.g., ["--tensor-parallel-size", "4"] for multiple gpus
-vllm_extra_args: [] # each param as a separate string, e.g. ["--quantization", "bitsandbytes"]
+vllm_extra_args: [] #["--quantization", "bitsandbytes"]
 vllm_env:                  # env vars for the vLLM process
   # VLLM_USE_V1: "1"  # may be needed for amd gpus
 
@@ -45,11 +45,11 @@ generation_api_key: "xxx" # API key for the vLLM server
 
 # --- Core Generation Settings ---
 generation_max_new_tokens: 1000
-generation_threads: 200 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
-generation_max_prompts: 2000 # Number of samples to generate from the prompts in the dataset
+generation_threads: 50 # Number of parallel threads for API queries in antislop-vllm. Note: vllm can become very inefficient if you go over some concurrency threshold (depending on vram)
+generation_max_prompts: 4000 # Number of samples to generate from the prompts in the dataset
 
 # --- Dataset & Chat Template ---
-generation_hf_dataset_name: 'Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT'
+generation_hf_dataset_name: 'sam-paech/essays-creative-writing-prompts'
 generation_hf_dataset_split: 'train'
 # A huggingface model id or local dir containing the tokeniser you want to use to apply chat templates.
 # This is important if you are generating a ftpo dataset for later training.
@@ -91,7 +91,7 @@ generation_ngram_language: "english"
 
 # --- Refusal Detection ---
 # Detects refusals & doesn't include them in the training dataset. Uses about 3GB extra VRAM.
-generation_refusal_detection: true
+generation_refusal_detection: false
 
 ################################################################################
 # N-GRAM ANALYSIS & BANNING (within auto-antislop)
@@ -105,19 +105,19 @@ top_k_trigrams: 5000
 
 # --- N-gram Banning Quotas (per iteration) ---
 # Bigrams
-dict_bigrams_initial: 400     # How many of the top over-represented dictionary bigrams to
+dict_bigrams_initial: 300     # How many of the top over-represented dictionary bigrams to
                               # ban in the first antislop iteration.
                               # "Dictionary" means the bigrams were also found in the human
                               # writing corpus.
-dict_bigrams_subsequent: 70   # How many to ban in each subsequent iteration
-nodict_bigrams_initial: 800   # "Nodict" here means the n-grams were not found at all in the
+dict_bigrams_subsequent: 0   # How many to ban in each subsequent iteration
+nodict_bigrams_initial: 200   # "Nodict" here means the n-grams were not found at all in the
                               # human corpus.
-nodict_bigrams_subsequent: 100
+nodict_bigrams_subsequent: 0
 # Trigrams
 dict_trigrams_initial: 300
-dict_trigrams_subsequent: 50
-nodict_trigrams_initial: 800
-nodict_trigrams_subsequent: 100
+dict_trigrams_subsequent: 0
+nodict_trigrams_initial: 200
+nodict_trigrams_subsequent: 0
 
 # --- User-Defined N-gram Bans ---
 # User-supplied extra n-grams to always ban (processed by auto-antislop)
@@ -132,14 +132,14 @@ compute_overrep_words: true
 top_k_words_for_overrep_analysis: 200000
 
 # --- Quotas for Adding Over-represented Words to Slop Phrase Ban List ---
-dict_overrep_initial: 800       # How many of the top over-represented dictionary words to
+dict_overrep_initial: 2000       # How many of the top over-represented dictionary words to
                                 # ban in the first antislop iteration.
                                 # "Dictionary" means the words were also found in the human
                                 # writing corpus.
-dict_overrep_subsequent: 200    # How many to ban in each subsequent iteration
-nodict_overrep_initial: 80      # "Nodict" here means the n-grams were not found at all in the
+dict_overrep_subsequent: 0    # How many to ban in each subsequent iteration
+nodict_overrep_initial: 120      # "Nodict" here means the n-grams were not found at all in the
                                 # human corpus.
-nodict_overrep_subsequent: 20
+nodict_overrep_subsequent: 0
 
 ################################################################################
 # SLOP PHRASE BANNING
@@ -148,8 +148,8 @@ nodict_overrep_subsequent: 20
 # Slop phrases are over-represented whole phrases extracted from the generated texts.
 enable_slop_phrase_ban: true
 min_phrase_freq_to_keep: 2 # Min frequency for a new phrase from slop-forensics to be considered
-top_n_initial_slop_ban: 600 # New slop phrases from slop-forensics to ban in iter 0
-top_n_subsequent_slop_ban: 100 # New slop phrases from slop-forensics to ban in later iters
+top_n_initial_slop_ban: 0 # New slop phrases from slop-forensics to ban in iter 0
+top_n_subsequent_slop_ban: 0 # New slop phrases from slop-forensics to ban in later iters
 
 # --- User-Defined Slop Phrase Bans ---
 # User supplied list of strings to always ban
@@ -187,15 +187,15 @@ whitelist_strings: [
 extra_regex_patterns: [
   # These ones ban "it's not x, it's y" type patterns:
     
-  #"\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
+  "\\b(?:\\w+n(?:['’]t)|not\\s+(?:just|only|merely|because))\\s+(?:(?![.;:?!…]).){1,100}?[.;:?!…]\\s*(?:it|they|you)(?:['’](?:s|re|m))?\\b(?!\\s+(?:was|were|is|are|wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t)\\b)(?:\\s*[*…]?\\s*)?(?!when\\b|then\\b|but\\b|and\\b|yet\\b)(?!right\\b)(?!normal\\b)(?!true\\b)(?!sure\\b)(?!only\\b)(?!still\\b)(?!rarely\\b)(?!already\\b)(?!wrong\\b)(?!want\\b)(?!just\\b)(?!couldn\\b)(?!could\\b)(?!saw\\b)(?!started\\b)(?!remember\\b)(?!struggled\\b)(?!watched\\b)(?!goal\\b)(?!took\\b)(?!kept\\b)(?!reminded\\b)(?!time\\b)(?!have\\b)(?!acted\\b)(?!smiled\\b)(?!think\\b)(?!give\\b)(?!grab\\b)(?!gave\\b)(?!turn\\b)(?!justify\\b)(?!\\w+ly\\b)(?=[a-z]{4,}\\b)[a-z]+\\w*",
 
-  #"\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
+  "\\b(?:\\w+n(?:['’]t)|not)\\s+(?:just|only|merely)?\\s*(?:(?![-–—]|[.?!…]).){1,80}?[-–—]{1,2}\\s*\\w+(?:['’]\\w+)?\\s+",
 
-  #"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
+  "\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?!\\b(?:minute|minutes|hour|hours|day|days|year|years|second|seconds)\\b)(?!with\\b)(?!even\\b)(?:(?![.;:?!…]).){2,120}?[.;:?!…]\\s*(?:it|they|you|that)(?:\\s+(?:was|were|is|are)\\b(?:\\s+[*_~]?\\w+[*_~]?)?|(?:['’](?:s|re|m))\\b(?:\\s+[*_~]?\\w+[*_~]?)?)",
 
-  #"\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
+  "\\bno\\s+longer\\s+(?:just|only|merely)?\\s+[^.;:?!…]{1,120}[.;:?!…]\\s*(?:it|they|you)\\s+(?:is|are|was|were)\\b(?:\\s+[*_~]?\\w+[*_~]?)?",
 
-  #"\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
+  "\\b(?:wasn['’]t|weren['’]t|isn['’]t|aren['’]t|ain['’]t|not)\\s+(?:just|only|merely)?\\s*(?:(?!\\bbut\\b|[.?!…]).){1,80}?[,;:\\-–—]\\s*but\\s+(?!I\\b)(?:also\\s+)?"
 
 ]
 
@@ -203,10 +203,10 @@ extra_regex_patterns: [
 # FINETUNING
 ################################################################################
 finetune_enabled: true
-
+#finetune_attention_implementation: eager
 # --- General Finetuning Setup ---
-finetune_use_unsloth: false
-finetune_mode: "ftpo" # ftpo / dpo / dpo_final_token
+finetune_use_unsloth: true
+finetune_mode: "ftpo" # dpo / ftpo (final token preference optimisation)
 finetune_ftpo_dataset: ""   # you can specify an existing ftpo dataset, or leave unset to let the
                             # pipeline use the one produced in the generation step
 finetune_base_model_id: null # Base model for DPO (if unset, uses model_id)
@@ -214,33 +214,33 @@ finetune_max_seq_length: 2500 # this may truncate some outputs
 finetune_load_in_4bit: true # qlora
 
 # --- Early Stopping ---
-finetune_early_stopping_wins: 0.85  # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
+finetune_early_stopping_wins: 0.88  # Early stopping threshold for fraction of *chosen* completions that are selected over *rejected*.
                                     # More than 0.85 may be overtrained. Set to > 1.0 to disable early stopping.
 finetune_early_stopping_loss: null  # Loss threshold for early stopping. Set to null to disable.
 
 # --- LoRA Configuration ---
-finetune_lora_r: 128 # the ftpo trainer works best with a high lora rank
-finetune_lora_alpha: 128
+finetune_lora_r: 512 # the ftpo trainer works best with a high lora rank
+finetune_lora_alpha: 256
 finetune_lora_dropout: 0.05
 finetune_weight_decay: 0.01
 finetune_target_modules: ["up_proj", "down_proj", "lm_head"] #["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
 
 # --- Layer Freezing ---
 finetune_freeze_early_layers: true
-finetune_n_layers_unfrozen: 5
+finetune_n_layers_unfrozen: 3
 
 # --- Training Process ---
 finetune_gradient_checkpointing: "unsloth"
 finetune_chat_template: "" # e.g. "gemma-3" -- get the chat template from unsloth's helper if required, otherwise leave the string blank to use the tokeniser's chat template
-finetune_batch_size: 1
-finetune_gradient_accumulation_steps: 16
+finetune_batch_size: 2
+finetune_gradient_accumulation_steps: 6
 finetune_warmup_ratio: 0.1
 finetune_num_epochs: 1
 
 # --- Learning Rate ---
 finetune_learning_rate: 0.000001
 finetune_auto_learning_rate: true  # true: automatically determine learning rate based on dataset size, effective batch size & lora rank
-finetune_auto_learning_rate_adjustment_scaling: 0.08 # scale the auto-lr by this factor
+finetune_auto_learning_rate_adjustment_scaling: 0.065 # scale the auto-lr by this factor
 
 # --- DPO/FTPO Specific ---
 finetune_beta: 0.1 # DPO beta
@@ -257,8 +257,9 @@ finetune_shuffle_seed: 666
 # --- FTPO Sample Regularization ---
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
-ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+ftpo_sample_rejected_regularisation_strength: 0.6
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 
@@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

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
@@ -258,7 +258,8 @@ finetune_shuffle_seed: 666
 # 0 = off; 0.9 strongly downsamples overrepresented rule violations
 # (this is useful because the raw generated dataset is typically very skewed)
 ftpo_sample_rejected_regularisation_strength: 0.8
-ftpo_sample_chosen_regularisation_strength: 0.2
+# 0 = off; positive values trim globally overrepresented chosen-token slots
+ftpo_sample_chosen_regularisation_strength: 0.0
 ftpo_sample_min_chosen_tokens: 4 # filter out ftpo samples that have fewer than this number in the chosen tokens list
 
 
@@ -276,4 +277,4 @@ ftpo_tau_mse_target: 0.5       # Grace bandwidth (logits) before the above MSE l
 # MSE loss term 2: stronger mse term applied to remaining (non-target) vocab
 ftpo_lambda_mse: 0.4
 
-ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"
+ftpo_clip_epsilon_logits: 2     # For a chosen token: "after winning vs rejected token by this margin, preference loss turns off"

core/analysis.py

@@ -156,10 +156,23 @@ def select_overrep_words_for_ban(dict_words: list[str],
     for w in dict_words:
         if len(selected) >= dict_q: break
         if w.lower() not in whitelist: selected.append(w)
+
+    n_dict = len(selected)
     for w in nodict_words:
-        if len(selected) >= dict_q + nodict_q: break
+        if len(selected) - n_dict >= nodict_q: break
         if w.lower() not in whitelist: selected.append(w)
-    logger.info(f"Selected {len(selected)} over-rep words for ban ({dict_q}/{nodict_q} quotas).")
+
+    n_nodict = len(selected) - n_dict
+    logger.info(
+        "Selected %d dict + %d non-dict over-rep words for ban "
+        "(quotas %d/%d; pools %d/%d).",
+        n_dict,
+        n_nodict,
+        dict_q,
+        nodict_q,
+        len(dict_words),
+        len(nodict_words),
+    )
     return selected
 
 
@@ -223,6 +236,9 @@ def update_banned_slop_phrases(
         except Exception:
             pass
 
+    # Track original size for logging
+    original_size = len(existing)
+
     # keep requested quota only
     cand_phrases = cand_phrases[:how_many_new]
     if over_represented_words:
@@ -230,10 +246,15 @@ def update_banned_slop_phrases(
             if w not in whitelist:
                 existing.add(w)
 
+    # Re-add extra_slop_phrases_to_ban from config (matches behavior of update_banned_ngrams_list)
+    for phrase in config.get('extra_slop_phrases_to_ban', []):
+        if phrase and not is_whitelisted(phrase):
+            existing.add(phrase)
+
     merged = sorted((existing | set(cand_phrases)) - whitelist)
     json_path.write_text(json.dumps([[p, 1] for p in merged], indent=2, ensure_ascii=False), "utf-8")
     logger.info(f"🚫 Slop-phrase ban list now {len(merged)} entries "
-                f"(+{len(merged)-len(existing)} this iter)")
+                f"(+{len(merged)-original_size} this iter)")
 
 
 # --- N-Gram Analysis ---
@@ -467,4 +488,4 @@ def calculate_repetition_score(gen_texts: list, total_chars: int, iteration_dfs:
         for tg in current_trigrams:
             if tg in target_ngrams: total_repetition_instances += 1
             
-    return norm_per_freq_denom(total_repetition_instances, float(total_chars), freq_norm_denom)
+    return norm_per_freq_denom(total_repetition_instances, float(total_chars), freq_norm_denom)

core/finetuning.py

@@ -29,7 +29,12 @@ import math
 import json
 from datetime import datetime, timezone
 from utils.dataset_helpers import load_ftpo_multi_dataset
-from utils.model_helpers import fix_gemma3_checkpoint, detie_lm_head, prepare_gemma3_for_save
+from utils.model_helpers import (
+    fix_gemma3_checkpoint,
+    detie_lm_head,
+    prepare_gemma3_for_save,
+    unwrap_clippable_linears,
+)
 # Import the new dataloader function
 from utils.trainer_dataloaders import load_and_prepare_dataset
 logger = logging.getLogger(__name__)
@@ -999,6 +1004,9 @@ def run_dpo_finetune(config: dict, experiment_run_dir: Path):
                 device_map  = {"": "cpu"},
                 trust_remote_code = True,
             )
+        unwrapped = unwrap_clippable_linears(base_fp16)
+        if unwrapped:
+            logger.info("Unwrapped %d clippable linear modules before PEFT merge.", unwrapped)
         model_fp16 = PeftModel.from_pretrained(
             base_fp16,
             lora_dir,                           # plug in the saved adapter
@@ -1051,4 +1059,4 @@ def run_dpo_finetune(config: dict, experiment_run_dir: Path):
         except Exception as e:
             logger.error(f"Error saving model: {e}", exc_info=True)
 
-    logger.info("Finetuning process completed.")
+    logger.info("Finetuning process completed.")

core/ftpo_trainer.py

@@ -209,7 +209,7 @@ class FTPOTrainer(DPOTrainer):
         attn  = inputs["attention_mask"].to(device)           # [B,L]
         B, L  = ids.shape
         
-        seq_len  = attn.sum(1)                         
+        seq_len  = attn.sum(1)
         pad_off  = (L - seq_len).unsqueeze(1)
         arange_L = torch.arange(L, device=ids.device).unsqueeze(0)
         pos_full = (arange_L - pad_off).clamp(min=0)
@@ -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,
         )
@@ -232,27 +233,24 @@ class FTPOTrainer(DPOTrainer):
         logp_bad = logp_all.gather(-1, rejected.unsqueeze(-1)).squeeze(-1)  
 
         batch_rows = torch.arange(B, device=logp_all.device).unsqueeze(1)
-        gathered      = logits_last[batch_rows, ch_ids]   
-        logit_bad     = logits_last.gather(-1, rejected.unsqueeze(-1))
-        margin        = gathered - logit_bad
-        weights = torch.clamp((clip_epsilon_logits - margin) / clip_epsilon_logits, 0.0, 1.0) * ch_mask
+        delta_tok = logits_last[batch_rows, ch_ids] - logits_last.gather(
+            -1, rejected.unsqueeze(-1)
+        )
+        weights = (
+            torch.clamp(
+                (clip_epsilon_logits - delta_tok) / clip_epsilon_logits,
+                0.0,
+                1.0,
+            )
+            * ch_mask
+        )
 
-        zero_row = weights.sum(dim=-1, keepdim=True) < 1e-12
-        weights  = torch.where(zero_row, ch_mask.float(), weights)
+        tau = 1.0
+        gap = clip_epsilon_logits - delta_tok
+        per_tok_loss = F.softplus(gap / tau)
 
-        weights_sum = weights.sum(dim=-1, keepdim=True)               
-        batch_rows  = torch.arange(B, device=ids.device).unsqueeze(1)
-
-        l_chosen    = logits_last[batch_rows, ch_ids]             
-        l_bad       = logits_last.gather(-1, rejected.unsqueeze(-1))   
-        delta_tok   = l_chosen - l_bad                            
-
-        margin  = clip_epsilon_logits
-        tau     = 1.0        
-        gap     = margin - delta_tok          
-        per_tok_loss = F.softplus(gap / tau)  
-
-        pref_loss = (per_tok_loss * weights).sum() / weights_sum.sum()   
+        chosen_counts = ch_mask.sum(dim=-1).clamp(min=1)
+        pref_loss = ((per_tok_loss * weights).sum(dim=-1) / chosen_counts).mean()
 
         extra_metrics = {}
 
@@ -262,21 +260,23 @@ 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, :]
 
-            freeze_mask = torch.ones_like(logits_last, dtype=torch.bool)
+            tether_mask = torch.ones_like(logits_last, dtype=torch.bool)
             rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
-            freeze_mask[rows[ch_mask], ch_ids[ch_mask]] = False
-            freeze_mask.scatter_(1, rejected.unsqueeze(-1), False)
+            tether_mask[rows[ch_mask], ch_ids[ch_mask]] = False
+            tether_mask.scatter_(1, rejected.unsqueeze(-1), False)
 
             diff        = logits_last - ref_logits_last
-            mse_elem_raw = (freeze_mask * diff.pow(2)).sum() / freeze_mask.sum()
+            mse_elem_raw = (tether_mask * diff.pow(2)).sum() / tether_mask.sum()
 
             tgt_mask = torch.zeros_like(logits_last, dtype=torch.bool)
             rows = torch.arange(B, device=ch_ids.device).unsqueeze(1).expand_as(ch_ids)
@@ -312,9 +312,23 @@ class FTPOTrainer(DPOTrainer):
         frac_win  = wins_tok.float().sum(-1) / ch_mask.sum(-1).clamp(min=1e-8)  
         chosen_win = frac_win.mean().detach()                 
 
+        active_delta = delta_tok[ch_mask]
+        active_weights = weights[ch_mask]
+        margin_win = (
+            ((delta_tok >= clip_epsilon_logits) & ch_mask).float().sum()
+            / ch_mask.float().sum().clamp(min=1e-8)
+        ).detach()
+        mean_delta = active_delta.mean().detach()
+        median_delta = active_delta.median().detach()
+        active_weight = active_weights.mean().detach()
+
         metrics = {
             "pref_loss":  pref_loss.detach(),
             "chosen_win": chosen_win,
+            "margin_win": margin_win,
+            "mean_delta": mean_delta,
+            "median_delta": median_delta,
+            "active_weight": active_weight,
             **extra_metrics,
         }
         self.store_metrics(metrics, train_eval="train")
@@ -328,4 +342,4 @@ class FTPOTrainer(DPOTrainer):
 
     # ----------------------------------------------------------
     def _prepare_dataset(self, dataset, *args, **_):
-        return dataset
+        return dataset

core/orchestration.py

@@ -318,7 +318,20 @@ def orchestrate_pipeline(config: Dict[str, Any], experiment_dir: Path, resume_mo
             if not _p.exists():
                 _p.write_text("[]", encoding="utf-8")        # write an empty JSON array
 
-        
+        # --- Merge user-defined bans from config (always, not just on initial run) ---
+        # This ensures extra_ngrams_to_ban and extra_slop_phrases_to_ban are included
+        # before any iteration starts, regardless of resume mode.
+        # This is idempotent since merge_custom_bans_into_file uses set union.
+        if config['enable_ngram_ban'] and config.get('extra_ngrams_to_ban'):
+            merge_custom_bans_into_file(banned_ngrams_json_path,
+                                        config['extra_ngrams_to_ban'])
+            logger.info(f"📝 Merged {len(config['extra_ngrams_to_ban'])} user-defined n-grams into {banned_ngrams_json_path.name}")
+        if config['enable_slop_phrase_ban'] and config.get('extra_slop_phrases_to_ban'):
+            merge_custom_bans_into_file(banned_slop_phrases_json_path,
+                                        config['extra_slop_phrases_to_ban'])
+            logger.info(f"📝 Merged {len(config['extra_slop_phrases_to_ban'])} user-defined slop phrases into {banned_slop_phrases_json_path.name}")
+
+
         # --- Regex Blocklist (user-supplied, written once if provided, used from iter 1+) ---
         # This file is created before the loop, but only passed to generation from iter 1.
         user_regex_blocklist_file: Optional[Path] = None # Renamed for clarity
@@ -432,16 +445,6 @@ def orchestrate_pipeline(config: Dict[str, Any], experiment_dir: Path, resume_mo
                     if user_regex_blocklist_file and user_regex_blocklist_file.exists(): # User-defined regex
                         regex_file_for_generation = user_regex_blocklist_file
 
-                    # If we are resuming and this is the first iteration after the resume,
-                    # force-merge any new YAML bans into the existing files *before* generation.
-                    if resume_mode and iter_idx == start_iter_idx:
-                        if config['enable_ngram_ban'] and config.get('extra_ngrams_to_ban'):
-                            merge_custom_bans_into_file(banned_ngrams_json_path,
-                                                        config['extra_ngrams_to_ban'])
-                        if config['enable_slop_phrase_ban'] and config.get('extra_slop_phrases_to_ban'):
-                            merge_custom_bans_into_file(banned_slop_phrases_json_path,
-                                                        config['extra_slop_phrases_to_ban'])
-                            
                     _copy_if_exists(ngram_file_for_generation,
                                     iter_analysis_dir / "banned_ngrams_used.json")
                     _copy_if_exists(slop_file_for_generation,

tests/__init__.py

@@ -0,0 +1 @@
+

tests/test_analysis.py

@@ -0,0 +1,42 @@
+import sys
+import unittest
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "slop-forensics"))
+
+from core.analysis import select_overrep_words_for_ban
+
+
+class SelectOverrepWordsForBanTests(unittest.TestCase):
+    def test_unused_dictionary_quota_does_not_spill_into_non_dictionary_pool(self):
+        config = {
+            "dict_overrep_initial": 10,
+            "nodict_overrep_initial": 2,
+            "dict_overrep_subsequent": 1,
+            "nodict_overrep_subsequent": 1,
+        }
+
+        with self.assertLogs("core.analysis", level="INFO") as logs:
+            selected = select_overrep_words_for_ban(
+                ["dict-a", "dict-b", "dict-c"],
+                ["nodict-a", "nodict-b", "nodict-c", "nodict-d", "nodict-e"],
+                True,
+                config,
+                whitelist=set(),
+            )
+
+        self.assertEqual(
+            selected,
+            ["dict-a", "dict-b", "dict-c", "nodict-a", "nodict-b"],
+        )
+        self.assertIn(
+            "Selected 3 dict + 2 non-dict over-rep words for ban "
+            "(quotas 10/2; pools 3/5).",
+            logs.output[0],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/dataset_helpers.py

@@ -16,6 +16,82 @@ logger = logging.getLogger(__name__)
 _WATCH = [" nodded", " leaned"]
 
 
+def _chosen_target_quotas(
+    chosen_counts: Counter[str],
+    strength: float,
+) -> dict[str, int]:
+    """Return per-token occurrence caps for chosen-token regularisation."""
+    if not chosen_counts or strength <= 0:
+        return dict(chosen_counts)
+
+    # Prevent the largest outliers from dominating the distribution before
+    # applying the smoother median-threshold regularisation.
+    if len(chosen_counts) >= 10:
+        cap_value = sorted(chosen_counts.values(), reverse=True)[9]
+        capped = {
+            token: min(count, cap_value)
+            for token, count in chosen_counts.items()
+        }
+    else:
+        capped = dict(chosen_counts)
+
+    median = float(np.median(list(capped.values())))
+    return {
+        token: int(round(
+            count
+            if count <= median
+            else count * (median / count) ** strength
+        ))
+        for token, count in capped.items()
+    }
+
+
+def _trim_chosen_to_quotas(
+    rows: list[dict],
+    quotas: dict[str, int],
+    rng: np.random.Generator,
+) -> list[dict]:
+    """Uniformly retain chosen-token occurrences up to their global quotas."""
+    counts = Counter(
+        token
+        for row in rows
+        for token in (row.get("multi_chosen_decoded") or [])
+    )
+    if all(quotas.get(token, count) >= count for token, count in counts.items()):
+        return [dict(row) for row in rows]
+
+    # Randomise both row and slot traversal so quota allocation does not
+    # systematically favour early examples or a token's probability rank.
+    seen: Counter[str] = Counter()
+    trimmed: list[dict | None] = [None] * len(rows)
+    for row_idx_raw in rng.permutation(len(rows)):
+        row_idx = int(row_idx_raw)
+        row = rows[row_idx]
+        decoded = row.get("multi_chosen_decoded") or []
+        keep: set[int] = set()
+        for slot_idx_raw in rng.permutation(len(decoded)):
+            slot_idx = int(slot_idx_raw)
+            token = decoded[slot_idx]
+            if seen[token] < max(0, quotas.get(token, counts[token])):
+                keep.add(slot_idx)
+                seen[token] += 1
+
+        new_row = dict(row)
+        new_row["multi_chosen_decoded"] = [
+            token for slot_idx, token in enumerate(decoded) if slot_idx in keep
+        ]
+
+        raw = row.get("multi_chosen_raw")
+        if isinstance(raw, list) and len(raw) == len(decoded):
+            new_row["multi_chosen_raw"] = [
+                token for slot_idx, token in enumerate(raw) if slot_idx in keep
+            ]
+
+        trimmed[row_idx] = new_row
+
+    return [row for row in trimmed if row is not None]
+
+
 def load_ftpo_multi_dataset(
     path: Path,
     tokenizer,
@@ -108,24 +184,10 @@ def load_ftpo_multi_dataset(
 
     _log_top(chosen_cts_orig, "ORIGINAL CHOSEN TOKENS")
 
-    # Trim the peak: cap top tokens to match the 10th highest count
-    if len(chosen_cts_orig) >= 10:
-        top_counts = sorted(chosen_cts_orig.values(), reverse=True)
-        cap_value = top_counts[9]  # 10th highest count
-        chosen_cts_capped = Counter()
-        for tok, cnt in chosen_cts_orig.items():
-            chosen_cts_capped[tok] = min(cnt, cap_value)
-    else:
-        chosen_cts_capped = chosen_cts_orig.copy()
-
-    # Now calculate regularization on the capped distribution
-    med_chosen = float(np.median(list(chosen_cts_capped.values())))
-    w_chosen = {tok: 1.0 if c <= med_chosen
-            else (med_chosen / c) ** chosen_reg_strength
-            for tok, c in chosen_cts_capped.items()}
-
-    tgt_chosen = {tok: int(round(c * w_chosen.get(tok, 1.0)))
-                for tok, c in chosen_cts_capped.items()}
+    tgt_chosen = _chosen_target_quotas(
+        chosen_cts_orig,
+        chosen_reg_strength,
+    )
 
     # Log the target quotas
     quota_items = sorted(tgt_chosen.items(), key=lambda x: x[1], reverse=True)[:20]
@@ -137,7 +199,13 @@ def load_ftpo_multi_dataset(
         _WATCH[1], tgt_chosen.get(_WATCH[1], 0), chosen_cts_orig.get(_WATCH[1], 0),
     )
 
-    _log_top(Counter(r["rejected_decoded"] for r in rows), "POST-CHOSEN")
+    rows = _trim_chosen_to_quotas(rows, tgt_chosen, rng)
+    chosen_cts_trimmed = Counter(
+        token
+        for row in rows
+        for token in (row["multi_chosen_decoded"] or [])
+    )
+    _log_top(chosen_cts_trimmed, "POST-CHOSEN")
 
     # ────────────────────────────────────────────────────────────────
     # 3️⃣  Apply min_chosen_tokens row filter
@@ -206,6 +274,13 @@ def load_ftpo_multi_dataset(
     
     rows = selected
 
+    final_chosen_counts = Counter(
+        token
+        for row in rows
+        for token in (row["multi_chosen_decoded"] or [])
+    )
+    _log_top(final_chosen_counts, "FINAL CHOSEN TOKENS")
+
     # ── Dump the final row subset exactly as it was read (no tokenisation) ──
     if experiment_run_dir is not None:
         ts = datetime.now(timezone.utc).astimezone()\

utils/model_helpers.py

@@ -7,6 +7,30 @@ from safetensors.torch import safe_open, save_file
 
 log = logging.getLogger(__name__)
 
+
+def unwrap_clippable_linears(model) -> int:
+    """
+    Replace remote-code clippable linear wrappers with their underlying
+    torch.nn.Linear modules so vanilla PEFT can inject adapters.
+
+    Some Gemma loaders expose projection modules as Gemma*ClippableLinear
+    wrappers with the real Linear stored on `.linear`. PEFT's LoRA injection
+    only accepts the inner Linear type.
+    """
+    count = 0
+
+    for child_name, child in list(model.named_children()):
+        count += unwrap_clippable_linears(child)
+
+        inner = getattr(child, "linear", None)
+        if inner is None:
+            continue
+        if child.__class__.__name__.endswith("ClippableLinear"):
+            setattr(model, child_name, inner)
+            count += 1
+
+    return count
+
 def fix_gemma3_checkpoint(ckpt_dir: str | Path) -> None:
     """
     If `ckpt_dir` is a Gemma-3 checkpoint whose tensor keys look like

utils/test_dataset_helpers.py

@@ -0,0 +1,154 @@
+import json
+import tempfile
+import unittest
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+from utils.dataset_helpers import (
+    _chosen_target_quotas,
+    _trim_chosen_to_quotas,
+    load_ftpo_multi_dataset,
+)
+
+
+def _counts(rows):
+    return Counter(
+        token
+        for row in rows
+        for token in row["multi_chosen_decoded"]
+    )
+
+
+class _Tokenized:
+    def __init__(self, input_ids):
+        self.input_ids = input_ids
+
+
+class _FakeTokenizer:
+    truncation_side = "right"
+
+    def __call__(self, text, **_kwargs):
+        def encode(value):
+            if value.startswith("context-"):
+                return [1, 2]
+            return [100 + sum(value.encode("utf-8"))]
+
+        if isinstance(text, list):
+            return _Tokenized([encode(value) for value in text])
+        return _Tokenized(encode(text))
+
+
+class ChosenRegularisationTests(unittest.TestCase):
+    def test_zero_strength_disables_all_trimming(self):
+        counts = Counter({f"token-{i}": 20 - i for i in range(12)})
+
+        self.assertEqual(_chosen_target_quotas(counts, 0), dict(counts))
+
+    def test_positive_strength_caps_and_regularises_outliers(self):
+        counts = Counter({f"token-{i}": 100 - 5 * i for i in range(12)})
+
+        quotas = _chosen_target_quotas(counts, 0.2)
+
+        tenth_highest = sorted(counts.values(), reverse=True)[9]
+        self.assertLess(quotas["token-0"], counts["token-0"])
+        self.assertLessEqual(quotas["token-0"], tenth_highest)
+        self.assertEqual(quotas["token-11"], counts["token-11"])
+
+    def test_trimming_enforces_quotas_and_keeps_raw_fields_aligned(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", " rare-a", " common"],
+                "multi_chosen_raw": ["raw-common-1", "raw-rare-a", "raw-common-2"],
+            },
+            {
+                "multi_chosen_decoded": [" common", " rare-b"],
+                "multi_chosen_raw": ["raw-common-3", "raw-rare-b"],
+            },
+        ]
+        quotas = {" common": 2, " rare-a": 1, " rare-b": 1}
+
+        trimmed = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(3407)
+        )
+
+        self.assertEqual(_counts(trimmed), Counter(quotas))
+        for row in trimmed:
+            self.assertEqual(
+                len(row["multi_chosen_decoded"]),
+                len(row["multi_chosen_raw"]),
+            )
+        # The helper does not mutate the source rows.
+        self.assertEqual(rows[0]["multi_chosen_decoded"].count(" common"), 2)
+
+    def test_trimming_is_reproducible(self):
+        rows = [
+            {
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": [" common", f" unique-{i}"],
+            }
+            for i in range(20)
+        ]
+        quotas = {" common": 5, **{f" unique-{i}": 1 for i in range(20)}}
+
+        first = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+        second = _trim_chosen_to_quotas(
+            rows, quotas, np.random.default_rng(123)
+        )
+
+        self.assertEqual(first, second)
+        self.assertEqual(_counts(first)[" common"], 5)
+
+    def test_minimum_filter_applies_after_trimming(self):
+        rows = [
+            {"multi_chosen_decoded": [" common", " keep"]},
+            {"multi_chosen_decoded": [" common", " other"]},
+        ]
+
+        trimmed = _trim_chosen_to_quotas(
+            rows,
+            {" common": 1, " keep": 1, " other": 1},
+            np.random.default_rng(7),
+        )
+        surviving = [
+            row for row in trimmed if len(row["multi_chosen_decoded"]) >= 2
+        ]
+
+        self.assertEqual(len(surviving), 1)
+
+    def test_loader_filters_rows_after_applying_chosen_quotas(self):
+        rows = [
+            {
+                "context_with_chat_template": f"context-{i}",
+                "rejected_decoded": " reject",
+                "multi_chosen_decoded": [" common", f" unique-{i}"],
+                "multi_chosen_raw": ["raw-common", f"raw-unique-{i}"],
+            }
+            for i in range(4)
+        ]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir) / "ftpo.jsonl"
+            with path.open("w", encoding="utf-8") as handle:
+                for row in rows:
+                    handle.write(json.dumps(row) + "\n")
+
+            dataset = load_ftpo_multi_dataset(
+                path,
+                _FakeTokenizer(),
+                chosen_reg_strength=1.0,
+                min_chosen_tokens=2,
+                num_proc=1,
+            )
+
+        # " common" is trimmed from four occurrences to one, so only its
+        # containing row still meets the two-chosen-token minimum.
+        self.assertEqual(len(dataset), 1)
+        self.assertEqual(len(dataset[0]["chosen_ids"]), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/trainer_dataloaders.py

@@ -110,7 +110,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             # balance *rejected* tokens
             rejected_reg_strength = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
             # balance *chosen* tokens
-            chosen_reg_strength   = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength   = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             # hard floor on |chosen|
             min_chosen_tokens     = config.get("ftpo_sample_min_chosen_tokens", 3),
             # overall training-set cap (used for per-token quotas too)
@@ -158,7 +158,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len                  = max_seq_length,
             rejected_reg_strength        = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength          = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength          = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens            = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples           = config.get("finetune_max_train_examples"),
         )
@@ -247,7 +247,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len               = max_seq_length,
             rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples        = config.get("finetune_max_train_examples"),
         )
@@ -299,7 +299,7 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
             experiment_run_dir = experiment_run_dir,
             max_seq_len               = max_seq_length,
             rejected_reg_strength     = config.get("ftpo_sample_rejected_regularisation_strength", 0.8),
-            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.2),
+            chosen_reg_strength       = config.get("ftpo_sample_chosen_regularisation_strength", 0.0),
             min_chosen_tokens         = config.get("ftpo_sample_min_chosen_tokens", 3),
             max_train_examples        = config.get("finetune_max_train_examples"),
         )
@@ -345,4 +345,4 @@ def load_and_prepare_dataset(config: dict, experiment_run_dir: Path, tokenizer:
         logger.error(f"Unknown finetune_mode '{mode}'. Use 'dpo' or 'ftpo'.")
         return None
 
-    return dpo_dataset_hf
+    return dpo_dataset_hf

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:

utils/whitelist.py

@@ -75,7 +75,7 @@ class WhitelistBuilder:
             tokenizer.pad_token,
             tokenizer.cls_token,
             tokenizer.sep_token,
-            *(tokenizer.additional_special_tokens or []),
+            *(getattr(tokenizer, "additional_special_tokens", []) or []),
         ]
         for raw_text in special_token_texts:
             if not raw_text: