Add torchao quant for mixtral and qwen_moe (#1418)

python/sglang/srt/layers/torchao_utils.py

@@ -2,10 +2,20 @@
 Common utilities for torchao.
 """
 
+from typing import Dict, Set
+
 import torch
 
 
-def torchao_quantize_param_data(param, torchao_config):
+def torchao_quantize_param_data(param: torch.Tensor, torchao_config: str):
+    """Quantize a Tensor with torchao quantization specified by torchao_config
+
+    Args:
+       `param`: weight parameter of the linear module
+       `torchao_config`: type of quantization and their arguments we want to use to
+        quantize the Tensor, e.g. int4wo-128 means int4 weight only quantization with group_size
+        128
+    """
     # Lazy import to suppress some warnings
     from torchao.quantization import (
         int4_weight_only,
@@ -36,3 +46,30 @@ def torchao_quantize_param_data(param, torchao_config):
         # [rank0]: AssertionError: fp8e4nv data type is not supported on CUDA arch < 89
         quantize_(dummy_linear, float8_weight_only())
     return dummy_linear.weight
+
+
+def apply_torchao_config_(
+    self: torch.nn.Module,
+    params_dict: Dict[str, torch.Tensor],
+    param_suffixes: Set[str],
+) -> None:
+    """A util function used for quantizing the weight parameters after they are loaded if
+       self.torchao_config is specified
+
+    Args:
+      `self`: the model we want to quantize
+      `params_dict`: dictionary mapping from param_name to the parameter Tensor
+      `param_suffixes`: a set of suffixes, we'll quantize the Tensor matching these suffixes
+
+    Returns:
+       None, the `params_dict` is modified inplace and the weights of `self` model are quantized
+    """
+    if self.torchao_config:
+        for param_suffix in param_suffixes:
+            for name in params_dict:
+                param = params_dict[name]
+                if param_suffix in name and param.ndim == 2:
+                    params_dict[name] = torchao_quantize_param_data(
+                        param, self.torchao_config
+                    )
+        self.load_state_dict(params_dict, assign=True)