CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-08-04 14:56:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

x/create: quantize lm_head at 8-bit in the requested family

Browse source 
The lm_head rule was asymmetric: the fp modes kept an untied head at
source precision (even under mxfp8, leaving it the only bf16 matmul in
the model), while int4 quantized it at 4 bits with no promotion. The
tied-embedding overrides (gemma4, cohere2moe) already resolve the head
to the 8-bit family type and hold quality close to bf16.

Apply the same decision to untied heads: the 8-bit type in the
requested family when it fits the shape, source precision otherwise.
int4 now promotes the head to int8, and the fp modes quantize it to
mxfp8 instead of keeping bf16.
This commit is contained in:
Jesse Gross 2026-07-23 12:39:53 -07:00
parent fce745fe5e
commit 83d4311ffe
2 changed files with 10 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
x/create/create.go
View file

@ -295,8 +295,12 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string {
return ""
}
// lm_head is too sensitive for the fp quant modes; keep it at source precision.
if strings.HasSuffix(name, "lm_head.weight") && (quantNorm == "nvfp4" || quantNorm == "mxfp4" || quantNorm == "mxfp8") {
// lm_head is too sensitive for 4-bit types; the 8-bit type in the requested
// family keeps quality close to bf16 while saving decode bandwidth.
if strings.HasSuffix(name, "lm_head.weight") {
if e := eightBit(quantNorm); isAligned(shape, e) {
return e
}
return ""
}

8
x/create/create_test.go
View file

@ -598,10 +598,10 @@ func TestGetTensorQuantization_MixedPrecisionPromotion(t *testing.T) {
// int8: already 8-bit, no promotion
{"v_proj int8 stays", "model.layers.0.self_attn.v_proj.weight", aligned, "int8", "int8"},
// lm_head stays at source precision for fp modes, quantizes for affine
{"lm_head nvfp4 kept", "lm_head.weight", aligned, "nvfp4", ""},
{"lm_head mxfp8 kept", "lm_head.weight", aligned, "mxfp8", ""},
{"lm_head int4 stays", "lm_head.weight", aligned, "int4", "int4"},
// lm_head resolves to the 8-bit type in the requested family
{"lm_head nvfp4 to mxfp8", "lm_head.weight", aligned, "nvfp4", "mxfp8"},
{"lm_head mxfp8 uniform", "lm_head.weight", aligned, "mxfp8", "mxfp8"},
{"lm_head int4 promoted", "lm_head.weight", aligned, "int4", "int8"},
// Expert tensors: down_proj also promoted for int4
{"expert down_proj int4", "model.layers.0.mlp.experts.down_proj.weight", []int32{128, 4096, 2816}, "int4", "int8"},