[diffusion] fix: fix z-Image online fp8 quantization crash with dit_cpu_offload (#29903)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -364,6 +364,11 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
         Batched prompts use stored text lengths. SP mode builds image caches for
         the local spatial shard.
         """
+        if rotary_emb is None:
+            raise ValueError(
+                "Z-Image transformer has no `rotary_emb`. It likely loaded via the "
+                "native diffusers fallback; check the load logs for the real error."
+            )
 
         def create_coordinate_grid(size, start=None, device=None):
             if start is None:

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

@@ -85,7 +85,12 @@ class TransformerLoader(ComponentLoader):
         component_server_args = _server_args_for_transformer_component(
             server_args, component_name
         )
-        return component_server_args.transformer_weights_path is not None
+        # Don't let a quantized load quietly fall back to the unquantized native
+        # model. That would drop the requested precision and bury the real error.
+        return (
+            component_server_args.transformer_weights_path is not None
+            or component_server_args.quantization is not None
+        )
 
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
@@ -165,6 +170,10 @@ class TransformerLoader(ComponentLoader):
             reduce_dtype=torch.float32,
             output_dtype=None,
             strict=False,
+            defer_cpu_offload_until_after_weight_processing=(
+                component_server_args.dit_cpu_offload
+                and quant_spec.requires_device_weight_processing
+            ),
         )
 
         # post-hooks (e.g., patch scales (nunchaku))

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

@@ -196,6 +196,7 @@ def maybe_load_fsdp_model(
     output_dtype: torch.dtype | None = None,
     pin_cpu_memory: bool = True,
     strict: bool = True,
+    defer_cpu_offload_until_after_weight_processing: bool = False,
 ) -> torch.nn.Module:
     """Load a model with optional FSDP (Fully Sharded Data Parallel) support.
 
@@ -206,6 +207,9 @@ def maybe_load_fsdp_model(
             - Weight loading and casting
         reduce_dtype: Data type for gradient reduction in FSDP mixed precision.
         strict: If True, enforce strict state dict loading (all keys must match).
+        defer_cpu_offload_until_after_weight_processing: If True, keep weights
+            on device until process_weights_after_loading completes, then apply
+            non-FSDP CPU offload.
     """
     # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
     # manually casting the inputs to the model
@@ -232,6 +236,17 @@ def maybe_load_fsdp_model(
         use_fsdp = False
         logger.info("Disabling FSDP for MPS platform as it's not compatible")
 
+    defer_cpu_offload = bool(
+        cpu_offload and defer_cpu_offload_until_after_weight_processing
+    )
+    if defer_cpu_offload and use_fsdp:
+        logger.warning(
+            "Ignoring deferred CPU offload for FSDP loading; keeping the existing "
+            "FSDP offload policy."
+        )
+        defer_cpu_offload = False
+    load_cpu_offload = cpu_offload and not defer_cpu_offload
+
     if use_fsdp:
         model._pre_fsdp_weight_loader_params = {
             n: p
@@ -251,7 +266,7 @@ def maybe_load_fsdp_model(
         )
         shard_model(
             model,
-            cpu_offload=cpu_offload,
+            cpu_offload=load_cpu_offload,
             reshard_after_forward=True,
             mp_policy=mp_policy,
             mesh=device_mesh,
@@ -280,7 +295,7 @@ def maybe_load_fsdp_model(
         device,
         param_dtype,
         strict=strict,
-        cpu_offload=cpu_offload,
+        cpu_offload=load_cpu_offload,
         param_names_mapping=param_names_mapping_fn,
     )
     if bnb_quant_states:
@@ -301,6 +316,8 @@ def maybe_load_fsdp_model(
             if _is_npu:
                 torch.npu.empty_cache()
     model.post_load_weights()
+    if defer_cpu_offload:
+        model.to("cpu")
 
     for n, p in chain(model.named_parameters(), model.named_buffers()):
         if p.is_meta:

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

@@ -122,6 +122,7 @@ class TransformerQuantLoadSpec:
     quant_config: Optional[QuantizationConfig]
     nunchaku_config: Optional[NunchakuConfig]
     param_dtype: Optional[torch.dtype]
+    requires_device_weight_processing: bool = False
     post_load_hooks: list[PostLoadHook] = field(default_factory=list)
 
     @property
@@ -255,6 +256,7 @@ class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter):
 
         quant_name_getter = getattr(type(quant_config), "get_name", None)
         quant_name = quant_name_getter() if callable(quant_name_getter) else None
+
         if quant_name != "modelopt_fp8":
             return
 
@@ -480,10 +482,25 @@ def resolve_transformer_quant_load_spec(
         quant_config=quant_config,
         nunchaku_config=nunchaku_config,
         param_dtype=param_dtype,
+        requires_device_weight_processing=_requires_device_weight_processing(
+            quant_config
+        ),
         post_load_hooks=post_load_hooks,
     )
 
 
+def _requires_device_weight_processing(
+    quant_config: Optional[QuantizationConfig],
+) -> bool:
+    """Return whether post-load weight processing needs CUDA/NPU tensors."""
+    quant_name = _get_quant_config_name(quant_config)
+    if quant_name == "fp8":
+        return not getattr(quant_config, "is_checkpoint_fp8_serialized", False)
+    if quant_name == "mxfp4":
+        return not getattr(quant_config, "is_checkpoint_mxfp4_serialized", False)
+    return False
+
+
 def _build_transformer_quant_adapters(
     *,
     cls_name: str,

python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 
 import torch
 
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp8Config,
 )
@@ -296,6 +297,21 @@ def test_modelopt_fp8_adapter_keeps_layerwise_offload_enabled():
     assert server_args.dit_layerwise_offload is True
 
 
+def test_modelopt_fp8_adapter_does_not_change_online_fp8_offload():
+    server_args = SimpleNamespace(
+        dit_cpu_offload=True,
+        dit_layerwise_offload=False,
+        quantization="fp8",
+    )
+
+    _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes(
+        server_args=server_args,
+        quant_config=Fp8Config(),
+    )
+
+    assert server_args.dit_cpu_offload is True
+
+
 def test_layerwise_capability_selects_layerwise_strategy_for_any_component():
     module = _LayerwiseComponent(enabled=True)
 

python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

@@ -49,6 +49,7 @@ from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
 from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
     NunchakuConfig,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
     _prepare_nvfp4_weight_bytes,
@@ -56,6 +57,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
 from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
     _filter_duplicate_precision_variant_safetensors,
     _Flux2Nvfp4FallbackAdapter,
+    _requires_device_weight_processing,
     resolve_transformer_quant_load_spec,
     resolve_transformer_safetensors_to_load,
 )
@@ -182,6 +184,14 @@ class TestTransformerQuantHelpers(unittest.TestCase):
 
         self.assertEqual(resolved, files)
 
+    def test_online_fp8_requires_device_weight_processing(self):
+        self.assertTrue(_requires_device_weight_processing(Fp8Config()))
+        self.assertFalse(
+            _requires_device_weight_processing(
+                Fp8Config(is_checkpoint_fp8_serialized=True)
+            )
+        )
+
     @patch(
         "sglang.multimodal_gen.runtime.loader.transformer_load_utils.build_nvfp4_config_from_safetensors_list",
         return_value=None,