Add bounded bitsandbytes INT8 and NF4 controls

README.md

@@ -16,3 +16,7 @@ A bounded ZeroGPU control lane for the matched local MLX embedding
 quantization experiment. The first milestone reproduces the frozen
 Qwen3-Embedding-0.6B BF16 vectors on CUDA and compares them with the saved MLX
 BF16 vectors. CUDA results do not reproduce MLX/Metal performance.
+
+The second milestone adds CUDA-native bitsandbytes INT8 and NF4 controls.
+They are deliberately reported as separate quantizers and are not treated as
+equivalents of MLX Q, oQ, or oQe formats.

app.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import json
 import os
 import time
+import gc
 from pathlib import Path
 
 import gradio as gr
@@ -10,7 +11,7 @@ import numpy as np
 import spaces
 import torch
 import torch.nn.functional as F
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 
 
 MODEL_PATH = Path("/models/qwen3-embedding-0.6b")
@@ -18,6 +19,7 @@ DATA_ROOT = Path(os.environ.get("REPRO_DATA_ROOT", "/data"))
 INPUT_PATH = DATA_ROOT / "inputs/retrieval_pairs.json"
 MLX_REFERENCE = DATA_ROOT / "local-reference/qwen3-embedding-0.6b/bf16.npz"
 OUTPUT_DIR = DATA_ROOT / "cloud-results/qwen3-embedding-0.6b"
+CUDA_BF16_REFERENCE = OUTPUT_DIR / "cuda-bf16.npz"
 TASK = (
     "Given a natural-language search query, retrieve the single passage "
     "that best answers it"
@@ -36,15 +38,19 @@ def detailed_instruction(query: str) -> str:
     return f"Instruct: {TASK}\nQuery:{query}"
 
 
-def encode(text: str) -> np.ndarray:
+def encode_with(active_model, text: str) -> np.ndarray:
     batch = tokenizer(text, return_tensors="pt", truncation=True, max_length=32768)
     batch = {key: value.to("cuda") for key, value in batch.items()}
     with torch.inference_mode():
-        output = model(**batch)
+        output = active_model(**batch)
         vector = F.normalize(output.last_hidden_state[:, -1, :].float(), p=2, dim=-1)
     return vector[0].cpu().numpy()
 
 
+def encode(text: str) -> np.ndarray:
+    return encode_with(model, text)
+
+
 def retrieval_metrics(queries: np.ndarray, documents: np.ndarray) -> tuple[dict, np.ndarray, np.ndarray]:
     scores = queries @ documents.T
     order = np.argsort(-scores, axis=1)
@@ -123,6 +129,117 @@ def run_bf16_control() -> dict:
     return result
 
 
+def compare_vectors(
+    queries: np.ndarray,
+    documents: np.ndarray,
+    scores: np.ndarray,
+    ranks: np.ndarray,
+    reference_path: Path,
+    prefix: str,
+) -> dict | None:
+    if not reference_path.exists():
+        return None
+    reference = np.load(reference_path)
+    ref_all = np.concatenate([reference["queries"], reference["documents"]])
+    candidate_all = np.concatenate([queries, documents])
+    aligned = np.sum(ref_all * candidate_all, axis=1)
+    score_delta = scores - reference["scores"]
+    return {
+        f"mean_aligned_cosine_{prefix}": float(aligned.mean()),
+        f"minimum_aligned_cosine_{prefix}": float(aligned.min()),
+        f"score_rmse_{prefix}": float(np.sqrt(np.mean(score_delta ** 2))),
+        f"queries_with_rank_change_{prefix}": int(
+            np.count_nonzero(ranks - reference["ranks"])
+        ),
+    }
+
+
+QUANTIZERS = {
+    "bnb-int8": BitsAndBytesConfig(load_in_8bit=True),
+    "bnb-nf4": BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=False,
+    ),
+}
+
+
+@spaces.GPU(duration=600)
+def run_quantized_control(variant: str) -> dict:
+    if variant not in QUANTIZERS:
+        raise ValueError(f"unsupported quantizer: {variant}")
+    pairs = json.loads(INPUT_PATH.read_text())
+    query_texts = [detailed_instruction(item["query"]) for item in pairs]
+    document_texts = [item["document"] for item in pairs]
+
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    allocation_before = int(torch.cuda.memory_allocated())
+    load_started = time.perf_counter()
+    quantized_model = AutoModel.from_pretrained(
+        MODEL_PATH,
+        quantization_config=QUANTIZERS[variant],
+        device_map={"": 0},
+        trust_remote_code=True,
+    ).eval()
+    torch.cuda.synchronize()
+    load_seconds = time.perf_counter() - load_started
+    allocation_after_load = int(torch.cuda.memory_allocated())
+
+    encode_started = time.perf_counter()
+    queries = np.stack([encode_with(quantized_model, text) for text in query_texts])
+    documents = np.stack([
+        encode_with(quantized_model, text) for text in document_texts
+    ])
+    torch.cuda.synchronize()
+    encode_seconds = time.perf_counter() - encode_started
+    metrics, scores, ranks = retrieval_metrics(queries, documents)
+    metrics.update({
+        "lane": f"cuda-zerogpu-{variant}",
+        "model": "Qwen/Qwen3-Embedding-0.6B",
+        "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
+        "quantizer": variant,
+        "quantization_config": QUANTIZERS[variant].to_dict(),
+        "load_seconds": load_seconds,
+        "encode_seconds": encode_seconds,
+        "texts_per_second": len(query_texts + document_texts) / encode_seconds,
+        "torch_version": torch.__version__,
+        "cuda_device": torch.cuda.get_device_name(0),
+        "cuda_allocation_before_load": allocation_before,
+        "cuda_allocation_after_load": allocation_after_load,
+        "cuda_incremental_model_allocation": allocation_after_load - allocation_before,
+        "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),
+    })
+    result = {
+        "metrics": metrics,
+        "cuda_bf16_comparison": compare_vectors(
+            queries, documents, scores, ranks, CUDA_BF16_REFERENCE, "vs_cuda_bf16"
+        ),
+        "mlx_bf16_comparison": compare_vectors(
+            queries, documents, scores, ranks, MLX_REFERENCE, "vs_mlx_bf16"
+        ),
+    }
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    stem = f"cuda-{variant}"
+    np.savez_compressed(
+        OUTPUT_DIR / f"{stem}.npz",
+        queries=queries,
+        documents=documents,
+        scores=scores,
+        ranks=ranks,
+        metrics=np.array(json.dumps(metrics)),
+    )
+    (OUTPUT_DIR / f"{stem}.json").write_text(json.dumps(result, indent=2) + "\n")
+
+    del quantized_model
+    gc.collect()
+    torch.cuda.empty_cache()
+    return result
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         "# Embedding Quantization CUDA Control\n"
@@ -137,6 +254,22 @@ with gr.Blocks() as demo:
         concurrency_limit=1,
         api_name="run_bf16_control",
     )
+    gr.Markdown(
+        "## CUDA-native quantizer controls\n"
+        "These are bitsandbytes INT8/NF4 controls, not MLX Q/oQ/oQe replicas."
+    )
+    quantizer = gr.Dropdown(
+        choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"
+    )
+    quant_button = gr.Button("Run bounded CUDA quantizer control")
+    quant_output = gr.JSON(label="Quantized result")
+    quant_button.click(
+        fn=run_quantized_control,
+        inputs=quantizer,
+        outputs=quant_output,
+        concurrency_limit=1,
+        api_name="run_quantized_control",
+    )
 
 
 if __name__ == "__main__":

requirements.txt

@@ -2,3 +2,5 @@ transformers>=4.51,<5
 numpy>=2,<3
 spaces>=0.45
 torch>=2.8
+accelerate>=1.2,<2
+bitsandbytes>=0.48,<1