WeMM-Embedding-2B-FP8

WeMM-Embedding-2B — FP8 Quantization

FP8 (8-bit float, E4M3) quantization of tencent/WeMM-Embedding-2B, intended for vLLM / SGLang-class backends with FP8 support (RTX 4090 / Ada and newer have native FP8 tensor cores).

This is a separate repo from the GGUF build — GGUF targets llama.cpp; this FP8 build targets GPU inference servers.

Files

• model.fp8.safetensors — weights stored as float8_e4m3fn (per-tensor scale).
• fp8_scales.json — per-layer dequant scale (layer name -> scalar).
• Supporting files (config, tokenizer, custom WeMMEmbedding modeling, chat templates) are mirrored from the base model so AutoModel can load it.

Note: the saved weights are raw fp8 + scale. To run, dequantize at load time (fp8 -> bf16) or serve through a backend that natively consumes fp8. See Usage below.

Evaluation (STS-B, same engine)

All numbers use the same engine (transformers / torch AutoModel) for both the BF16 baseline and the FP8 model, so Δρ is pure FP8 rounding error.

Metrics:

• STS-B Spearman ρ: rank correlation between model cosine similarities and human similarity scores (0-5). Higher is better.
• Δρ vs BF16: relative drop of ρ against the BF16 baseline. Negative = scored slightly above baseline (within noise).
• Emb Cosine vs BF16: mean cosine between each sentence embedding and its BF16 counterpart (space fidelity). 1.0 = identical.

Conclusion: FP8 causes negligible quality loss (Δρ = +0.12%, Emb Cosine = 0.999) on STS-B. This is the recommended format when serving on FP8-capable GPUs.

Usage

Dequantize to bf16 at load time (transformers)

import torch, json, safetensors.torch as st
from transformers import AutoModel, AutoProcessor

sd = st.load_file("model.fp8.safetensors")
scales = json.load(open("fp8_scales.json"))
for k, s in scales.items():          # k like 'model.xxx.weight'
    sd[k] = (sd[k].to(torch.float32) * s).to(torch.bfloat16)   # dequant
# save a runnable bf16 copy, or load directly:
model = AutoModel.from_pretrained(".", trust_remote_code=True, dtype=torch.bfloat16)

vLLM / SGLang

These backends expect a compressed-tensors / native fp8 checkpoint. The raw fp8+scale layout here is not yet wrapped for direct vLLM loading; repackaging into the backend's fp8 format (or quantizing the base model with the backend's own fp8 path) is required before serving. The quality numbers above already demonstrate the FP8 format itself is near-lossless.

Notes

• FP8 was produced by manual per-tensor fp8 (E4M3) quantization of all 285 Linear layers (including the vision tower). llmcompressor oneshot did not actually quantize the qwen3_5 custom Linear layers, so the manual path is used for the reported numbers.
• The GGUF variants (Q8_0 / Q6_K / Q5_K_M / IQ4_XS / IQ3_M) live in a separate repo.