anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion][Bug-Fix]: Fix Ulysses SP Accurary Error (#377)

Browse Source 
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
This commit is contained in:
Didan DengandJiangyun Zhu authored and GitHub committed 2025-12-22 17:48:39 +08:00
1 parent b06a5148f6
commit ee2b78e826
12 files changed
+747 -199
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/worker/gpu_worker.py
+28 -25
View File
@@ -23,6 +23,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.forward_context import set_forward_context
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -65,31 +66,33 @@ class GPUWorker:
vllm_config = VllmConfig()
vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
self.vllm_config = vllm_config
with (
set_current_omni_diffusion_config(self.od_config),
set_current_vllm_config(vllm_config),
):
init_distributed_environment(world_size=world_size, rank=rank)
logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
parallel_config = self.od_config.parallel_config
initialize_model_parallel(
data_parallel_size=parallel_config.data_parallel_size,
cfg_parallel_size=parallel_config.cfg_parallel_size,
sequence_parallel_size=parallel_config.sequence_parallel_size,
ulysses_degree=parallel_config.ulysses_degree,
ring_degree=parallel_config.ring_degree,
tensor_parallel_size=parallel_config.tensor_parallel_size,
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
)
with set_current_omni_diffusion_config(self.od_config):
with set_current_vllm_config(vllm_config):
init_distributed_environment(world_size=world_size, rank=rank)
logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
parallel_config = self.od_config.parallel_config
initialize_model_parallel(
data_parallel_size=parallel_config.data_parallel_size,
cfg_parallel_size=parallel_config.cfg_parallel_size,
sequence_parallel_size=parallel_config.sequence_parallel_size,
ulysses_degree=parallel_config.ulysses_degree,
ring_degree=parallel_config.ring_degree,
tensor_parallel_size=parallel_config.tensor_parallel_size,
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
load_config = LoadConfig()
model_loader = DiffusersPipelineLoader(load_config)
time_before_load = time.perf_counter()
with DeviceMemoryProfiler() as m:
self.pipeline = model_loader.load_model(
od_config=self.od_config,
load_device=f"cuda:{rank}",
)
load_config = LoadConfig()
model_loader = DiffusersPipelineLoader(load_config)
time_before_load = time.perf_counter()
with DeviceMemoryProfiler() as m:
self.pipeline = model_loader.load_model(
od_config=self.od_config,
load_device=f"cuda:{rank}",
)
time_after_load = time.perf_counter()
time_after_load = time.perf_counter()
logger.info(
"Model loading took %.4f GiB and %.6f seconds",
@@ -116,8 +119,8 @@ class GPUWorker:
# Refresh cache context if needed
if self.cache_backend is not None and self.cache_backend.is_enabled():
self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
output = self.pipeline.forward(req)
with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
output = self.pipeline.forward(req)
return output
def shutdown(self) -> None: