[diffusion] feat: support nunchaku for Z-Image-Turbo and flux.1 (int4) (#18959)

Author: mickqian

python/sglang/multimodal_gen/configs/models/dits/flux.py

@@ -23,9 +23,36 @@ class FluxArchConfig(DiTArchConfig):
 
     stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
 
+    # nunchaku checkpoint uses different weight names; map to sglang flux layout
     param_names_mapping: dict = field(
         default_factory=lambda: {
-            r"transformer\.(\w*)\.(.*)$": r"\1.\2",
+            # HF diffusers format
+            r"^transformer\.(\w*)\.(.*)$": r"\1.\2",
+            # transformer_blocks nunchaku format (raw export - before internal conversion)
+            r"^transformer_blocks\.(\d+)\.mlp_fc1\.(.*)$": r"transformer_blocks.\1.ff.net.0.proj.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_fc2\.(.*)$": r"transformer_blocks.\1.ff.net.2.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_context_fc1\.(.*)$": r"transformer_blocks.\1.ff_context.net.0.proj.\2",
+            r"^transformer_blocks\.(\d+)\.mlp_context_fc2\.(.*)$": r"transformer_blocks.\1.ff_context.net.2.\2",
+            r"^transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_qkv.\2",
+            r"^transformer_blocks\.(\d+)\.qkv_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
+            r"^transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"transformer_blocks.\1.attn.to_out.0.\2",
+            r"^transformer_blocks\.(\d+)\.out_proj_context\.(.*)$": r"transformer_blocks.\1.attn.to_add_out.\2",
+            r"^transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"transformer_blocks.\1.attn.norm_q.\2",
+            r"^transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"transformer_blocks.\1.attn.norm_k.\2",
+            r"^transformer_blocks\.(\d+)\.norm_added_q\.(.*)$": r"transformer_blocks.\1.attn.norm_added_q.\2",
+            r"^transformer_blocks\.(\d+)\.norm_added_k\.(.*)$": r"transformer_blocks.\1.attn.norm_added_k.\2",
+            # transformer_blocks nunchaku format (already converted with convert_flux_state_dict)
+            r"^transformer_blocks\.(\d+)\.attn\.add_qkv_proj\.(.*)$": r"transformer_blocks.\1.attn.to_added_qkv.\2",
+            # single_transformer_blocks nunchaku format (raw export - before internal conversion)
+            r"^single_transformer_blocks\.(\d+)\.qkv_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_qkv.\2",
+            r"^single_transformer_blocks\.(\d+)\.out_proj\.(.*)$": r"single_transformer_blocks.\1.attn.to_out.0.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm_q\.(.*)$": r"single_transformer_blocks.\1.attn.norm_q.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm_k\.(.*)$": r"single_transformer_blocks.\1.attn.norm_k.\2",
+            # nunchaku quantization parameter name conversions (apply to all blocks)
+            r"^(.*)\.smooth_orig$": r"\1.smooth_factor_orig",
+            r"^(.*)\.smooth$": r"\1.smooth_factor",
+            r"^(.*)\.lora_down$": r"\1.proj_down",
+            r"^(.*)\.lora_up$": r"\1.proj_up",
         }
     )
 

python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py

@@ -16,27 +16,6 @@
 
 logger = init_logger(__name__)
 
-SVDQ_W4A4_LAYER_PATTERNS = [
-    "attn.to_qkv",
-    "attn.to_out",
-    "attn.add_qkv_proj",
-    "attn.to_add_out",
-    "img_mlp",
-    "txt_mlp",
-]
-
-AWQ_W4A16_LAYER_PATTERNS = [
-    "img_mod",
-    "txt_mod",
-]
-
-SKIP_QUANTIZATION_PATTERNS = [
-    "norm",
-    "embed",
-    "rotary",
-    "pos_embed",
-]
-
 
 @lru_cache(maxsize=1)
 def is_nunchaku_available() -> bool:
@@ -61,13 +40,15 @@ class NunchakuConfig(QuantizationConfig):
         group_size: Quantization group size (automatically set based on precision)
         act_unsigned: Use unsigned activation quantization
         quantized_model_path: Path to pre-quantized model weights (.safetensors)
+        model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules()
     """
 
-    precision: str = "int4"  # "int4" or "nvfp4"
+    precision: str = "int4"
     rank: int = 32
     group_size: Optional[int] = None
     act_unsigned: bool = False
     quantized_model_path: Optional[str] = None
+    model_cls: Optional[type] = None
 
     @classmethod
     def get_name(cls) -> str:
@@ -99,15 +80,27 @@ def from_config(cls, config: dict[str, Any]) -> "NunchakuConfig":
     def get_quant_method(
         self, layer: torch.nn.Module, prefix: str
     ) -> Optional[QuantizeMethodBase]:
-
         if not isinstance(layer, LinearBase):
             return None
 
-        for pattern in SKIP_QUANTIZATION_PATTERNS:
+        # get quantization rules from model class
+        quant_rules = self._get_quant_rules()
+
+        # priority: skip > awq_w4a16 > svdq_w4a4 > default
+        skip_patterns = quant_rules.get("skip", [])
+        for pattern in skip_patterns:
             if pattern in prefix.lower():
                 return None
 
-        for pattern in SVDQ_W4A4_LAYER_PATTERNS:
+        awq_patterns = quant_rules.get("awq_w4a16", [])
+        for pattern in awq_patterns:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuAWQLinearMethod
+
+                return NunchakuAWQLinearMethod(group_size=64)
+
+        svdq_patterns = quant_rules.get("svdq_w4a4", [])
+        for pattern in svdq_patterns:
             if pattern in prefix:
                 from ..nunchaku_linear import NunchakuSVDQLinearMethod
 
@@ -117,14 +110,7 @@ def get_quant_method(
                     act_unsigned=self.act_unsigned,
                 )
 
-        for pattern in AWQ_W4A16_LAYER_PATTERNS:
-            if pattern in prefix:
-                from ..nunchaku_linear import NunchakuAWQLinearMethod
-
-                return NunchakuAWQLinearMethod(
-                    group_size=64,
-                )
-
+        # default: apply svdq_w4a4 to all remaining linear layers
         from ..nunchaku_linear import NunchakuSVDQLinearMethod
 
         return NunchakuSVDQLinearMethod(
@@ -133,6 +119,13 @@ def get_quant_method(
             act_unsigned=self.act_unsigned,
         )
 
+    def _get_quant_rules(self) -> dict[str, list[str]]:
+        if self.model_cls is not None and hasattr(
+            self.model_cls, "get_nunchaku_quant_rules"
+        ):
+            return self.model_cls.get_nunchaku_quant_rules()
+        return {}
+
     def __post_init__(self):
         if self.group_size is None:
             if self.precision == "nvfp4":

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -94,8 +94,9 @@ def load_customized(
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         nunchaku_config = server_args.nunchaku_config
-
         if nunchaku_config is not None:
+            nunchaku_config.model_cls = model_cls
+
             # respect dtype from checkpoint
             # TODO: improve the condition
             param_dtype = None
@@ -158,7 +159,7 @@ def load_customized(
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
 
         # considering the existent of mixed-precision models (e.g., nunchaku)
-        if next(model.parameters()).dtype != param_dtype:
+        if next(model.parameters()).dtype != param_dtype and param_dtype:
             logger.warning(
                 f"Model dtype does not match expected param dtype, {next(model.parameters()).dtype} vs {param_dtype}"
             )

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -238,7 +238,8 @@ def load_model_from_full_model_state_dict(
     """
     meta_sd = model.state_dict()
     param_dict = dict(model.named_parameters())
-    sharded_sd = {}
+
+    # map names from checkpoint to customized names
     custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
         full_sd_iterator, param_names_mapping
     )  # type: ignore
@@ -250,7 +251,7 @@ def load_model_from_full_model_state_dict(
     # sort parameter names to ensure all ranks process parameters in the same order
     sorted_param_names = sorted(custom_param_sd.keys())
 
-    requires_grad = False
+    sharded_sd = {}
 
     # shard from loaded state_dict, custom_param_sd -> sharded_sd
     for target_param_name in sorted_param_names:

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -31,7 +31,7 @@ def set_default_torch_dtype(dtype: torch.dtype):
 
 
 def get_param_names_mapping(
-    mapping_dict: dict[str, str],
+    mapping_dict: dict[str, str | tuple[str, int, int]],
 ) -> Callable[[str], tuple[str, Any, Any]]:
     """
     Creates a mapping function that transforms parameter names using regex patterns.
@@ -44,21 +44,50 @@ def get_param_names_mapping(
     """
 
     def mapping_fn(name: str) -> tuple[str, Any, Any]:
-        # Try to match and transform the name using the regex patterns in mapping_dict
-        for pattern, replacement in mapping_dict.items():
-            match = re.match(pattern, name)
-            if match:
-                merge_index = None
-                total_split_params = None
+        # support chained conversions, e.g.:
+        # transformer.xxx.lora_down -> xxx.lora_down -> xxx.proj_down
+        merge_index = None
+        total_split_params = None
+        max_steps = max(8, len(mapping_dict) * 2)
+        applied_patterns: set[str] = set()
+        visited_names: set[str] = {name}
+
+        for _ in range(max_steps):
+            transformed = False
+            for pattern, replacement in mapping_dict.items():
+                # avoid re-applying the same rule on its own output
+                if pattern in applied_patterns:
+                    continue
+                if re.match(pattern, name) is None:
+                    continue
+
+                curr_merge_index = None
+                curr_total_split_params = None
                 if isinstance(replacement, tuple):
-                    merge_index = replacement[1]
-                    total_split_params = replacement[2]
+                    curr_merge_index = replacement[1]
+                    curr_total_split_params = replacement[2]
                     replacement = replacement[0]
-                name = re.sub(pattern, replacement, name)
-                return name, merge_index, total_split_params
 
-        # If no pattern matches, return the original name
-        return name, None, None
+                new_name = re.sub(pattern, replacement, name)
+
+                if new_name != name:
+                    if curr_merge_index is not None:
+                        merge_index = curr_merge_index
+                        total_split_params = curr_total_split_params
+
+                    name = new_name
+                    applied_patterns.add(pattern)
+                    if name in visited_names:
+                        transformed = False
+                        break
+                    visited_names.add(name)
+                    transformed = True
+                    break
+
+            if not transformed:
+                break
+
+        return name, merge_index, total_split_params
 
     return mapping_fn
 
@@ -150,25 +179,5 @@ def _list_safetensors_files(model_path: str) -> list[str]:
 
 BYTES_PER_GB = 1024**3
 
-
-def get_memory_usage_of_component(module) -> float | None:
-    """
-    returned value is in GB, rounded to 2 decimal digits
-    """
-    if not isinstance(module, nn.Module):
-        return None
-    if hasattr(module, "get_memory_footprint"):
-        usage = module.get_memory_footprint() / BYTES_PER_GB
-    else:
-        # manually
-        param_size = sum(p.numel() * p.element_size() for p in module.parameters())
-        buffer_size = sum(b.numel() * b.element_size() for b in module.buffers())
-
-        total_size_bytes = param_size + buffer_size
-        usage = total_size_bytes / (1024**3)
-
-    return round(usage, 2)
-
-
 # component name ->  ComponentLoader class
 component_name_to_loader_cls: Dict[str, Type[Any]] = {}

python/sglang/multimodal_gen/runtime/models/dits/base.py

@@ -107,3 +107,17 @@ class CachableDiT(TeaCacheMixin, BaseDiT):
     def __init__(self, config: DiTConfig, **kwargs) -> None:
         super().__init__(config, **kwargs)
         self._init_teacache_state()
+
+    @classmethod
+    def get_nunchaku_quant_rules(cls) -> dict[str, dict[str, Any]]:
+        """
+        Get quantization rules for Nunchaku quantization.
+
+        Returns a dict mapping layer name patterns to quantization configs:
+        {
+            "skip": [list of patterns to skip quantization],
+            "svdq_w4a4": [list of patterns for SVDQ W4A4],
+            "awq_w4a16": [list of patterns for AWQ W4A16],
+        }
+        """
+        return {}