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

feat(minimax-h3): enable 2x RTX 4090/5090 with DLO (#5764)

Browse Source 
Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>
This commit is contained in:
SYLARandlishunyang12 authored and GitHub committed 2026-08-06 09:02:29 +08:00
1 parent b4581b293b
commit 1c2a81f6d8
28 files changed
+1658 -162
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/offloader/base.py
+20 -7
View File
@@ -27,7 +27,10 @@ class OffloadConfig:
pin_cpu_memory: bool = True
use_hsdp: bool = False
dp_size: int = 1 # derived from parallel_config, not user-configurable
dlo_use_allgather: bool = True # if True: shard + AllGather; if False: full weights H2D only
# True: add DP sharding + AllGather. False: stream the standard loader's
# rank-local tensors (including TP-local shards) with H2D only.
dlo_use_allgather: bool = True
dlo_resident_layers: int = 0 # leading DiT layers kept on device
model_path: str | None = None # checkpoint path for mmap weight loading
@classmethod
@@ -94,17 +97,26 @@ class OffloadConfig:
else:
strategy = OffloadStrategy.NONE
# When dlo_use_allgather=False, each rank loads full weights (no sharding,
# no AllGather). This avoids AllGather synchronization requirements
# (concurrent requests, dummy run skip) at the cost of N× CPU memory.
# With dlo_use_allgather=False, do not add another DP shard. Each rank
# streams the tensors produced by the standard loader, which may
# already be TP-local shards. This avoids AllGather synchronization
# requirements (concurrent requests, dummy run skip).
dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
if dlo_resident_layers < 0:
raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
if dlo_resident_layers and dlo_use_allgather:
raise ValueError(
"dlo_resident_layers currently requires --dlo-no-use-allgather so "
"resident blocks use weights prepared by the standard TP-aware loader"
)
# If dlo_use_allgather=False, force dp_size=1 (each rank independent)
if enable_distributed_layerwise_offload and not dlo_use_allgather:
dp_size = 1
logger.info(
"Distributed layerwise offload: dlo_use_allgather=False, "
"each rank loads full weights (no shard, no AllGather)"
"streaming standard-loader rank-local weights (no DP shard or AllGather)"
)
# HSDP already shards parameters into DTensors. Running distributed
@@ -114,8 +126,8 @@ class OffloadConfig:
raise ValueError(
"Distributed layerwise offload with AllGather is incompatible with "
"HSDP: HSDP parameters are already sharded DTensors, and the offloader "
"would double-shard them. Use --dlo-no-use-allgather (full weights per "
"rank) or disable HSDP."
"would double-shard them. Use --dlo-no-use-allgather (standard-loader "
"rank-local weights) or disable HSDP."
)
return cls(
@@ -124,6 +136,7 @@ class OffloadConfig:
use_hsdp=use_hsdp,
dp_size=dp_size,
dlo_use_allgather=dlo_use_allgather,
dlo_resident_layers=dlo_resident_layers,
model_path=getattr(od_config, "model", None),
)