From 4730d7b9478346ff1d69a8922eef367f71823d3a Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 4 Sep 2026 17:32:43 +0200 Subject: [PATCH 1/2] global scale for nvfp4 moes --- mlx_lm/models/switch_layers.py | 63 ++++++++++++++++++++++++++-------- mlx_lm/utils.py | 15 ++++++++ tests/test_models.py | 42 +++++++++++++++++++++++ 3 files changed, 105 insertions(+), 15 deletions(-) diff --git a/mlx_lm/models/switch_layers.py b/mlx_lm/models/switch_layers.py index 1fe5d917e..2f3d518c8 100644 --- a/mlx_lm/models/switch_layers.py +++ b/mlx_lm/models/switch_layers.py @@ -24,6 +24,18 @@ def _scatter_unsort(x, inv_order, shape=None): return x +def _quantize_experts(w, group_size, bits, mode): + """Quantize each expert with its own nvfp4 tensor scale.""" + if mode != "nvfp4": + raise ValueError(f"A global scale needs 'nvfp4' mode, got '{mode}'.") + gs = mx.max(mx.abs(w), axis=(-2, -1)).astype(mx.float32) + qs = [ + mx.quantize(w[e], group_size, bits, mode=mode, global_scale=gs[e]) + for e in range(w.shape[0]) + ] + return mx.stack([q for q, _ in qs]), mx.stack([s for _, s in qs]), gs + + class QuantizedSwitchLinear(nn.Module): def __init__( self, @@ -34,21 +46,29 @@ def __init__( group_size: int = 64, bits: int = 4, mode: str = "affine", + global_scale: bool = False, ): super().__init__() scale = math.sqrt(1 / input_dims) - self.weight, self.scales, *biases = mx.quantize( - mx.random.uniform( - low=-scale, - high=scale, - shape=(num_experts, output_dims, input_dims), - ), - group_size=group_size, - bits=bits, - mode=mode, + weight = mx.random.uniform( + low=-scale, + high=scale, + shape=(num_experts, output_dims, input_dims), ) - self.biases = biases[0] if biases else None + if global_scale: + self.weight, self.scales, self.global_scale = _quantize_experts( + weight, group_size, bits, mode + ) + self.biases = None + else: + self.weight, self.scales, *biases = mx.quantize( + weight, + group_size=group_size, + bits=bits, + mode=mode, + ) + self.biases = biases[0] if biases else None if bias: self.bias = mx.zeros((num_experts, output_dims)) @@ -83,6 +103,7 @@ def __call__(self, x, indices, sorted_indices=False): group_size=self.group_size, bits=self.bits, mode=self.mode, + global_scale=self.get("global_scale"), sorted_indices=sorted_indices, ) if "bias" in self: @@ -128,7 +149,13 @@ def __call__(self, x, indices, sorted_indices=False): x = x + mx.expand_dims(self["bias"][indices], -2) return x - def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"): + def to_quantized( + self, + group_size: int = 64, + bits: int = 4, + mode: str = "affine", + global_scale: bool = False, + ): num_experts, output_dims, input_dims = self.weight.shape ql = QuantizedSwitchLinear( input_dims, @@ -139,10 +166,16 @@ def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine" bits, mode=mode, ) - ql.weight, ql.scales, *biases = mx.quantize( - self.weight, group_size, bits, mode=mode - ) - ql.biases = biases[0] if biases else None + if global_scale: + ql.weight, ql.scales, ql.global_scale = _quantize_experts( + self.weight, group_size, bits, mode + ) + ql.biases = None + else: + ql.weight, ql.scales, *biases = mx.quantize( + self.weight, group_size, bits, mode=mode + ) + ql.biases = biases[0] if biases else None if "bias" in self: ql.bias = self.bias diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ff8252d8e..18d3ed824 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -830,12 +830,19 @@ def save_model( ) +def _takes_global_scale(module: nn.Module) -> bool: + """Only some layers hold an nvfp4 tensor scale.""" + params = inspect.signature(module.to_quantized).parameters + return "global_scale" in params + + def quantize_model( model: nn.Module, config: dict, group_size: Optional[int], bits: Optional[int], mode: str = "affine", + global_scale: bool = False, quant_predicate: Optional[Callable[[str, nn.Module], Union[bool, dict]]] = None, ) -> Tuple[nn.Module, dict]: """ @@ -847,6 +854,8 @@ def quantize_model( group_size (Optional[int]): Group size for quantization. bits (Optional[int]): Bits per weight for quantization. mode (str): The quantization mode. + global_scale (bool): Use one ``nvfp4`` tensor scale per expert on the + switch layers. Only these layers support it for now. quant_predicate (Callable): A callable that decides how to quantize each layer based on the path. Accepts the layer `path` and the `module`. Returns either a bool to signify quantize/no quantize or @@ -887,6 +896,12 @@ def wrapped_predicate(path, module): bool_or_params = True if quant_predicate is not None: bool_or_params = quant_predicate(path, module) + # The scale is a parameter of the layer, so the config must record it + # for the loader to rebuild the same shapes. + if global_scale and bool_or_params and _takes_global_scale(module): + if not isinstance(bool_or_params, dict): + bool_or_params = dict(quant_params) + bool_or_params["global_scale"] = True if isinstance(bool_or_params, dict): quantized_config["quantization"][path] = bool_or_params elif fine_grained_config and bool_or_params: diff --git a/tests/test_models.py b/tests/test_models.py index 0a0c53028..1ce7474d9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -406,6 +406,48 @@ def model_test_runner(self, model, model_type, vocab_size, num_layers): # Make sure the model can be copied / pickled copy.deepcopy(model) + @unittest.skipIf( + not mx.metal.is_available(), "Global scale is only supported on Metal backend" + ) + def test_switch_linear_global_scale(self): + from mlx_lm.models.switch_layers import SwitchLinear + + mx.random.seed(0) + E, N, K = 8, 128, 256 + layer = SwitchLinear(K, N, E, bias=False) + # Scale the experts apart so a mixed up scale shows in the output. + layer.weight = layer.weight * mx.array( + [[1, 2, 4, 8][e % 4] for e in range(E)], mx.float32 + ).reshape((E, 1, 1)) + + ql = layer.to_quantized(group_size=16, bits=4, mode="nvfp4", global_scale=True) + self.assertEqual(ql.global_scale.shape, (E,)) + self.assertEqual(ql.global_scale.dtype, mx.float32) + + x = mx.random.normal((16, 1, K)) + indices = mx.random.randint(0, E, (16,)) + w_hat = mx.stack( + [ + mx.dequantize( + ql.weight[e], + ql.scales[e], + mode="nvfp4", + global_scale=ql.global_scale[e], + dtype=x.dtype, + ) + for e in range(E) + ] + ) + expected = x @ w_hat[indices].swapaxes(-1, -2) + self.assertTrue(mx.allclose(ql(x, indices), expected, atol=1e-4)) + + # Without it the layer keeps the plain nvfp4 parameters + self.assertNotIn( + "global_scale", layer.to_quantized(group_size=16, bits=4, mode="nvfp4") + ) + with self.assertRaises(ValueError): + layer.to_quantized(group_size=32, bits=4, mode="mxfp4", global_scale=True) + def test_bailing_moe_v3(self): from dataclasses import replace From a6ff1bf443eb3be1ffdd5a8640ff546f0a5644ad Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 11 Sep 2026 15:34:51 +0200 Subject: [PATCH 2/2] handle externally quantized models --- mlx_lm/utils.py | 38 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 18d3ed824..a294a0bd5 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -378,6 +378,14 @@ def load_model( if "quantization_config" in text_config: config["quantization_config"] = text_config["quantization_config"] + if "quantization_config" not in config: + # NVIDIA ModelOpt exports keep their quantization metadata in a separate + # file rather than in config.json. + hf_quant_config = model_path / "hf_quant_config.json" + if hf_quant_config.exists(): + with open(hf_quant_config, "r") as fid: + config["quantization_config"] = json.load(fid) + model_args = model_args_class.from_dict(config) model = model_class(model_args) @@ -392,7 +400,20 @@ def class_predicate(p, m): return config["quantization"][p] if not hasattr(m, "to_quantized"): return False - return f"{p}.scales" in weights + if f"{p}.scales" not in weights: + return False + # An nvfp4 tensor scale is a parameter of the layer, so the layer has + # to be built to hold one. Whether it needs one is a property of the + # checkpoint, so read it from the weights rather than the config: + # externally produced nvfp4 does not carry per-layer entries. + if f"{p}.global_scale" in weights and _takes_global_scale(m): + return { + "group_size": quantization["group_size"], + "bits": quantization["bits"], + "mode": quantization.get("mode", "affine"), + "global_scale": True, + } + return True nn.quantize( model, @@ -427,6 +448,19 @@ def class_predicate(p, m): config["quantization"] = quantization config["quantization_config"] = quantization _quantize(quantization) + elif quant_method == "modelopt": + # NVIDIA ModelOpt. + algo = quantization_config["quantization"]["quant_algo"] + if algo != "NVFP4": + raise ValueError(f"Unsupported modelopt quant_algo: {algo}") + quantization = { + "group_size": quantization_config["quantization"].get("group_size", 16), + "bits": 4, + "mode": "nvfp4", + } + config["quantization"] = quantization + config["quantization_config"] = quantization + _quantize(quantization) elif quant_method in ("awq", "gptq"): # Transform AutoAWQ/GPTQ packed weights to MLX format weights, quantization = _transform_awq_weights(weights, quantization_config) @@ -886,7 +920,7 @@ def defaults_for_mode(mode, group_size, bits): fine_grained_config = True else: fine_grained_config = False - quantized_config["quantization"] = quant_params + quantized_config["quantization"] = dict(quant_params) def wrapped_predicate(path, module): if not hasattr(module, "to_quantized"):