[Diffusion] Register runtime cache for direct DLO H2D
Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup. Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
1 parent
67f15e73cd
commit
23ab14a15b
18 files changed
+976
-74
No files matched your search
@@ -492,6 +492,59 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch,
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(
|
||||
use_hsdp=False,
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
data_parallel_size=1,
|
||||
cfg_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
),
|
||||
quantization_config=None,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=False,
|
||||
dlo_enable_runtime_cache=True,
|
||||
dlo_runtime_cache_dir=str(tmp_path),
|
||||
dlo_runtime_cache_pin_limit_gib=1.0,
|
||||
model="unused",
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
model = nn.Module()
|
||||
model.transformer = nn.Linear(2, 2, bias=False)
|
||||
events: list[str] = []
|
||||
runtime_plan = HostWeightPlan(
|
||||
backing_kind="runtime_cache",
|
||||
bindings={},
|
||||
runtime_layout_key="layout-key",
|
||||
post_load_complete=True,
|
||||
)
|
||||
|
||||
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
|
||||
loader.load_weights = lambda _model: events.append("load") # type: ignore[method-assign]
|
||||
loader._process_weights_after_loading = lambda *_args: events.append("process") # type: ignore[method-assign]
|
||||
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
|
||||
|
||||
def unexpected_checkpoint_plan(*_args, **_kwargs):
|
||||
pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
|
||||
|
||||
def build_cache(_pipeline, **_kwargs):
|
||||
events.append("cache")
|
||||
return HostWeightPlanResult(runtime_plan)
|
||||
|
||||
monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
|
||||
monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
|
||||
|
||||
assert loader.load_model(load_device="cpu") is model
|
||||
assert events == ["load", "process", "cache"]
|
||||
assert loader.take_host_weight_plan() is runtime_plan
|
||||
|
||||
|
||||
def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
|
||||
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
|
||||
import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
|
||||
|
||||
Reference in new issue
Block a user