[Diffusion] Simplify DLO host plan setup

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-08-19 23:00:31 +08:00
1 parent 6bcd350301
commit 58eebee14b
7 files changed
+77 -37

No files matched your search

@@ -310,6 +310,7 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch
assert loader.load_model(load_device="cpu") is model
assert not loaded_ordinary_weights
assert plan.requires_backend
assert loader.take_host_weight_plan() is plan
assert loader.take_host_weight_plan() is None
@@ -449,7 +450,6 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
post_load_complete=True,
)
def load_weights(_model):
@@ -530,7 +530,6 @@ def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeyp
backing_kind="host_weight_cache",
bindings={},
runtime_layout_key="layout-key",
post_load_complete=True,
)
loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
@@ -90,7 +90,7 @@ def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
assert first.plan is not None
assert second.plan is not None
assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
assert first.plan.post_load_complete
assert not first.plan.requires_backend
assert {binding.file_path for binding in first.plan.bindings.values()} == {
binding.file_path for binding in second.plan.bindings.values()
}