mlx: tighten up gemma4 moe loading code (#16964)
This change allows .experts.gate_proj / .up_proj / .down_proj tensor names to each be used for both quantized (i.e. nvfp4 and mxfp8) and non-quantized (bf16) models. Previous to this only non-quantized models used that tensor naming scheme.
This commit is contained in:
Patrick Devine
+57
-57
@@ -243,6 +243,46 @@ func collectExpertProjection(tensors map[string]*mlx.Array, cfg *TextConfig, pre
|
||||
return out
|
||||
}
|
||||
|
||||
// loadFusedExperts populates moe with a stacked-3D fused gate_up + down expert
|
||||
// pair, choosing the quantized GatherQMM path when scale sidecars are present
|
||||
// and the dense GatherMM path otherwise. gateUpW/downW are the stacked tensors
|
||||
// and gateUpKey/downKey are their map keys, used to locate the _scale/_qbias
|
||||
// sidecars and per-tensor quant metadata. Shared by the .experts./.moe.
|
||||
// pre-stacked names and the .moe.switch_mlp. names so all fused layouts load
|
||||
// identically.
|
||||
func loadFusedExperts(tensors map[string]*mlx.Array, cfg *TextConfig, moe *MoEBlock, gateUpW *mlx.Array, gateUpKey string, downW *mlx.Array, downKey string) {
|
||||
moe.UseFusedGateUp = true
|
||||
|
||||
gateUpScales := firstNonNil(tensors, gateUpKey+"_scale", gateUpKey+".scale")
|
||||
downScales := firstNonNil(tensors, downKey+"_scale", downKey+".scale")
|
||||
if gateUpScales == nil || downScales == nil {
|
||||
// Unquantized: keep fused and transpose for GatherMM ([experts, in, out]).
|
||||
moe.GateUpWeight = transposeForGatherMM(gateUpW)
|
||||
moe.DownWeight = transposeForGatherMM(downW)
|
||||
return
|
||||
}
|
||||
|
||||
// Quantized (affine/nvfp4/mxfp8): keep fused as a single tensor for one
|
||||
// GatherQMM call instead of two. No transpose. nvfp4/mxfp8 carry no qbias,
|
||||
// and this path does not apply nvfp4's optional double-scale global_scale.
|
||||
moe.UseQuantized = true
|
||||
moe.GateUpWeightQ = gateUpW
|
||||
moe.GateUpScales = gateUpScales
|
||||
moe.GateUpBiases = firstNonNil(tensors, gateUpKey+"_qbias", gateUpKey+".bias")
|
||||
moe.DownWeightQ = downW
|
||||
moe.DownScales = downScales
|
||||
moe.DownBiases = firstNonNil(tensors, downKey+"_qbias", downKey+".bias")
|
||||
|
||||
moe.GateUpGroupSize, moe.GateUpBits, moe.QuantMode = model.ResolveLinearQuantParams(
|
||||
cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode,
|
||||
cfg.TensorQuant, gateUpKey, gateUpW, gateUpScales,
|
||||
)
|
||||
moe.DownGroupSize, moe.DownBits, moe.DownQuantMode = model.ResolveLinearQuantParams(
|
||||
cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode,
|
||||
cfg.TensorQuant, downKey, downW, downScales,
|
||||
)
|
||||
}
|
||||
|
||||
// Router implements Gemma 4's expert routing mechanism.
|
||||
type Router struct {
|
||||
Proj nn.LinearLayer // [hidden_size -> num_experts]
|
||||
@@ -761,32 +801,32 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
|
||||
|
||||
moe := &MoEBlock{PerExpertScale: perExpertScale}
|
||||
|
||||
// Check for pre-stacked tensors (unquantized HF format).
|
||||
// Try .experts. first (new weight drop), fall back to .moe. (old format).
|
||||
gateUpW := tensors[layerPrefix+".experts.gate_up_proj"]
|
||||
// Check for pre-stacked tensors (BF16 or quantized HF format).
|
||||
// Try .experts. first (new weight drop), fall back to .moe. (old
|
||||
// format). Track the matched key so loadFusedExperts can find the
|
||||
// quant sidecars.
|
||||
gateUpKey := layerPrefix + ".experts.gate_up_proj"
|
||||
gateUpW := tensors[gateUpKey]
|
||||
if gateUpW == nil {
|
||||
gateUpW = tensors[layerPrefix+".moe.gate_up_proj"]
|
||||
gateUpKey = layerPrefix + ".moe.gate_up_proj"
|
||||
gateUpW = tensors[gateUpKey]
|
||||
}
|
||||
gateW := tensors[layerPrefix+".experts.gate_proj"]
|
||||
if gateW == nil {
|
||||
gateW = tensors[layerPrefix+".moe.gate_proj"]
|
||||
}
|
||||
if gateUpW != nil {
|
||||
// Fused gate+up: split along dim 1, transpose for GatherMM.
|
||||
dims := gateUpW.Dims()
|
||||
half := int32(dims[1] / 2)
|
||||
gateSlice := sliceAxis1(gateUpW, 0, half)
|
||||
upSlice := sliceAxis1(gateUpW, half, int32(dims[1]))
|
||||
moe.GateWeight = transposeForGatherMM(gateSlice)
|
||||
moe.UpWeight = transposeForGatherMM(upSlice)
|
||||
downW := tensors[layerPrefix+".experts.down_proj"]
|
||||
// Fused gate+up (dense or quantized): split happens at matmul time.
|
||||
downKey := layerPrefix + ".experts.down_proj"
|
||||
downW := tensors[downKey]
|
||||
if downW == nil {
|
||||
downW = tensors[layerPrefix+".moe.down_proj"]
|
||||
downKey = layerPrefix + ".moe.down_proj"
|
||||
downW = tensors[downKey]
|
||||
}
|
||||
if downW == nil {
|
||||
return fmt.Errorf("layer %d: missing MoE down_proj with fused gate_up_proj", i)
|
||||
}
|
||||
moe.DownWeight = transposeForGatherMM(downW)
|
||||
loadFusedExperts(tensors, m.TextConfig, moe, gateUpW, gateUpKey, downW, downKey)
|
||||
} else if gateW != nil {
|
||||
// Separate gate_proj and up_proj (older format). Transpose for GatherMM.
|
||||
moe.GateWeight = transposeForGatherMM(gateW)
|
||||
@@ -814,8 +854,8 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
|
||||
return fmt.Errorf("layer %d: missing switch_mlp down_proj", i)
|
||||
}
|
||||
|
||||
// Check for quantized weights (scales present).
|
||||
// The scale key depends on whether the tensor has .weight suffix.
|
||||
// Resolve base keys: the scale/bias suffix attaches to whichever
|
||||
// of .weight / bare matched the weight tensor.
|
||||
gateUpKey := layerPrefix + ".moe.switch_mlp.gate_up_proj.weight"
|
||||
if tensors[gateUpKey] == nil {
|
||||
gateUpKey = layerPrefix + ".moe.switch_mlp.gate_up_proj"
|
||||
@@ -824,47 +864,7 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
|
||||
if tensors[downKey] == nil {
|
||||
downKey = layerPrefix + ".moe.switch_mlp.down_proj"
|
||||
}
|
||||
gateUpScales := firstNonNil(tensors, gateUpKey+"_scale", gateUpKey+".scale")
|
||||
downScales := firstNonNil(tensors, downKey+"_scale", downKey+".scale")
|
||||
|
||||
if gateUpScales != nil && downScales != nil {
|
||||
// Quantized: keep fused gate_up as single tensor for GatherQMM.
|
||||
// One fused call instead of two separate gate+up calls.
|
||||
gateUpBiases := firstNonNil(tensors, gateUpKey+"_qbias", gateUpKey+".bias")
|
||||
downBiases := firstNonNil(tensors, downKey+"_qbias", downKey+".bias")
|
||||
|
||||
moe.GateUpWeightQ = switchGateUp
|
||||
moe.GateUpScales = gateUpScales
|
||||
moe.GateUpBiases = gateUpBiases
|
||||
moe.DownWeightQ = switchDown
|
||||
moe.DownScales = downScales
|
||||
if downBiases != nil {
|
||||
moe.DownBiases = downBiases
|
||||
}
|
||||
|
||||
groupSize, bits, mode := model.ResolveLinearQuantParams(
|
||||
m.QuantGroupSize, m.QuantBits, m.QuantMode,
|
||||
m.TensorQuant, gateUpKey, switchGateUp, gateUpScales,
|
||||
)
|
||||
moe.UseQuantized = true
|
||||
moe.UseFusedGateUp = true
|
||||
moe.GateUpGroupSize = groupSize
|
||||
moe.GateUpBits = bits
|
||||
moe.QuantMode = mode
|
||||
|
||||
dGroupSize, dBits, dMode := model.ResolveLinearQuantParams(
|
||||
m.QuantGroupSize, m.QuantBits, m.QuantMode,
|
||||
m.TensorQuant, downKey, switchDown, downScales,
|
||||
)
|
||||
moe.DownGroupSize = dGroupSize
|
||||
moe.DownBits = dBits
|
||||
moe.DownQuantMode = dMode
|
||||
} else {
|
||||
// Unquantized switch_mlp: keep fused and transpose for GatherMM.
|
||||
moe.GateUpWeight = transposeForGatherMM(switchGateUp)
|
||||
moe.UseFusedGateUp = true
|
||||
moe.DownWeight = transposeForGatherMM(switchDown)
|
||||
}
|
||||
loadFusedExperts(tensors, m.TextConfig, moe, switchGateUp, gateUpKey, switchDown, downKey)
|
||||
} else {
|
||||
// Per-expert tensors (from create path).
|
||||
// Try separate gate_proj/up_proj first, then fused gate_up_proj.
|
||||
|
||||
Reference in New Issue
Block a user