| from __future__ import annotations |
|
|
| import json |
| import os |
| import time |
| import gc |
| from pathlib import Path |
|
|
| import gradio as gr |
| import numpy as np |
| import spaces |
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig |
|
|
|
|
| MODEL_PATH = Path("/models/qwen3-embedding-0.6b") |
| 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" |
| ) |
|
|
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side="left") |
| model = AutoModel.from_pretrained( |
| MODEL_PATH, |
| torch_dtype=torch.bfloat16, |
| trust_remote_code=True, |
| ).to("cuda").eval() |
|
|
|
|
| def detailed_instruction(query: str) -> str: |
| return f"Instruct: {TASK}\nQuery:{query}" |
|
|
|
|
| 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 = 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) |
| ranks = np.array([ |
| int(np.where(order[index] == index)[0][0]) + 1 |
| for index in range(len(queries)) |
| ]) |
| positive = np.diag(scores) |
| negative = scores.copy() |
| np.fill_diagonal(negative, -np.inf) |
| margins = positive - negative.max(axis=1) |
| return { |
| "pair_count": len(queries), |
| "top1": float(np.mean(ranks == 1)), |
| "recall_at_5": float(np.mean(ranks <= 5)), |
| "mrr": float(np.mean(1.0 / ranks)), |
| "mean_margin": float(margins.mean()), |
| "minimum_margin": float(margins.min()), |
| "mean_rank": float(ranks.mean()), |
| "worst_rank": int(ranks.max()), |
| }, scores, ranks |
|
|
|
|
| @spaces.GPU(duration=300) |
| def run_bf16_control() -> dict: |
| 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] |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.reset_peak_memory_stats() |
| torch.cuda.synchronize() |
| started = time.perf_counter() |
| queries = np.stack([encode(text) for text in query_texts]) |
| documents = np.stack([encode(text) for text in document_texts]) |
| if torch.cuda.is_available(): |
| torch.cuda.synchronize() |
| elapsed = time.perf_counter() - started |
| metrics, scores, ranks = retrieval_metrics(queries, documents) |
| metrics.update({ |
| "lane": "cuda-zerogpu-bf16", |
| "model": "Qwen/Qwen3-Embedding-0.6B", |
| "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3", |
| "elapsed_seconds": elapsed, |
| "texts_per_second": len(query_texts + document_texts) / elapsed, |
| "torch_version": torch.__version__, |
| "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, |
| "cuda_peak_bytes": torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None, |
| }) |
|
|
| comparison = None |
| if MLX_REFERENCE.exists(): |
| reference = np.load(MLX_REFERENCE) |
| ref_all = np.concatenate([reference["queries"], reference["documents"]]) |
| cuda_all = np.concatenate([queries, documents]) |
| aligned = np.sum(ref_all * cuda_all, axis=1) |
| score_delta = scores - reference["scores"] |
| comparison = { |
| "mean_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.mean()), |
| "minimum_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.min()), |
| "score_rmse_cuda_vs_mlx_bf16": float(np.sqrt(np.mean(score_delta ** 2))), |
| "queries_with_rank_change": int(np.count_nonzero(ranks - reference["ranks"])), |
| } |
|
|
| result = {"metrics": metrics, "mlx_bf16_comparison": comparison} |
| OUTPUT_DIR.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed( |
| OUTPUT_DIR / "cuda-bf16.npz", |
| queries=queries, |
| documents=documents, |
| scores=scores, |
| ranks=ranks, |
| metrics=np.array(json.dumps(metrics)), |
| ) |
| (OUTPUT_DIR / "cuda-bf16.json").write_text(json.dumps(result, indent=2) + "\n") |
| 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", "bnb-nf4") |
|
|
|
|
| def quantization_config(variant: str) -> BitsAndBytesConfig: |
| if variant == "bnb-int8": |
| return BitsAndBytesConfig(load_in_8bit=True) |
| if variant == "bnb-nf4": |
| return BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=False, |
| ) |
| raise ValueError(f"unsupported quantizer: {variant}") |
|
|
|
|
| @spaces.GPU(duration=300) |
| def run_quantized_control(variant: str) -> dict: |
| if variant not in QUANTIZERS: |
| raise ValueError(f"unsupported quantizer: {variant}") |
| config = quantization_config(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=config, |
| 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": config.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" |
| "Runs one bounded 48-text BF16 control and writes the result to the attached private bucket. " |
| "The requested GPU duration is capped at five minutes." |
| ) |
| run_button = gr.Button("Run 0.6B BF16 CUDA control", variant="primary") |
| output = gr.JSON(label="Result") |
| run_button.click( |
| fn=run_bf16_control, |
| outputs=output, |
| 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__": |
| demo.queue(default_concurrency_limit=1).launch() |
|
|
