Add global scale support to quantized layers#426
Conversation
| public let mode: QuantizationMode | ||
| public let scales: MLXArray | ||
| public let biases: MLXArray? | ||
| public let globalScale: MLXArray? |
There was a problem hiding this comment.
One thing we need to be careful of here is that there is currently no python support for this. That means:
- maybe we want to set the key as
global_scaleorglobal_scalesto match the likely python naming - if this is non-nil after quantizing the model, it will be required when loading weights
- will this cause a problem or is nvfp4 just not used?
There was a problem hiding this comment.
Good point. I changed the serialised parameter key to global_scale
davidkoski
left a comment
davidkoski
left a comment
There was a problem hiding this comment.
I like it, but see my questions -- I think these need some thought before merging.
Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations. Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there. Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface.
| if globalScale != nil { | ||
| return matmul(x, dequantizedWeight.T) |
There was a problem hiding this comment.
I think this is correct -- quantizedMM doesn't support globalScales. But it does seem like it defeats the purpose of the quantized path a little. I guess the other option is caching the dequantized weight, which is for sure the opposite intent :-)
I think this may be the best we can do until there is wider support (quantizedMM + metal).
There was a problem hiding this comment.
You are right, and I miss that, I am sorry. I thought about it awhile and I have added another commit that should really be the best that can be done here
There was a problem hiding this comment.
Oh this is a really good approach!
davidkoski
left a comment
davidkoski
left a comment
There was a problem hiding this comment.
Looks good. I worry that nvfp4 has sketchy support -- this would appear to load tensors quantized that way but I am not sure it will work in practice unless people do run on cuda/cpu.
I am OK with merging it, but just want to double check with you. It sounded like you had a use case for this and dealt with the restrictions.
Since MLX encodes nvfp4 group scales pre-multiplied by (448 * 6) / globalScale, the global scale factors out of dequantization linearly. Exploit this by always running quantizedMM (supported on Metal without globalScale) and applying globalScale / (448 * 6) to the output, instead of dequantizing the full weight and falling back to a plain matmul. This keeps weights quantized end-to-end on all backends, removes the dequantizedWeight fallback from QuantizedLinear and QuantizedEmbedding, and applies the scale before the bias add so the bias is unaffected. Add Metal tests that validate the GPU quantized path against CPU dequantized(globalScale:) references for QuantizedLinear (matmul + bias) and QuantizedEmbedding (asLinear and index lookup).
davidkoski
left a comment
davidkoski
left a comment
There was a problem hiding this comment.
Looks good, just waiting on CI. Thank you!
2 participants
Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations.
Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there.
Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface.
Proposed changes
Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes