Info
Last Execution: 2026-09-03
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Quantization¶
A checkpoint that does not fit on your GPU can be held in fewer bits per weight.
Doing that through transformers normally means picking a quantizer backend and
building a config object, which is a lot of ceremony for a choice that amounts to
how wide is a weight. In nnsight the format goes in the dtype slot:
TransformersModel(repo_id, dtype="nf4") # where you would write "bfloat16"
This page runs every claim it makes on meta-llama/Llama-3.2-1B and
openai-community/gpt2, on one RTX A6000. It needs bitsandbytes and
accelerate, neither of which is installed with nnsight:
pip install bitsandbytes accelerate
import gc
import torch
from nnsight.modeling.transformers import TransformersModel
REPO = "meta-llama/Llama-3.2-1B"
PROMPT = "The Eiffel Tower is in the city of"
def free():
gc.collect()
torch.cuda.empty_cache()
print("compute capability:", torch.cuda.get_device_capability())
compute capability: (8, 6)
Naming a Format¶
dtype= takes a quantization name wherever it would take a torch dtype. Nothing
else about the constructor changes, and there is nothing to import.
free()
before = torch.cuda.memory_allocated()
model = TransformersModel(REPO, task="text-generation", dtype="nf4",
dispatch=True, device_map="cuda")
nf4_weights = (torch.cuda.memory_allocated() - before) / 1024**3
linear = model.model.layers[0].mlp.gate_proj._module
print("linear class:", type(linear).__name__)
print(f"weights on GPU: {nf4_weights:.2f} GB")
assert type(linear).__name__ == "Linear4bit"
linear class: Linear4bit weights on GPU: 1.00 GB
The names, and what each one loads:
| Name | What you get | Bytes/weight |
|---|---|---|
nf4, int4, 4bit |
bitsandbytes 4-bit, NF4 | 0.5 |
fp4 |
bitsandbytes 4-bit, FP4 | 0.5 |
int8, 8bit |
bitsandbytes LLM.int8() | 1.0 |
fp8 |
transformers block-wise FP8, compute capability 8.9+ | 1.0 |
Three names for 4-bit on purpose: people reach for whichever of int4, 4bit
and nf4 they last read about, and two of the three being an error would buy
nothing. All three load NF4, which is the format bitsandbytes recommends. fp4
is reached only by asking for it, and the accuracy table below shows why.
Your own quantization_config=BitsAndBytesConfig(...) still works. Passing it
with a quantization name raises, since the two can disagree about how the
weights are held.
fp8 Is Refused Below Compute Capability 8.9¶
fp8 is transformers' own quantizer rather than bitsandbytes, and it needs a
4090, an L40S, an H100 or later — an A100 is 8.0 and does not qualify. nnsight
refuses the load with a ValueError before any weights are fetched.
The refusal is there because transformers itself does not raise on older
hardware: it logs a warning, sets dequantize on the quantization config, and
loads bfloat16 at twice the width you budgeted, with the quantizer object still
attached — so a model that inspects hf_quantizer is told it is quantized when
it is not. This card is compute capability 8.6, so the refusal is what we can
show:
try:
TransformersModel(REPO, task="text-generation", dtype="fp8",
dispatch=True, device_map="cuda")
except ValueError as e:
print(type(e).__name__ + ":", e)
refused = True
assert refused and torch.cuda.get_device_capability() < (8, 9)
ValueError: dtype="fp8" needs a GPU of compute capability 8.9 or better (4090, L40S, H100 and later); on this machine transformers would silently dequantize and load bfloat16 at twice the width asked for. Use dtype="int8" or dtype="nf4" here instead.
config.quantization_config.dequantize is what tells the two cases apart. True
means the weights are bfloat16 whatever the name asked for.
Your Intervention Code Does Not Change¶
The module tree is identical. A quantized linear is a different class holding a differently shaped weight, but it sits at the same path with the same children, so every module reference, envoy and remote request lines up. Activations are ordinary 16-bit tensors of the usual shape.
with model.trace(PROMPT):
gate = model.model.layers[5].mlp.gate_proj.output.save()
logits = model.lm_head.output.save()
print("gate:", tuple(gate.shape), gate.dtype)
print("next token:", repr(model.tokenizer.decode(logits[0, -1].argmax())))
assert gate.shape[-1] == 8192 and gate.dtype == torch.bfloat16
assert model.tokenizer.decode(logits[0, -1].argmax()) == " Paris"
[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer TokenizersBackend. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.
gate: (1, 11, 8192) torch.bfloat16 next token: ' Paris'
Raw weights are the exception. A 4-bit weight really is stored packed, so
reading one gives a uint8 blob rather than the matrix. Read activations, or
load unquantized when the weights themselves are the object of study.
weight = model.model.layers[0].mlp.gate_proj._module.weight
n_params = sum(p.numel() for p in model._module.parameters())
print("stored weight:", tuple(weight.shape), weight.dtype)
print(f"sum(p.numel()): {n_params:,}")
assert weight.dtype == torch.uint8 and weight.shape == (8192 * 2048 // 2, 1)
del model
free()
stored weight: (8388608, 1) torch.uint8 sum(p.numel()): 749,275,136
That parameter count is the storage, not the model: an unquantized Llama-3.2-1B
reports 1,235,814,400 parameters where the 4-bit one reports 749,275,136, because
two weights share a byte. Anything that sizes a model by sum(p.numel()) gets a
number that is not its parameter count.
What It Costs¶
Three things move when you quantize: memory, accuracy, and speed. The loop below loads Llama-3.2-1B four times and measures all three against the bfloat16 run.
Accuracy here is KL between the two models' next-token distributions and
top-1 agreement, how often they pick the same next token, over the 86
positions of a fixed passage. A hidden-state norm, which is the easy thing to
report, moves by about 3% at nf4 and tells you nothing about either.
import math
import time
PASSAGE = (
"The Eiffel Tower, a wrought-iron lattice tower on the Champ de Mars in Paris, "
"was designed by Gustave Eiffel and completed in 1889 for the World's Fair. "
"It stands 330 metres tall and was the tallest structure in the world until the "
"Chrysler Building was finished in New York in 1930. Visitors reach the summit by "
"lift or by climbing more than 1,600 steps."
)
def measure(dtype):
free()
before = torch.cuda.memory_allocated()
m = TransformersModel(REPO, task="text-generation", dtype=dtype,
dispatch=True, device_map="cuda")
weights = (torch.cuda.memory_allocated() - before) / 1024**3
ids = m.tokenizer(PASSAGE, return_tensors="pt").input_ids
with m.trace(PASSAGE):
out = m.output.logits.save()
logprobs = torch.log_softmax(out[0, :-1].float(), dim=-1).cpu()
for _ in range(3): # warm the kernels before timing
with m.trace(PROMPT):
pass
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(20):
with m.trace(PROMPT):
pass
torch.cuda.synchronize()
forward_ms = (time.perf_counter() - start) / 20 * 1000
del m
free()
return weights, logprobs, forward_ms
results = {}
for name in ["bfloat16", "int8", "nf4", "fp4"]:
results[name] = measure(name)
print(f"{name:9} loaded")
bfloat16 loaded
int8 loaded
nf4 loaded
fp4 loaded
reference = results["bfloat16"][1]
print(f"{'dtype':9} {'weights':>9} {'vs bf16':>8} {'mean KL':>9} {'max KL':>7} "
f"{'top-1':>7} {'forward':>9}")
for name, (weights, logprobs, forward_ms) in results.items():
ratio = weights / results["bfloat16"][0]
kl = (reference.exp() * (reference - logprobs)).sum(-1)
agree = (reference.argmax(-1) == logprobs.argmax(-1)).float().mean()
shown = (" - - -" if name == "bfloat16"
else f"{kl.mean():9.3f} {kl.max():7.2f} {agree:6.1%}")
print(f"{name:9} {weights:6.2f} GB {ratio:7.2f}x {shown} {forward_ms:7.1f} ms")
if name == "nf4":
assert 0.4 < ratio < 0.5, "4-bit should be well under half of bfloat16"
assert 0.80 < agree < 0.95, f"nf4 top-1 agreement was {agree:.3f}"
if name == "int8":
assert agree > 0.90, f"int8 top-1 agreement was {agree:.3f}"
assert forward_ms > 2 * results["bfloat16"][2], "int8 is the slow format"
if name == "fp4":
nf4_agree = (reference.argmax(-1) == results["nf4"][1].argmax(-1)).float().mean()
assert agree < nf4_agree, "fp4 should agree less often than nf4"
dtype weights vs bf16 mean KL max KL top-1 forward bfloat16 2.30 GB 1.00x - - - 13.7 ms int8 1.40 GB 0.61x 0.011 0.09 95.3% 50.8 ms nf4 1.00 GB 0.43x 0.143 1.64 87.1% 19.8 ms fp4 1.00 GB 0.43x 0.182 1.37 82.4% 19.6 ms
4-bit changes the argmax next token more than one time in ten. Treat a
quantized run as a different model rather than a cheaper copy of the same one:
do not compare activations across widths, and do not report a 4-bit result as
though it were the checkpoint's. fp4 disagrees more often than nf4 at exactly
the same size, which is why the unqualified names point at NF4.
The damage shrinks as the model grows. The same measurement on Llama-3.1-8B
(14.96 GB in bfloat16) gives nf4 5.65 GB at mean KL 0.049 and 92.9% agreement,
against 87.1% on the 1B.
int8 is the accurate format and the slow one at once. Unpacking weights on the
way into each matmul costs a little at 4-bit and a lot at 8-bit, and for a sweep
over hundreds of forwards that is usually what decides between them. The card
here is shared, so the exact ratio moves: it ranged from 2.9x to 4.5x over four
runs, which is why the check in the cell above only asks for 2x.
Memory saves less than the arithmetic predicts, because the format leaves
embeddings, norms and the LM head in 16 bits and stores a scale per block. At 0.5
bytes per weight, nf4 on this model "should" be 0.58 GB. Counting the
embeddings at 2 bytes and the rest at the format's width predicts 0.94, which is
close enough to budget from. Better still, measure.
The Compute Dtype¶
Everything the format leaves alone, and everything the model computes in, is
bfloat16. The exception is int8: bitsandbytes implements LLM.int8() in
float16 and casts anything else on the way in, warning once per matmul as it
does. So an int8 model's activations arrive as float16.
compute_dtype= overrides it (bnb_4bit_compute_dtype= is accepted as a synonym
for anyone arriving from the bitsandbytes documentation).
For a float32 checkpoint this is the sentence that bites: quantizing moves the whole model to the compute dtype, so activations that were float32 come back at half the width. GPT-2 is float32.
for dtype in [None, "nf4", "int8"]:
free()
kwargs = {} if dtype is None else {"dtype": dtype}
gpt2 = TransformersModel("openai-community/gpt2", task="text-generation",
dispatch=True, device_map="cuda", **kwargs)
with gpt2.trace(PROMPT):
attn = gpt2.transformer.h[5].attn.c_attn.output.save()
stored = gpt2.transformer.h[0].attn.c_attn._module.weight
print(f"{str(dtype):9} activation {str(attn.dtype):16} stored weight {stored.dtype}")
if dtype is None:
assert attn.dtype == torch.float32
if dtype == "nf4":
assert attn.dtype == torch.bfloat16
if dtype == "int8":
assert attn.dtype == torch.float16
del gpt2
free()
None activation torch.float32 stored weight torch.float32
nf4 activation torch.bfloat16 stored weight torch.uint8
int8 activation torch.float16 stored weight torch.int8
Gradients Through an int8 Model Can Overflow¶
float16 carries a much narrower exponent than bfloat16, and a backward pass over a large-magnitude loss runs out of it. Nothing raises: the NaNs propagate into whatever you compute from the gradient, which for attribution patching is the attribution score itself.
def grad_norm(model, loss_fn):
with model.trace(PROMPT):
activation = model.transformer.h[0].output
loss = loss_fn(model)
with loss.backward():
grad = activation.grad.save()
return grad.dtype, float(grad.float().norm())
for dtype in [None, "nf4", "int8"]:
free()
kwargs = {} if dtype is None else {"dtype": dtype}
gpt2 = TransformersModel("openai-community/gpt2", task="text-generation",
dispatch=True, device_map="cuda", **kwargs)
summed = grad_norm(gpt2, lambda m: m.output.logits.sum())
normalized = grad_norm(gpt2, lambda m: torch.log_softmax(m.output.logits[0, -1], -1).max())
print(f"{str(dtype):9} logits.sum() -> {summed[0]} {summed[1]:12.2f} "
f"log-prob -> {normalized[0]} {normalized[1]:.4f}")
if dtype == "int8":
assert math.isnan(summed[1]) # the sum overflows float16
assert normalized[1] > 0 # a normalized loss does not
else:
assert summed[1] > 0
del gpt2
free()
None logits.sum() -> torch.float32 934879.56 log-prob -> torch.float32 0.5625
nf4 logits.sum() -> torch.bfloat16 1079000.00 log-prob -> torch.bfloat16 0.5912
int8 logits.sum() -> torch.float16 nan log-prob -> torch.float16 0.5209
Use a loss on a scale float16 can hold, such as a log-probability, or use nf4,
which computes in bfloat16 and survives either loss.
What Quantization Will Not Do¶
bitsandbytes swaps nn.Linear and nothing else, which decides what a given
checkpoint saves.
A mixture-of-experts model barely shrinks. transformers 5 holds the experts
as stacked 3-D parameters on one module rather than as linears, so bitsandbytes
leaves them at the compute dtype and quantizes only the attention projections,
the router and the shared layers. Those are the minority of an MoE's weights.
from collections import Counter
MOE = "yujiepan/qwen3-moe-tiny-random"
sizes = {}
for dtype in [None, "nf4"]:
free()
before = torch.cuda.memory_allocated()
kwargs = {} if dtype is None else {"dtype": dtype}
moe = TransformersModel(MOE, task="text-generation", dispatch=True,
device_map="cuda", **kwargs)
sizes[dtype] = (torch.cuda.memory_allocated() - before) / 1024**2
kinds = Counter(type(m).__name__ for m in moe._module.modules()
if "Linear" in type(m).__name__ or "Experts" in type(m).__name__)
expert = dict(moe._module.named_parameters())["model.layers.1.mlp.experts.gate_up_proj"]
print(f"{str(dtype):5} {sizes[dtype]:6.2f} MB {dict(kinds)}")
print(f" experts.gate_up_proj {tuple(expert.shape)} {expert.dtype}")
assert expert.dtype == torch.bfloat16 # never quantized, either way
del moe
free()
print(f"\nsaving: {1 - sizes['nf4'] / sizes[None]:.1%}")
assert sizes["nf4"] / sizes[None] > 0.9 # under 10% saved
None 19.02 MB {'Linear': 12, 'Qwen3MoeExperts': 1}
experts.gate_up_proj (8, 256, 64) torch.bfloat16
nf4 18.96 MB {'Linear4bit': 11, 'Qwen3MoeExperts': 1, 'Linear': 1}
experts.gate_up_proj (8, 256, 64) torch.bfloat16
saving: 0.3%
Quantizing Qwen1.5-MoE-A2.7B at nf4 takes it from 12.89 GiB to 12.53 GiB, under
3%, while still perturbing every routing decision the attention output feeds. For
an MoE that does not fit, tensor parallelism is the answer; see the sharding
section of Loading a Model.
A narrow torch dtype that nothing can load falls back to float32. torch.int1
through torch.int7 exist, so dtype="int3" is not rejected as a name.
transformers tries it, fails, and loads float32 instead — twice the width of
the default, in answer to a request for something narrower. Nothing raises.
free()
before = torch.cuda.memory_allocated()
int3 = TransformersModel(REPO, task="text-generation", dtype="int3",
dispatch=True, device_map="cuda")
int3_weights = (torch.cuda.memory_allocated() - before) / 1024**3
print("config.dtype: ", int3._module.config.dtype)
print(f"weights on GPU: {int3_weights:.2f} GB "
f"({int3_weights / results['bfloat16'][0]:.1f}x the bfloat16 model)")
assert int3._module.config.dtype == torch.float32
assert int3_weights > results["bfloat16"][0]
del int3
free()
[transformers] Falling back to torch.float32 because loading with the original dtype failed on the target device.
config.dtype: torch.float32 weights on GPU: 4.61 GB (2.0x the bfloat16 model)
load_in_4bit= and load_in_8bit= are not transformers 5 arguments. They
are what tutorials written against transformers 4 pass, and what an LLM asked
for 4-bit loading will usually write. In transformers 5 they reach the model
class as a stray keyword, and the pipeline reports it well down a nested
traceback.
try:
TransformersModel(REPO, task="text-generation", load_in_4bit=True,
dispatch=True, device_map="cuda")
except ValueError as error:
lines = [line for line in str(error).splitlines() if "load_in_4bit" in line]
print(str(error).splitlines()[0][:90], "...")
print(lines[0].strip())
assert "unexpected keyword argument 'load_in_4bit'" in lines[0]
free()
Could not load model meta-llama/Llama-3.2-1B with any of the following classes: (<class 't ... TypeError: LlamaForCausalLM.__init__() got an unexpected keyword argument 'load_in_4bit'
Use dtype="nf4" or dtype="int8".
Requirements¶
pip install bitsandbytes accelerate. Neither is a dependency of nnsight, and
transformers raises ImportError naming whichever is missing once the load
reaches the quantizer.
A GPU is not required: bitsandbytes 0.50 quantizes and runs on CPU, reaching the
same layer-5 norm as the GPU to five significant figures. What the quantizers do
reject is the meta device, which is why dispatch=False still works — the meta
build ignores the quantization and builds the architecture at the compute dtype.
That is what makes the lazy path work, and what lets a client model a checkpoint
a server holds quantized. Weights are quantized only when they are actually
loaded.
On NDIF the same names configure a deployment, so a server can hold a model 4-bit with nothing client-side changing. A client cannot request it: a remote model key is the repo id and revision, and says nothing about how the weights are held.