llama-3.2-nv-embedqa-1b-v2-onnx

NVIDIA Llama Nemotron Embedding 1B v2 — ONNX

ONNX export of nvidia/llama-3.2-nv-embedqa-1b-v2 for use with fastembed-rs and onnxruntime.

Model details

Files

• model.onnx — ONNX graph definition
• model.onnx.data — External weight data (~4.6 GB, float32)
• tokenizer.json — Tokenizer (Llama 3.2 BPE)

Usage with onnxruntime

import onnxruntime as ort
import numpy as np
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("kzzalews/llama-3.2-nv-embedqa-1b-v2-onnx")
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

def embed(texts: list[str]) -> np.ndarray:
    enc = tokenizer(texts, return_tensors="np", padding=True, truncation=True)
    hidden = sess.run(["last_hidden_state"], {
        "input_ids": enc["input_ids"],
        "attention_mask": enc["attention_mask"],
    })[0]
    # mean pooling
    mask = enc["attention_mask"][:, :, np.newaxis].astype(np.float32)
    return (hidden * mask).sum(axis=1) / mask.sum(axis=1)

# query prefix required for retrieval
embeddings = embed(["query: How does Kubernetes handle OOMKill events?"])
print(embeddings.shape)  # (1, 2048)

fastembed-rs integration (pending PR)

Two model variants for fastembed-rs:

float32 variant:

// src/models/text_embedding.rs
ModelInfo {
    model: EmbeddingModel::NvidiaNemotronEmbedQA1BV2,
    dim: 2048,
    description: String::from("NVIDIA Llama Nemotron Embedding 1B v2 for retrieval"),
    model_code: String::from("kzzalews/llama-3.2-nv-embedqa-1b-v2-onnx"),
    model_file: String::from("model.onnx"),
    additional_files: vec!["model.onnx.data".to_string()],
    output_key: None,
},

// src/text_embedding/impl.rs
EmbeddingModel::NvidiaNemotronEmbedQA1BV2 => Some(Pooling::Mean),

float16 variant (2.8 GB, identical quality, same speed on CPU):

ModelInfo {
    model: EmbeddingModel::NvidiaNemotronEmbedQA1BV2Fp16,
    dim: 2048,
    description: String::from("NVIDIA Llama Nemotron Embedding 1B v2 (float16)"),
    model_code: String::from("kzzalews/llama-3.2-nv-embedqa-1b-v2-onnx"),
    model_file: String::from("fp16/model.onnx"),
    additional_files: vec!["fp16/model.onnx.data".to_string()],
    output_key: None,
},

// src/text_embedding/impl.rs
EmbeddingModel::NvidiaNemotronEmbedQA1BV2Fp16 => Some(Pooling::Mean),

Variants

Note: INT8 dynamic quantization was tested and rejected (cos_sim = 0.65–0.70 across all configs without calibration data — unusable for retrieval). Float16 saves space but does not improve CPU inference speed (CPU emulates fp16 without AVX-512 FP16). Speed benefit requires GPU or AVX-512 FP16 hardware.

Performance (CPU: Intel i9-13950HX, 16 threads, no GPU)

Numerical validation (ONNX vs PyTorch): max diff = 1.76e-05, mean diff = 2.49e-06.

Prompts

This model uses instruction prefixes for retrieval:

• Queries: "query: <text>"
• Documents: "passage: <text>"