Add native 4-bit/8-bit mixed-precision loading in joint_schema_model.py and GGUF documentation

README.md

@@ -5,6 +5,8 @@ license: apache-2.0
 pipeline_tag: text-generation
 tags:
   - fastbrain
+  - clef
+  - gguf
   - systemone
   - decision-index
   - tool-calling
@@ -82,19 +84,25 @@ Evaluated on the official **Cloudflare Decision Index (v0.2.1)** across **10 ben
 
 ## Quickstart
 
-`TypeScript/FastBrain-9B` is a drop-in replacement for `Cloudflare/clef-flash` using the standard `joint_schema_model` runtime:
+### 1. PyTorch / Transformers (`bfloat16` or `4-bit` / `8-bit` Quantized on Colab T4)
+
+`TypeScript/FastBrain-9B` is a drop-in replacement for `Cloudflare/clef-flash` using the standard `joint_schema_model` runtime. It supports both full-precision `bfloat16` (`18.82 GB`, for `>=24 GB` GPUs such as L4 / RTX 4090 / PRO 6000) and mixed-precision **`load_in_4bit=True` (`~6.8 GB` VRAM, runs on a free Google Colab T4 16GB GPU)** or **`load_in_8bit=True` (`~11.2 GB` VRAM)** while keeping the `116M` `JointSchemaHead` in native `bfloat16`/`float16`:
 
 ```python
 from huggingface_hub import snapshot_download
 import sys
 import torch
 
-model_dir = snapshot_download("TypeScript/FastBrain-9B")
+model_dir = snapshot_download("TypeScript/FastBrain-9B", ignore_patterns=["*.gguf"])
 sys.path.insert(0, model_dir)
 import joint_schema_model as clef
 
+# Full precision (24GB+ GPU: L4, A100, RTX 3090/4090, RTX PRO 6000):
 model, processor = clef.load_release_model(model_dir, device="cuda", dtype=torch.bfloat16)
 
+# Or 4-bit NF4 mixed-precision (8GB–16GB GPU, e.g., free Google Colab T4):
+# model, processor = clef.load_release_model(model_dir, device="cuda", dtype=torch.float16, load_in_4bit=True)
+
 response = clef.systemone(model, processor, {
     "model": "TypeScript/FastBrain-9B",
     "state": {
@@ -116,14 +124,55 @@ response = clef.systemone(model, processor, {
 print(response)
 ```
 
+### 2. `llama.cpp` GGUF (`Q4_K_M` / `Q8_0` / `BF16`) via `/v1/systemone`
+
+This repository includes official `llama.cpp` (`b11430`+ / `ClefModel`) GGUF quantizations that pack both the `Qwen3.5-9B` hybrid backbone and the `116M` `JointSchemaHead` (`dec.*` tensors preserved at `Q8_0` precision):
+
+```bash
+# Start llama-server on the /v1/systemone decision endpoint (6.04 GB Q4_K_M quant)
+llama-server -hf TypeScript/FastBrain-9B:Q4_K_M --port 8080
+
+# Query the single-forward-pass /v1/systemone endpoint
+curl http://127.0.0.1:8080/v1/systemone \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "FastBrain-9B",
+    "state": "Customer message: I was charged twice for my order last week and nobody has replied.",
+    "questions": {
+      "route": {
+        "type": "choice",
+        "instructions": "Which team should handle this?",
+        "criteria": {"billing": "Payments or invoices", "technical": "Bugs or outages"}
+      },
+      "urgent": {"type": "noul", "instructions": "Is the customer requesting urgent resolution?"}
+    }
+  }'
+```
+
+---
+
+## Hardware Footprint, Quantization Tiers & Inference Speed
+
+| Format / Precision Tier | Weight Footprint | Peak VRAM | Minimum Hardware | Typical Latency (`500–1k tok`) |
+| :--- | :---: | :---: | :--- | :---: |
+| **`FastBrain-9B-Q4_K_M.gguf`** *(4-bit + `Q8_0` Head)* | `6.04 GB` | `~7.2 GB` | **8GB GPU / Colab Free T4 (16GB) / Mac M1–M4** | `~90–220 ms` |
+| **PyTorch `load_in_4bit=True`** *(NF4 + `fp16` Head)* | `~5.5 GB` | `~6.8 GB` | **Colab Free T4 (16GB) / RTX 3060–4060 (8GB+)** | `~180–280 ms` |
+| **`FastBrain-9B-Q8_0.gguf`** *(Full 8-bit GGUF)* | `9.68 GB` | `~11.0 GB` | **12GB–16GB GPU / Colab T4 / Mac 16GB+** | `~80–180 ms` |
+| **PyTorch `load_in_8bit=True`** *(INT8 + `bf16` Head)* | `~9.6 GB` | `~11.2 GB` | **Colab T4 (16GB) / RTX 4070 Ti (12GB+)** | `~130–210 ms` |
+| **Native `bfloat16` (`safetensors` / `BF16.gguf`)** | `18.82 GB` (`18.16 GB`) | `20.3–24.7 GB` | **Colab L4 (24GB) / RTX 3090–4090 / PRO 6000 (96GB)** | **`38.8 ms` (`B=1`) · `61.7 ms` (`B=16` on G4 6000)** |
+
 ---
 
 ## Repository Structure
 
-| File | Description |
-| :--- | :--- |
-| `model-00001-of-00004.safetensors` .. `00004` | Full `9B` backbone weights (`bfloat16`, in-place merged, no LoRA required) |
-| `joint_head.safetensors` & `joint_head_config.json` | `116M`-parameter `JointSchemaHead` weights (`122` tensors) and configuration |
-| `joint_schema_model.py` | `ClefModel`, `JointSchemaHead`, `load_release_model()`, and `systemone()` runtime |
-| `config.json`, `tokenizer.json`, `processor_config.json` | Standard Hugging Face model, tokenizer, and processor configurations |
-| `assets/*.png` | High-resolution technical benchmark and calibration charts |
+| File | Size | Description |
+| :--- | :---: | :--- |
+| `model-00001-of-00004.safetensors` .. `00004` | `18.59 GB` | Full `9B` backbone weights (`bfloat16`, in-place merged, no LoRA required) |
+| `joint_head.safetensors` & `joint_head_config.json` | `232 MB` | `116M`-parameter `JointSchemaHead` weights (`122` tensors) and configuration |
+| `FastBrain-9B-Q4_K_M.gguf` | `6.04 GB` | Recommended 4-bit GGUF (`dec.*` decision head, `ssm_`, and `attn_` preserved at `Q8_0`) |
+| `FastBrain-9B-Q8_0.gguf` | `9.68 GB` | High-precision 8-bit GGUF (`ClefModel` architecture for `llama-server`) |
+| `FastBrain-9B-BF16.gguf` | `18.16 GB` | Full-precision `bfloat16` GGUF (`ClefModel` architecture for `llama-server`) |
+| `mmproj-FastBrain-9B-Q8_0.gguf` & `BF16.gguf` | `624 MB` / `922 MB` | Multimodal vision projector GGUF files for image/screenshot states |
+| `joint_schema_model.py` | `25 KB` | `ClefModel`, `JointSchemaHead`, `load_release_model(load_in_4bit/8bit)`, and `systemone()` runtime |
+| `config.json`, `tokenizer.json`, `processor_config.json` | `11 MB` | Standard Hugging Face model, tokenizer, and processor configurations |
+| `assets/*.png` | `1.2 MB` | High-resolution technical benchmark and calibration charts |

joint_schema_model.py

@@ -482,12 +482,16 @@ class ClefModel(torch.nn.Module):
             return_dict=True,
             **media,
         )
+        head_dtype = next(self.head.parameters()).dtype
+        out_emb = base_model.get_output_embeddings().weight
+        if hasattr(out_emb, "dequantize"):
+            out_emb = out_emb.dequantize()
         return self.head(
-            outputs.last_hidden_state,
+            outputs.last_hidden_state.to(dtype=head_dtype),
             batch["input_ids"],
             batch["attention_mask"],
             batch["records"],
-            base_model.get_output_embeddings().weight,
+            out_emb.to(dtype=head_dtype),
         )
 
 
@@ -495,9 +499,17 @@ def load_release_model(
     model_path: str | Path,
     device: str | torch.device = "cuda",
     dtype: torch.dtype = torch.bfloat16,
+    load_in_4bit: bool = False,
+    load_in_8bit: bool = False,
     **from_pretrained_kwargs: Any,
 ) -> tuple[ClefModel, Any]:
-    """Load a Clef release (merged backbone, joint schema head, and processor)."""
+    """Load a FastBrain / Clef release (merged backbone, joint schema head, and processor).
+
+    Supports full-precision ``bfloat16``/``float16`` as well as mixed-precision
+    4-bit NF4 (``load_in_4bit=True``, ~6.8 GB VRAM on Colab T4) and 8-bit
+    (``load_in_8bit=True``, ~11.2 GB VRAM) quantization via ``bitsandbytes`` while
+    preserving the 116M ``JointSchemaHead`` in native ``dtype`` precision.
+    """
     from huggingface_hub import snapshot_download
     from safetensors.torch import load_file
     from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration
@@ -505,6 +517,23 @@ def load_release_model(
     path = Path(model_path)
     if not path.is_dir():
         path = Path(snapshot_download(str(model_path)))
+    from_pretrained_kwargs = dict(from_pretrained_kwargs)
+    if load_in_4bit or load_in_8bit:
+        from transformers import BitsAndBytesConfig
+
+        if load_in_4bit:
+            from_pretrained_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=dtype,
+                bnb_4bit_use_double_quant=True,
+                llm_int8_skip_modules=["lm_head", "visual"],
+            )
+        else:
+            from_pretrained_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_8bit=True,
+                llm_int8_skip_modules=["lm_head", "visual"],
+            )
     backbone = Qwen3_5ForConditionalGeneration.from_pretrained(
         path,
         dtype=dtype,
@@ -560,7 +589,7 @@ def systemone(model: ClefModel, processor: Any, request: dict[str, Any], max_len
         if question["type"] != "noul" and not question.get("criteria"):
             raise ValueError(f"{question_id}: criteria must not be empty")
     encoded = encode_record(processor.tokenizer, request, max_length=max_length, processor=processor)
-    device = next(model.parameters()).device
+    device = next(model.head.parameters()).device
     logits = model(collate_records([encoded], processor.tokenizer.pad_token_id, device))[0]
     answers = {
         question.question_id: systemone_answer(