minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: remove accelerate dependency for device mapping (#18026)

Browse Source 
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
This commit is contained in:
CHEN Xi
2026-02-02 17:24:19 +08:00
committed by GitHub
co-authored by Kangyan-Zhou
parent e3021b65fe
commit aa780a6258
1 changed files with 3 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
+3 -29
View File
@@ -21,6 +21,7 @@ from diffusers import DiffusionPipeline
from PIL import Image from PIL import Image
from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase, ComposedPipelineBase,
) )
@@ -391,12 +392,7 @@ class DiffusersPipeline(ComposedPipelineBase):
self.model_path = model_path self.model_path = model_path
dtype = self._get_dtype(server_args) dtype = self._get_dtype(server_args)
device_map = self._get_device_map(server_args) logger.info("Loading diffusers pipeline with dtype=%s", dtype)
logger.info(
"Loading diffusers pipeline with dtype=%s, device_map=%s",
dtype,
device_map,
)
# Build common kwargs for from_pretrained # Build common kwargs for from_pretrained
load_kwargs = { load_kwargs = {
@@ -405,11 +401,6 @@ class DiffusersPipeline(ComposedPipelineBase):
"revision": server_args.revision, "revision": server_args.revision,
} }
# Add device_map for direct GPU loading and parallel shard loading
# This warms up CUDA caching allocator and enables parallel loading via accelerate
if device_map is not None:
load_kwargs["device_map"] = device_map
# Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit) # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit)
config = server_args.pipeline_config config = server_args.pipeline_config
if config is not None: if config is not None:
@@ -458,22 +449,13 @@ class DiffusersPipeline(ComposedPipelineBase):
else: else:
raise raise
# Only move to device if device_map wasn't used (already on device) pipe = pipe.to(get_local_torch_device())
if device_map is None:
if torch.cuda.is_available():
pipe = pipe.to("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
pipe = pipe.to("mps")
# Apply VAE memory optimizations from pipeline config # Apply VAE memory optimizations from pipeline config
self._apply_vae_optimizations(pipe, server_args) self._apply_vae_optimizations(pipe, server_args)
# Apply attention backend if specified # Apply attention backend if specified
self._apply_attention_backend(pipe, server_args) self._apply_attention_backend(pipe, server_args)
# Apply cache-dit acceleration if configured # Apply cache-dit acceleration if configured
pipe = self._apply_cache_dit(pipe, server_args) pipe = self._apply_cache_dit(pipe, server_args)
logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__) logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__)
return pipe return pipe
@@ -581,14 +563,6 @@ class DiffusersPipeline(ComposedPipelineBase):
logger.info("Enabled cache-dit for diffusers pipeline") logger.info("Enabled cache-dit for diffusers pipeline")
return pipe return pipe
def _get_device_map(self, server_args: ServerArgs) -> str | None:
"""
Determine device_map for pipeline loading.
"""
if not torch.cuda.is_available():
return None
return "cuda"
def _get_dtype(self, server_args: ServerArgs) -> torch.dtype: def _get_dtype(self, server_args: ServerArgs) -> torch.dtype:
""" """
Determine the dtype to use for model loading. Determine the dtype to use for model loading.