SketchSSM calibration: Nemotron Nano 9B v2 (BF16 weights)
calibration.pt is a portable SketchSSM
calibration file for nvidia/NVIDIA-Nemotron-Nano-9B-v2. 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. It contains no model weights.
Calibration weights
Collected with BF16 weights: nvidia/NVIDIA-Nemotron-Nano-9B-v2 at revision 6533e8de2c68e4536bf7c411d7a3ce5734111476 (the original BF16 checkpoint). 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 |
| 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,325 bytes |
Layers are stored in the order of the model's recurrent layers; layer_ids lists the model layer index of each. 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 published bundle tables and the exported frames equal those of the bundle export path:
| Mean rank | Dense heads | Sketch heads |
|---|---|---|
| 4 | 12 | 3444 |
| 6 | 30 | 3426 |
| 10 | 230 | 3226 |
| 21 | 775 | 2681 |
Other mean ranks are allocated with the same rule but have no stored table to
compare with. manifest.json lists the SHA-256 of calibration.pt and these
results.
Usage
With the SketchSSM repository, export the frames for a mean rank:
hf download SketchSSM/Nemotron-Nano-9B-v2-BF16 calibration.pt --local-dir calibration
python -m sketchssm.calibration export --calibration calibration/calibration.pt \
--mean-rank 10 --out frames.pt
With the vLLM included in the SketchSSM repository (see its README for installation):
vllm serve nvidia/NVIDIA-Nemotron-Nano-9B-v2 --sketchssm SketchSSM/Nemotron-Nano-9B-v2-BF16 --sketchssm-mean-rank 10 \
--trust-remote-code --mamba-ssm-cache-dtype float32 --no-enable-prefix-caching
How it was produced
From the public calibration bundle sketchssm/calibration/example/nemotron_nano in the
SketchSSM repository:
python -m sketchssm.calibration package --bundle sketchssm/calibration/example/nemotron_nano --out calibration.pt
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.