fix(gemma4): diagnose malformed per-module quant overrides#695
Merged
Conversation
PR #690 added Gemma 4 per-module MLP quantization overrides via a flattened `overrides` map on `QuantizationArgs`, but a malformed override value (a non-object value, a wrong-typed/negative/zero `group_size`, or an unsupported `bits` width) silently fell back to the root defaults with no indication anything was wrong, which is a debugging hazard for quantization frontier work (issue #691). `QuantizationArgs::quant_params_for` now validates the entry matched by the queried module `prefix` before applying it: a non-object value, or a `group_size`/`bits` field with the wrong type, a negative value, a zero group size, or an unsupported bit width (mlxcel only supports the affine `{2,3,4,5,6,8}` set) is reported via `tracing::warn!` naming the offending key, then the loader falls back to the root defaults for that field exactly as before. Metadata keys that legitimately live alongside per-module overrides in the same flattened map, such as `mode`, `quant_method`, `quant_algo`, and `config_groups`, are untouched because they never coincide with a real module-prefix lookup. Added a `quant_override_diagnostics_tests` module covering valid full and partial overrides, absent overrides, a non-object override value, string-typed and negative/zero/unsupported-width numeric fields, sibling metadata keys, and a typo'd prefix that simply never matches, verifying valid configs keep loading unchanged while malformed ones are diagnosed instead of silently mis-routed.
Member
Author
PR Finalization CompleteSummary
No source changes were needed; working tree is unchanged from the prior commit. Ready for merge. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
PR #690 added Gemma 4 per-module MLP quantization overrides through a flattened
overridesmap onQuantizationArgs, but a malformed override value (a non-object value, a wrong-typed/negative/zerogroup_size, or an unsupportedbitswidth) silently fell back to the root defaults with no indication anything was wrong, which is a debugging hazard for quantization frontier work.What changed
QuantizationArgs::quant_params_forinsrc/models/gemma4.rsnow validates the override entry matched by the queried moduleprefixbefore applying it, reporting an actionabletracing::warn!naming the offending key for a non-object override value, a wrong-typed or negativegroup_size/bits, a zero group size, or an unsupported bit width (mlxcel only supports the affine{2,3,4,5,6,8}set), then falls back to the root defaults for that field exactly as before.validate_override_group_sizeandvalidate_override_bits, that produce the actionable reason strings used in the warnings.mode,quant_method,quant_algo,config_groups) are untouched because they never coincide with a real module-prefix lookup, so they are never mistaken for malformed overrides.quant_override_diagnostics_testsmodule insrc/models/gemma4.rscovering valid full and partial overrides, absent overrides, a non-object override value, string-typed and negative/zero/unsupported-width numeric fields, sibling metadata keys, and a typo'd prefix that simply never matches.Test plan
cargo check -p mlxcel --features cuda --libcargo test -p mlxcel --features cuda --lib models::gemma4::quant(24 tests, all pass, including the 11 new diagnostics tests and the pre-existingquant_scheme_tests)cargo test -p mlxcel --features cuda --lib models::gemma4_tests(9 pass, 8 ignored for serial MLX execution as before, including the existinggemma4_root_quantization_preserves_per_module_overridesregression test, confirming valid configs load unchanged)cargo clippy -p mlxcel --features cuda --lib --tests(no warnings touchinggemma4.rs; the crate has one pre-existing unrelatedredundant_closurewarning insrc/models/sanitize.rsthat predates this change)cargo fmt -p mlxcel -- --checkCloses #691