prism-quantized / README.md
jc-914's picture
Add files using upload-large-folder tool
2d2f883 verified
|
Raw History Blame Contribute Delete
8.03 kB
---
license: apache-2.0
base_model: sarvamai/prism
pipeline_tag: text-classification
language: [en, hi, bn, ta, te, mr, gu, kn, ml, or, pa]
tags: [decision-model, classification, calibration, int4, gptq, quantized, moe]
---
# Sarvam-Prism INT4 (≤16 GB GPUs)
An INT4 quantization of [sarvamai/prism](https://huggingface.co/sarvamai/prism), a decision model built on Sarvam-105B.
The original needs one FP8-capable GPU with ~120 GB of memory. This version runs on **any CUDA GPU with about 10 GB of
memory or more**, for example 16, 12 or 24 GB cards. It does not need FP8 tensor cores.
You give it a context and one or more questions, each with a fixed set of answers. It returns a probability for every
answer in one forward pass. The API is the same as the original's: `decide(state, questions)`.
## Requirements
- A CUDA GPU with **≥ 10 GB** of memory. sm_80+ (Ampere or newer) is recommended. sm_70/75 also work, on a slower path.
It was tested with GPU memory hard-capped at 15 GiB and at 11 GiB.
- **≥ 64 GB of host RAM.** About 54 GB of expert weights live in pinned host memory.
- Python 3.10+, PyTorch ≥ 2.4, Triton ≥ 3.0, `transformers==4.52.1` (see `requirements.txt`), and a C compiler
(Triton builds its kernel on first use).
- About 56 GB of disk.
## Quick start
```bash
hf download jc-914/prism-quantized --local-dir prism-int4
cd prism-int4
pip install -r requirements.txt
python example.py . # optional 2nd arg: GPU memory cap in GiB, e.g. `python example.py . 12`
```
```python
from prism_int4 import Prism
model = Prism.from_pretrained(".") # uses the free memory on cuda:0; or gpu_memory_gb=12
answers = model.decide(
state={"ticket": "My card was charged twice for the same order and I still haven't got a refund."},
questions={
"team": {"type": "choice", "instructions": "Which team should handle this ticket?",
"criteria": {"billing": "payments, charges and refunds", "tech": "app or website problems",
"delivery": "shipping and order status"}},
"urgent": {"type": "yes_no", "instructions": "Does this need a reply within the hour?"},
"anger": {"type": "score", "instructions": "How upset is the customer?",
"criteria": ["calm", "annoyed", "angry", "furious"]},
},
)
# Many requests at once (amortises the weight streaming, see Speed):
results = model.decide_batch([(state1, questions1), (state2, questions2), ...])
```
The request format (`choice` / `yes_no` / `score`) and the output format are the same as the
[original model card](https://huggingface.co/sarvamai/prism).
## What was quantized
| part | parameters | format here | vs. original |
|---|---|---|---|
| 128 routed experts × 31 MoE layers | ~100B | **INT4**, GPTQ, symmetric, group size 128, fp32 scales | was FP8 |
| attention, dense layer 0, shared experts | ~4B | FP8 (e4m3, one scale per output row), widened to bf16 at use | **unchanged** |
| embeddings, norms, MoE router, decision head | – | bf16 / fp32 | unchanged |
- **GPTQ calibration.** 262,533 tokens of requests in Prism's own input format (state, question, options and
answer, with the decision attention mask). The text is multilingual Wikipedia in all 11 supported languages, and
the requests cover all three question types. Each layer's MoE input was recorded from the original FP8 model.
Every expert's down projection is calibrated on the output of its *quantized* gate/up projection.
- **Rare experts.** Experts that were routed fewer than 512 calibration tokens (180 of 3,968) use round-to-nearest
with an MSE-optimal clip search instead of GPTQ, so their error feedback is not fitted to a handful of samples.
- **Layer error.** The Hessian-weighted relative output error, averaged over layers, is 0.0014 for gate/up and
0.0036 for down, against 0.0075 and 0.0109 for plain round-to-nearest. Per-layer numbers are in
`quantization/quant_stats.json`.
## Accuracy against the original FP8 model
The suite is 26 requests / 43 questions: English and all 10 Indic languages, all question types, and contexts from a
few tokens up to ~4k tokens. It is disjoint from the calibration data. The reference outputs come from the original
`prism.py` on an H200.
| metric | INT4 (15 GiB cap) | INT4 (11 GiB cap) |
|---|---|---|
| top-1 answer agreement | **43 / 43** | 43 / 43 |
| mean of the largest per-option probability difference | 0.012 | 0.012 |
| worst largest per-option probability difference | 0.111 | 0.111 |
| mean KL(original ‖ int4) | 0.0016 | 0.0016 |
| peak GPU memory | 12.3 GiB | 9.1 GiB |
The largest differences are on borderline questions. For example, a yes/no question at 0.53 / 0.47 in the original
comes out at 0.64 / 0.36 here. Clear-cut answers agree closely. This is a small suite: on any task where calibration
matters, validate against your own labelled data. Full outputs are in `quantization/eval_*.json`.
## How it runs on a small GPU
- Dense weights (~3.9 GiB) stay on the GPU.
- The INT4 experts (1.59 GiB per layer) stay in pinned host RAM. They are streamed to the GPU one layer at a time on
a side CUDA stream, double-buffered, so the next layer's copy overlaps the current layer's compute.
- As many layers as fit in `gpu_memory_gb − workspace_gb` stay resident. On a GPU with ~60 GB free, everything is
resident and nothing streams.
- The experts run in a Triton grouped W4A16 kernel: tensor-core dots with fp32 accumulation, using bf16 on sm_80+ and
fp16 on older GPUs. It compiles for sm_70 through sm_120.
- The 2 GB token-embedding table stays in host RAM.
`from_pretrained` knobs:
| knob | default | effect |
|---|---|---|
| `gpu_memory_gb` | all free memory | GPU memory the model may use |
| `workspace_gb` | 3.5 | memory kept for activations. Raise it for long contexts or big batches; lower it to keep more layers resident |
| `batch_size` | 16 | questions per forward pass |
| `max_batch_tokens` | 16384 | padded tokens per forward pass |
| `pin_memory` | `True` | set `False` if the host can't pin ~54 GB (copies become slower) |
## Speed
When layers stream, each forward pass moves the streamed expert weights over PCIe: ~50 GB with 2 resident layers,
53 GB with none. **Latency is therefore bounded by host-to-GPU bandwidth, not GPU compute.**
- Measured on an H200 NVL host (PCIe Gen5) with the GPU capped at 15 GiB and at 11 GiB: median **2.6 s per request**.
- Expected at roughly 25 GB/s (PCIe Gen4 x16): about 2–2.5 s per forward pass.
- Expected at roughly 12 GB/s (PCIe Gen4 x8, or Gen3 x16): about 4.5–5 s per forward pass.
A forward pass scores up to `batch_size` questions together, so put many questions, or many requests via
`decide_batch`, into one call. Throughput then scales almost linearly up to the batch limit. There are no CUDA graphs,
because weights move every pass. The first call also compiles the Triton kernel.
## Files
| file | contents |
|---|---|
| `model-layer-*.safetensors` | per layer: FP8 dense tensors (unchanged) plus `experts.{gate_up,down}_{qweight,scales}` (INT4) |
| `model-embed.safetensors`, `decision_head.safetensors`, `tokenizer*`, `config.json`, `*_sarvam_moe.py` | from the original release; `config.json` documents the quantization |
| `prism_int4.py`, `example.py`, `requirements.txt` | inference code |
| `quantization/` | the quantization script, calibration generator, per-layer stats and evaluation outputs |
INT4 packing: `qweight` is uint8 `[E, N, K/2]`. In each 128-wide group along K, byte *j* holds element *j* in its low
nibble and element *j+64* in its high nibble; value = nibble − 8. `scales` is fp32 `[E, N, K/128]`, and
weight = value × scale.
## License and credit
Apache 2.0, the same as the original. All model weights derive from
[sarvamai/prism](https://huggingface.co/sarvamai/prism) by Sarvam AI. This is an unofficial quantization: for the
original, its evaluation and the hosted API, see Sarvam's [release blog](https://www.sarvam.ai/blogs/sarvam-prism).