SketchSSM calibration: Nemotron Nano 9B v2 (FP8 weights)
calibration.pt is a portable SketchSSM
calibration file for nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8.
It contains the group-shared sketch basis and the Full-Gram allocation scores
from which the per-head rank table and the ordered frames for any mean rank are
derived, and the model layer index of each calibrated recurrent layer. It
contains no model weights.
Calibration weights
Collected with FP8 weights: nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8
at revision 8bc5eece2eb5514c4bca7f2ec655b91eb554f4c0. The basis and the allocation scores depend on the weights,
so use this file with these weights; for another precision or checkpoint,
calibrate with that checkpoint.
Contents
| Field | Value |
|---|---|
| Base model | Nemotron Nano 9B v2 (Mamba-2) |
| Recurrent layers | 27 (model layers 0, 2, 4, 6, 7, 9, 11, 13, 16, 18, 20, 23, 25, 27, 29, 32, 34, 36, 38, 41, 43, 44, 46, 48, 50, 52, 54) |
| State heads per layer | 128 |
| Key dim K / value dim V | 128 / 80 |
| Basis groups per layer | 8 |
| Window W | 16 |
| Erase factor | no |
| Allocation rank cap | 49 |
Basis omega |
float32, shape (27, 8, 79, 128) |
| File size | 10,952,133 bytes |
The file loads with torch.load(..., weights_only=True). The
calibration guide
documents its keys.
Verified mean ranks
For these mean ranks, the derived tables equal the bundle tables and the exported frames equal those of the bundle export path:
| Mean rank | Dense heads | Sketch heads |
|---|---|---|
| 2 | 0 | 3456 |
| 4 | 12 | 3444 |
| 8 | 103 | 3353 |
| 16 | 536 | 2920 |
Other mean ranks are allocated with the same rule. manifest.json lists the
SHA-256 of calibration.pt and these results.
Usage
With vLLM's SketchSSM support (--sketchssm), the ranks for the requested mean
rank are allocated when the model loads:
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2-FP8 \
--sketchssm SketchSSM/Nemotron-Nano-9B-v2-FP8 --sketchssm-mean-rank 8 \
--mamba-ssm-cache-dtype float32
With the SketchSSM package, export the frames for a mean rank:
hf download SketchSSM/Nemotron-Nano-9B-v2-FP8 calibration.pt --local-dir calibration
python -m sketchssm.calibration export --calibration calibration/calibration.pt \
--mean-rank 8 --out frames.pt
How it was produced
With the public calibration pipeline of the SketchSSM repository (release
0.1.15) and collect.yaml from this repository:
python -m sketchssm.calibration calibrate --config collect.yaml --out out
python -m sketchssm.calibration package --bundle out --out calibration.pt
python -m sketchssm.calibration manifest --bundle out --calibration calibration.pt --out manifest.json
package re-allocates every configured mean rank from the packaged contents
and fails unless each table equals the bundle table.
License
This calibration file is released under the Apache License 2.0, like the SketchSSM repository. It is derived from the base model's weights, so use it under the base model's license as well.