[Feature] Distributed layerwise offload with DP multi-concurrency + mmap weight… (#5397)

Signed-off-by: evanchueng <evanchueng@outlook.com>
Signed-off-by: david6666666 <530634352@qq.com>
Co-authored-by: david6666666 <530634352@qq.com>

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -362,17 +362,42 @@ class DiffusersPipelineLoader:
                 )
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
-                logger.debug("Loading weights on %s ...", load_device)
-                if load_format == "diffusers":
-                    # DiffusersAdapterPipeline.load_weights() calls
-                    # DiffusionPipeline.from_pretrained() internally; it does
-                    # not use our native customized pipeline classes.
-                    cast(DiffusersAdapterPipeline, model).load_weights()
+
+                # Skip load_weights only for DLO+AllGather when the model
+                # supports mmap loading and no online quantization is active.
+                # This condition MUST match the gate in
+                # DistributedLayerwiseOffloadBackend.enable() so that the
+                # loader skips ⟺ enable() uses mmap.
+                from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading
+
+                _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+                _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
+                _has_online_quant = self._has_online_quant(model)
+                _supports_mmap = supports_mmap_loading(model)
+
+                _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
+
+                if _skip_load:
+                    logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
                 else:
-                    self.load_weights(model)
-                # HSDP processes quantized weights before wrapping parameters as
-                # DTensors. The non-HSDP path can process them here as usual.
-                self._process_weights_after_loading(model, target_device)
+                    if _dist_offload and _use_ag and _has_online_quant:
+                        raise ValueError(
+                            "Online quantization is incompatible with DLO+AllGather: "
+                            "the sharding + AllGather mechanism flattens weights by "
+                            "dtype, which breaks quantized weight/scale layouts. "
+                            "Please use --dlo-no-use-allgather or disable online "
+                            "quantization."
+                        )
+                    if _dist_offload and _use_ag and not _supports_mmap:
+                        logger.info(
+                            "DLO+AllGather active but model does not support mmap: loading weights via regular loader."
+                        )
+                    logger.debug("Loading weights on %s ...", load_device)
+                    if load_format == "diffusers":
+                        cast(DiffusersAdapterPipeline, model).load_weights()
+                    else:
+                        self.load_weights(model)
+                    self._process_weights_after_loading(model, target_device)
 
             if offload_after_quant:
                 model.to("cpu")