Keep classifier head in fp32 so the checkpoint loads on Linux CPU

README.md

@@ -36,11 +36,17 @@ training data and evaluation.
 | Method | bitsandbytes NF4, weight-only |
 | Compute dtype | float32 |
 | Double quantization | off |
-| Quantized modules | all `torch.nn.Linear` (attention, dense, pooler, classifier) |
+| Quantized modules | 145 `torch.nn.Linear` layers (attention, dense, pooler) |
+| Kept in fp32 | classifier head, embeddings, LayerNorm |
 | Checkpoint size | ~699 MB (original fp32: ~1.74 GB) |
 | Resident memory | ~648 MB (original fp32: ~1,660 MB) |
 
-Embeddings, LayerNorm, and the sigmoid head remain in full precision.
+The classifier head is deliberately left in fp32. It is a `[1, 1024]` matrix, and
+bitsandbytes' packed CPU kernel asserts that each quantized layer's output
+dimension is divisible by its block size, which 1 is not. Quantizing the head
+works on Apple Silicon MPS but raises
+`AssertionError: N must be divisible by block_n` on Linux CPU, so the head is kept
+exact to make one checkpoint that loads everywhere.
 
 ## Usage
 
@@ -86,8 +92,8 @@ measurements below. It is not needed to use the model; see
 Measured against the fp32 weights of the original model on 16 held-in prompts
 spanning clearly human to clearly AI text:
 
-- mean absolute change in P(AI): 0.021
-- maximum absolute change in P(AI): 0.080
+- mean absolute change in P(AI): 0.020
+- maximum absolute change in P(AI): 0.076
 - verdict flips at a 0.5 decision threshold: 0
 
 Quantized scores are close but not identical to the fp32 model, and the largest
@@ -126,4 +132,4 @@ additional error on top of that. MIT license, inherited from the original model.
   journal={Working Notes of CLEF},
   year={2026}
 }
-```
+```

config.json

@@ -44,7 +44,9 @@
     "bnb_4bit_use_double_quant": false,
     "llm_int8_enable_fp32_cpu_offload": false,
     "llm_int8_has_fp16_weight": false,
-    "llm_int8_skip_modules": null,
+    "llm_int8_skip_modules": [
+      "classifier"
+    ],
     "llm_int8_threshold": 6.0,
     "load_in_4bit": true,
     "load_in_8bit": false,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ae345b48c7f26cde6edeaf860a107a3fc2f54fed0b055bfe9a50ffbf9b7c97f9
-size 698682667
+oid sha256:b6d34cd0e7f8744d9954f9594b320213ac847a399516182f79bf73f1b75504fe
+size 698685757

scripts/README.md

@@ -21,7 +21,10 @@ uploaded directly. Optional argument: output directory, defaulting to
 python scripts/quantize.py
 ```
 
-It finishes by reloading the saved checkpoint and printing a P(AI) sanity value.
+It keeps the classifier head in fp32, because bitsandbytes' packed CPU kernel
+requires each quantized layer's output dimension to divide evenly by its block
+size and the head is `[1, 1024]`. It then asserts that no quantized layer would
+break that kernel, and prints a reload sanity value.
 
 ## `bench_quant.py`
 

scripts/quantize.py

@@ -1,7 +1,7 @@
 """Save a 4-bit NF4 checkpoint of the detector for fast repeated loading.
 
-Writes to models/gradient-ai-text-detector-4bit, which app.py picks up
-automatically. The quantized checkpoint is roughly 650 MB instead of 1.7 GB.
+Writes to models/gradient-ai-text-detector-4bit by default. The quantized
+checkpoint is roughly 650 MB instead of 1.7 GB.
 
 Usage:
     .venv/bin/python scripts/quantize.py [output_dir]
@@ -22,6 +22,30 @@ DEFAULT_OUT = Path(__file__).parent.parent / "models" / "gradient-ai-text-detect
 
 SAMPLE = "In today's rapidly evolving digital landscape, organizations must unlock value."
 
+# bitsandbytes' packed CPU kernel asserts that each quantized weight's output
+# dimension divides evenly by its block size. The classifier head is [1, 1024],
+# so it must stay in fp32 or the checkpoint cannot run on Linux CPU (Spaces).
+SKIP_MODULES = ["classifier"]
+BLOCK = 64
+
+
+def check_cpu_compatible(model):
+    """Fail loudly rather than ship a checkpoint only macOS can load."""
+    quantized = [
+        module
+        for module in model.modules()
+        if type(module).__name__ == "Linear4bit"
+    ]
+    offenders = sorted(
+        {module.out_features for module in quantized if module.out_features % BLOCK}
+    )
+    if offenders:
+        raise SystemExit(
+            f"quantized layers with out_features {offenders} are not divisible by "
+            f"{BLOCK}; add them to SKIP_MODULES"
+        )
+    return len(quantized)
+
 
 def main(out_dir):
     out_dir = Path(out_dir)
@@ -33,10 +57,13 @@ def main(out_dir):
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.float32,
+            llm_int8_skip_modules=SKIP_MODULES,
         ),
     )
     model.eval()
     model.to(torch.device("cpu"))
+    print(f"quantized {check_cpu_compatible(model)} linear layers")
+    print(f"kept in fp32: {SKIP_MODULES} ({model.classifier.weight.dtype})")
     model.save_pretrained(out_dir, safe_serialization=True)
     AutoTokenizer.from_pretrained(MODEL_ID).save_pretrained(out_dir)
     AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False).save_pretrained(out_dir)