TiGa-RCE/embedding-quantization-cuda-control
0
1from __future__ import annotations2 3import json4import os5import time6import gc7import traceback8from pathlib import Path9 10import gradio as gr11import numpy as np12import spaces13import torch14import torch.nn.functional as F15from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig16 17 18MODEL_PATH = Path("/models/qwen3-embedding-0.6b")19DATA_ROOT = Path(os.environ.get("REPRO_DATA_ROOT", "/data"))20INPUT_PATH = DATA_ROOT / "inputs/retrieval_pairs.json"21MLX_REFERENCE = DATA_ROOT / "local-reference/qwen3-embedding-0.6b/bf16.npz"22OUTPUT_DIR = DATA_ROOT / "cloud-results/qwen3-embedding-0.6b"23CUDA_BF16_REFERENCE = OUTPUT_DIR / "cuda-bf16.npz"24TASK = (25 "Given a natural-language search query, retrieve the single passage "26 "that best answers it"27)28 29 30tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side="left")31# The completed 0.6B phase is intentionally not resident while the larger GTE32# startup path is validated. Packing both models together exceeded the33# current Space startup envelope at 6.16 GB.34model = None35 36# ZeroGPU optimizes CUDA placements made during module startup. The previous37# larger-family path loaded this model inside the decorated call and exhausted38# that reservation before producing an artifact.39GTE_MODEL_PATH = Path("/models/gte-qwen2-1.5b")40gte_tokenizer = None41gte_model = None42 43QWEN8_MODEL_PATH = Path("/models/qwen3-embedding-8b")44qwen8_tokenizer = AutoTokenizer.from_pretrained(QWEN8_MODEL_PATH, padding_side="left")45qwen8_model = None46 47 48def detailed_instruction(query: str) -> str:49 return f"Instruct: {TASK}\nQuery:{query}"50 51 52def encode_with(active_model, text: str, active_tokenizer=None) -> np.ndarray:53 active_tokenizer = active_tokenizer or tokenizer54 batch = active_tokenizer(55 text, return_tensors="pt", truncation=True, max_length=3276856 )57 batch = {key: value.to("cuda") for key, value in batch.items()}58 with torch.inference_mode():59 output = active_model(**batch, use_cache=False)60 vector = F.normalize(output.last_hidden_state[:, -1, :].float(), p=2, dim=-1)61 return vector[0].cpu().numpy()62 63 64def encode(text: str) -> np.ndarray:65 return encode_with(model, text)66 67 68def retrieval_metrics(queries: np.ndarray, documents: np.ndarray) -> tuple[dict, np.ndarray, np.ndarray]:69 scores = queries @ documents.T70 order = np.argsort(-scores, axis=1)71 ranks = np.array([72 int(np.where(order[index] == index)[0][0]) + 173 for index in range(len(queries))74 ])75 positive = np.diag(scores)76 negative = scores.copy()77 np.fill_diagonal(negative, -np.inf)78 margins = positive - negative.max(axis=1)79 return {80 "pair_count": len(queries),81 "top1": float(np.mean(ranks == 1)),82 "recall_at_5": float(np.mean(ranks <= 5)),83 "mrr": float(np.mean(1.0 / ranks)),84 "mean_margin": float(margins.mean()),85 "minimum_margin": float(margins.min()),86 "mean_rank": float(ranks.mean()),87 "worst_rank": int(ranks.max()),88 }, scores, ranks89 90 91@spaces.GPU(duration=300)92def run_bf16_control() -> dict:93 if model is None:94 raise RuntimeError("0.6B BF16 control is offline during the GTE phase")95 pairs = json.loads(INPUT_PATH.read_text())96 query_texts = [detailed_instruction(item["query"]) for item in pairs]97 document_texts = [item["document"] for item in pairs]98 99 if torch.cuda.is_available():100 torch.cuda.reset_peak_memory_stats()101 torch.cuda.synchronize()102 started = time.perf_counter()103 queries = np.stack([encode(text) for text in query_texts])104 documents = np.stack([encode(text) for text in document_texts])105 if torch.cuda.is_available():106 torch.cuda.synchronize()107 elapsed = time.perf_counter() - started108 metrics, scores, ranks = retrieval_metrics(queries, documents)109 metrics.update({110 "lane": "cuda-zerogpu-bf16",111 "model": "Qwen/Qwen3-Embedding-0.6B",112 "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",113 "elapsed_seconds": elapsed,114 "texts_per_second": len(query_texts + document_texts) / elapsed,115 "torch_version": torch.__version__,116 "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,117 "cuda_peak_bytes": torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None,118 })119 120 comparison = None121 if MLX_REFERENCE.exists():122 reference = np.load(MLX_REFERENCE)123 ref_all = np.concatenate([reference["queries"], reference["documents"]])124 cuda_all = np.concatenate([queries, documents])125 aligned = np.sum(ref_all * cuda_all, axis=1)126 score_delta = scores - reference["scores"]127 comparison = {128 "mean_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.mean()),129 "minimum_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.min()),130 "score_rmse_cuda_vs_mlx_bf16": float(np.sqrt(np.mean(score_delta ** 2))),131 "queries_with_rank_change": int(np.count_nonzero(ranks - reference["ranks"])),132 }133 134 result = {"metrics": metrics, "mlx_bf16_comparison": comparison}135 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)136 np.savez_compressed(137 OUTPUT_DIR / "cuda-bf16.npz",138 queries=queries,139 documents=documents,140 scores=scores,141 ranks=ranks,142 metrics=np.array(json.dumps(metrics)),143 )144 (OUTPUT_DIR / "cuda-bf16.json").write_text(json.dumps(result, indent=2) + "\n")145 return result146 147 148def compare_vectors(149 queries: np.ndarray,150 documents: np.ndarray,151 scores: np.ndarray,152 ranks: np.ndarray,153 reference_path: Path,154 prefix: str,155) -> dict | None:156 if not reference_path.exists():157 return None158 reference = np.load(reference_path)159 ref_all = np.concatenate([reference["queries"], reference["documents"]])160 candidate_all = np.concatenate([queries, documents])161 aligned = np.sum(ref_all * candidate_all, axis=1)162 score_delta = scores - reference["scores"]163 return {164 f"mean_aligned_cosine_{prefix}": float(aligned.mean()),165 f"minimum_aligned_cosine_{prefix}": float(aligned.min()),166 f"score_rmse_{prefix}": float(np.sqrt(np.mean(score_delta ** 2))),167 f"queries_with_rank_change_{prefix}": int(168 np.count_nonzero(ranks - reference["ranks"])169 ),170 }171 172 173QUANTIZERS = ("bnb-int8", "bnb-nf4")174 175 176def quantization_config(variant: str) -> BitsAndBytesConfig:177 if variant == "bnb-int8":178 return BitsAndBytesConfig(load_in_8bit=True)179 if variant == "bnb-nf4":180 return BitsAndBytesConfig(181 load_in_4bit=True,182 bnb_4bit_quant_type="nf4",183 bnb_4bit_compute_dtype=torch.bfloat16,184 bnb_4bit_use_double_quant=False,185 )186 raise ValueError(f"unsupported quantizer: {variant}")187 188 189@spaces.GPU(duration=300)190def run_quantized_control(variant: str) -> dict:191 if variant not in QUANTIZERS:192 raise ValueError(f"unsupported quantizer: {variant}")193 config = quantization_config(variant)194 pairs = json.loads(INPUT_PATH.read_text())195 query_texts = [detailed_instruction(item["query"]) for item in pairs]196 document_texts = [item["document"] for item in pairs]197 198 gc.collect()199 torch.cuda.empty_cache()200 torch.cuda.reset_peak_memory_stats()201 torch.cuda.synchronize()202 allocation_before = int(torch.cuda.memory_allocated())203 load_started = time.perf_counter()204 quantized_model = AutoModel.from_pretrained(205 MODEL_PATH,206 quantization_config=config,207 device_map={"": 0},208 trust_remote_code=True,209 ).eval()210 torch.cuda.synchronize()211 load_seconds = time.perf_counter() - load_started212 allocation_after_load = int(torch.cuda.memory_allocated())213 214 encode_started = time.perf_counter()215 queries = np.stack([encode_with(quantized_model, text) for text in query_texts])216 documents = np.stack([217 encode_with(quantized_model, text) for text in document_texts218 ])219 torch.cuda.synchronize()220 encode_seconds = time.perf_counter() - encode_started221 metrics, scores, ranks = retrieval_metrics(queries, documents)222 metrics.update({223 "lane": f"cuda-zerogpu-{variant}",224 "model": "Qwen/Qwen3-Embedding-0.6B",225 "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",226 "quantizer": variant,227 "quantization_config": config.to_dict(),228 "load_seconds": load_seconds,229 "encode_seconds": encode_seconds,230 "texts_per_second": len(query_texts + document_texts) / encode_seconds,231 "torch_version": torch.__version__,232 "cuda_device": torch.cuda.get_device_name(0),233 "cuda_allocation_before_load": allocation_before,234 "cuda_allocation_after_load": allocation_after_load,235 "cuda_incremental_model_allocation": allocation_after_load - allocation_before,236 "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),237 })238 result = {239 "metrics": metrics,240 "cuda_bf16_comparison": compare_vectors(241 queries, documents, scores, ranks, CUDA_BF16_REFERENCE, "vs_cuda_bf16"242 ),243 "mlx_bf16_comparison": compare_vectors(244 queries, documents, scores, ranks, MLX_REFERENCE, "vs_mlx_bf16"245 ),246 }247 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)248 stem = f"cuda-{variant}"249 np.savez_compressed(250 OUTPUT_DIR / f"{stem}.npz",251 queries=queries,252 documents=documents,253 scores=scores,254 ranks=ranks,255 metrics=np.array(json.dumps(metrics)),256 )257 (OUTPUT_DIR / f"{stem}.json").write_text(json.dumps(result, indent=2) + "\n")258 259 del quantized_model260 gc.collect()261 torch.cuda.empty_cache()262 return result263 264 265BF16_FAMILIES = {266 "gte-qwen2-1.5b": {267 "path": Path("/models/gte-qwen2-1.5b"),268 "source": "Alibaba-NLP/gte-Qwen2-1.5B-instruct",269 "revision": "a9af15a6372d7d6b25e9fb07c2ccb9e1fe645644",270 },271 "qwen3-embedding-8b": {272 "path": Path("/models/qwen3-embedding-8b"),273 "source": "Qwen/Qwen3-Embedding-8B",274 "revision": "1d8ad4ca9b3dd8059ad90a75d4983776a23d44af",275 },276}277 278 279def _run_family_bf16_control(family: str) -> dict:280 if family not in BF16_FAMILIES:281 raise ValueError(f"unsupported family: {family}")282 spec = BF16_FAMILIES[family]283 pairs = json.loads(INPUT_PATH.read_text())284 query_texts = [detailed_instruction(item["query"]) for item in pairs]285 document_texts = [item["document"] for item in pairs]286 287 torch.cuda.reset_peak_memory_stats()288 torch.cuda.synchronize()289 allocation_before = int(torch.cuda.memory_allocated())290 family_tokenizer = AutoTokenizer.from_pretrained(291 spec["path"], padding_side="left", trust_remote_code=True292 )293 load_started = time.perf_counter()294 family_model = AutoModel.from_pretrained(295 spec["path"],296 dtype=torch.bfloat16,297 trust_remote_code=True,298 low_cpu_mem_usage=True,299 device_map={"": "cuda"},300 ).eval()301 torch.cuda.synchronize()302 load_seconds = time.perf_counter() - load_started303 loading_strategy = "in-call-direct-cuda-device-map"304 owns_model = True305 allocation_after_load = int(torch.cuda.memory_allocated())306 307 encode_started = time.perf_counter()308 queries = np.stack([309 encode_with(family_model, text, family_tokenizer) for text in query_texts310 ])311 documents = np.stack([312 encode_with(family_model, text, family_tokenizer) for text in document_texts313 ])314 torch.cuda.synchronize()315 encode_seconds = time.perf_counter() - encode_started316 metrics, scores, ranks = retrieval_metrics(queries, documents)317 metrics.update({318 "lane": "cuda-zerogpu-bf16",319 "family": family,320 "model": spec["source"],321 "source_revision": spec["revision"],322 "load_seconds": load_seconds,323 "loading_strategy": loading_strategy,324 "encode_seconds": encode_seconds,325 "texts_per_second": len(query_texts + document_texts) / encode_seconds,326 "torch_version": torch.__version__,327 "cuda_device": torch.cuda.get_device_name(0),328 "cuda_allocation_before_load": allocation_before,329 "cuda_allocation_after_load": allocation_after_load,330 "cuda_incremental_model_allocation": allocation_after_load - allocation_before,331 "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),332 })333 local_reference = (334 DATA_ROOT / "local-results/full-q4-q6-q8" / family / "quality/bf16.npz"335 )336 result = {337 "metrics": metrics,338 "mlx_bf16_comparison": compare_vectors(339 queries, documents, scores, ranks, local_reference, "cuda_vs_mlx_bf16"340 ),341 "previous_cuda_bf16_comparison": compare_vectors(342 queries,343 documents,344 scores,345 ranks,346 DATA_ROOT / "cloud-results" / family / "cuda-bf16-root-pack.npz",347 "direct_vs_root_pack_bf16",348 ),349 }350 output_dir = DATA_ROOT / "cloud-results" / family351 output_dir.mkdir(parents=True, exist_ok=True)352 np.savez_compressed(353 output_dir / "cuda-bf16.npz",354 queries=queries,355 documents=documents,356 scores=scores,357 ranks=ranks,358 metrics=np.array(json.dumps(metrics)),359 )360 (output_dir / "cuda-bf16.json").write_text(json.dumps(result, indent=2) + "\n")361 362 if owns_model:363 del family_model, family_tokenizer364 gc.collect()365 torch.cuda.empty_cache()366 return result367 368 369@spaces.GPU(duration=300)370def run_family_bf16_control(family: str) -> dict:371 try:372 return _run_family_bf16_control(family)373 except Exception as error:374 return {375 "diagnostic_error": type(error).__name__,376 "diagnostic_message": str(error),377 "diagnostic_traceback": traceback.format_exc(),378 }379 380 381def _run_family_quantized_control(family: str, variant: str) -> dict:382 if family not in BF16_FAMILIES:383 raise ValueError(f"unsupported family: {family}")384 if variant not in QUANTIZERS:385 raise ValueError(f"unsupported quantizer: {variant}")386 387 spec = BF16_FAMILIES[family]388 config = quantization_config(variant)389 pairs = json.loads(INPUT_PATH.read_text())390 query_texts = [detailed_instruction(item["query"]) for item in pairs]391 document_texts = [item["document"] for item in pairs]392 family_tokenizer = AutoTokenizer.from_pretrained(393 spec["path"], padding_side="left", trust_remote_code=True394 )395 396 gc.collect()397 torch.cuda.empty_cache()398 torch.cuda.reset_peak_memory_stats()399 torch.cuda.synchronize()400 allocation_before = int(torch.cuda.memory_allocated())401 load_started = time.perf_counter()402 family_model = AutoModel.from_pretrained(403 spec["path"],404 quantization_config=config,405 device_map={"": 0},406 trust_remote_code=True,407 ).eval()408 torch.cuda.synchronize()409 load_seconds = time.perf_counter() - load_started410 allocation_after_load = int(torch.cuda.memory_allocated())411 412 encode_started = time.perf_counter()413 queries = np.stack([414 encode_with(family_model, text, family_tokenizer) for text in query_texts415 ])416 documents = np.stack([417 encode_with(family_model, text, family_tokenizer) for text in document_texts418 ])419 torch.cuda.synchronize()420 encode_seconds = time.perf_counter() - encode_started421 metrics, scores, ranks = retrieval_metrics(queries, documents)422 metrics.update({423 "lane": f"cuda-zerogpu-{variant}",424 "family": family,425 "model": spec["source"],426 "source_revision": spec["revision"],427 "quantizer": variant,428 "quantization_config": config.to_dict(),429 "load_seconds": load_seconds,430 "encode_seconds": encode_seconds,431 "texts_per_second": len(query_texts + document_texts) / encode_seconds,432 "torch_version": torch.__version__,433 "cuda_device": torch.cuda.get_device_name(0),434 "cuda_allocation_before_load": allocation_before,435 "cuda_allocation_after_load": allocation_after_load,436 "cuda_incremental_model_allocation": allocation_after_load - allocation_before,437 "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),438 })439 output_dir = DATA_ROOT / "cloud-results" / family440 result = {441 "metrics": metrics,442 "cuda_bf16_comparison": compare_vectors(443 queries,444 documents,445 scores,446 ranks,447 output_dir / "cuda-bf16.npz",448 "vs_cuda_bf16",449 ),450 "mlx_bf16_comparison": compare_vectors(451 queries,452 documents,453 scores,454 ranks,455 DATA_ROOT / "local-results/full-q4-q6-q8" / family / "quality/bf16.npz",456 "vs_mlx_bf16",457 ),458 }459 output_dir.mkdir(parents=True, exist_ok=True)460 stem = f"cuda-{variant}"461 np.savez_compressed(462 output_dir / f"{stem}.npz",463 queries=queries,464 documents=documents,465 scores=scores,466 ranks=ranks,467 metrics=np.array(json.dumps(metrics)),468 )469 (output_dir / f"{stem}.json").write_text(json.dumps(result, indent=2) + "\n")470 471 del family_model, family_tokenizer472 gc.collect()473 torch.cuda.empty_cache()474 return result475 476 477@spaces.GPU(duration=300)478def run_family_quantized_control(family: str, variant: str) -> dict:479 try:480 return _run_family_quantized_control(family, variant)481 except Exception as error:482 return {483 "diagnostic_error": type(error).__name__,484 "diagnostic_message": str(error),485 "diagnostic_traceback": traceback.format_exc(),486 }487 488 489with gr.Blocks() as demo:490 gr.Markdown(491 "# Embedding Quantization CUDA Control\n"492 "Runs one bounded 48-text BF16 control and writes the result to the attached private bucket. "493 "The requested GPU duration is capped at five minutes."494 )495 run_button = gr.Button("Run 0.6B BF16 CUDA control", variant="primary")496 output = gr.JSON(label="Result")497 run_button.click(498 fn=run_bf16_control,499 outputs=output,500 concurrency_limit=1,501 api_name="run_bf16_control",502 )503 gr.Markdown(504 "## CUDA-native quantizer controls\n"505 "These are bitsandbytes INT8/NF4 controls, not MLX Q/oQ/oQe replicas."506 )507 quantizer = gr.Dropdown(508 choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"509 )510 quant_button = gr.Button("Run bounded CUDA quantizer control")511 quant_output = gr.JSON(label="Quantized result")512 quant_button.click(513 fn=run_quantized_control,514 inputs=quantizer,515 outputs=quant_output,516 concurrency_limit=1,517 api_name="run_quantized_control",518 )519 gr.Markdown("## Larger-family BF16 cross-runtime controls")520 family = gr.Dropdown(521 choices=list(BF16_FAMILIES),522 value="gte-qwen2-1.5b",523 label="Model family",524 )525 family_button = gr.Button("Run bounded family BF16 control")526 family_output = gr.JSON(label="Family result")527 family_button.click(528 fn=run_family_bf16_control,529 inputs=family,530 outputs=family_output,531 concurrency_limit=1,532 api_name="run_family_bf16_control",533 )534 gr.Markdown("## Larger-family CUDA-native quantizer controls")535 quant_family = gr.Dropdown(536 choices=list(BF16_FAMILIES),537 value="qwen3-embedding-8b",538 label="Model family",539 )540 family_quantizer = gr.Dropdown(541 choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"542 )543 family_quant_button = gr.Button("Run bounded family quantizer control")544 family_quant_output = gr.JSON(label="Family quantized result")545 family_quant_button.click(546 fn=run_family_quantized_control,547 inputs=[quant_family, family_quantizer],548 outputs=family_quant_output,549 concurrency_limit=1,550 api_name="run_family_quantized_control",551 )552 553 554if __name__ == "__main__":555 demo.queue(default_concurrency_limit=1).launch()556