diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 5cc5b5fd9..8b4a6a4dc 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 888cea116..e53d70eee 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/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)) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index c80a88dbf..38b7c8b89 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index e6685cd6e..590e42287 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 2401c1006..35e4f1b94 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 5a62d7168..4d465cb55 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/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,