[Diffusion] Add loader-owned host-weight plans for DLO

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-15 23:32:27 +08:00
1 parent d1e230c95b
commit e2796b6168
17 files changed
+1717 -588

No files matched your search

+4 -6
View File
@@ -27,11 +27,10 @@ class OffloadConfig:
pin_cpu_memory: bool = True
use_hsdp: bool = False
dp_size: int = 1 # derived from parallel_config, not user-configurable
# True: add DP sharding + AllGather. False: stream the standard loader's
# rank-local tensors (including TP-local shards) with H2D only.
# True: add DP sharding + AllGather. False: stream complete rank-local
# blocks from the loader-selected host backing 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
def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -58,7 +57,6 @@ class OffloadConfig:
parallel_config = getattr(od_config, "parallel_config", None)
use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False
# Derive dp_size from parallel_config — not user-configurable.
# The offload adapts to whatever DP/SP is already configured.
dp_size = 1
@@ -116,7 +114,8 @@ class OffloadConfig:
dp_size = 1
logger.info(
"Distributed layerwise offload: dlo_use_allgather=False, "
"streaming standard-loader rank-local weights (no DP shard or AllGather)"
"streaming complete rank-local blocks (no DLO shard or AllGather); "
"the backend will select mmap or standard-loader host storage"
)
# HSDP already shards parameters into DTensors. Running distributed
@@ -137,7 +136,6 @@ class OffloadConfig:
dp_size=dp_size,
dlo_use_allgather=dlo_use_allgather,
dlo_resident_layers=dlo_resident_layers,
model_path=getattr(od_config, "model", None),
)