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

[Feat] Enable expert parallel for diffusion MoE layers (#1323)

Browse Source 
Signed-off-by: Semmer2 <semmer@live.cn>
This commit is contained in:
Ting FU authored and GitHub committed 2026-03-11 10:45:36 +08:00
1 parent 4725593a4b
commit 6ac85d4607
11 files changed
+361 -18
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/worker/diffusion_worker.py
+2
View File
@@ -106,6 +106,7 @@ class DiffusionWorker:
vllm_config = VllmConfig(compilation_config=CompilationConfig())
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
vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
self.vllm_config = vllm_config
# Initialize distributed environment
@@ -127,6 +128,7 @@ class DiffusionWorker:
pipeline_parallel_size=parallel_config.pipeline_parallel_size,
fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1,
hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1,
enable_expert_parallel=parallel_config.enable_expert_parallel,
)
init_workspace_manager(self.device)