[Diffusion] Add loader-owned host-weight plans for DLO
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
d1e230c95b
commit
e2796b6168
17 files changed
+1717
-588
No files matched your search
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user