Add global scale for nvfp4 quantized MoEs

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

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)
@@ -830,12 +864,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 +888,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
@@ -877,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"):
@@ -887,6 +930,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:

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