[Refactor][Diffusion] Complete generic backend plan consumption (#7313)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Anjie HouandClaude Opus 5 authored and GitHub committed 2026-09-15 08:40:04 +08:00
1 parent 5385dda4fe
commit 85da943391
7 files changed
+228 -56

No files matched your search

+29 -4
View File
@@ -41,7 +41,8 @@ Every backend implements `OffloadBackend`:
- `is_enabled()` reports lifecycle state.
`disable()` does not promise to restore the pipeline's original device
placement. The caller owns any subsequent rematerialization.
placement. It must restore usable parameter and buffer storage before releasing
backend-owned host weights. The caller owns subsequent device placement.
An enable failure must remove every partially installed hook. Rank-local
backends restore ordinary tensors so the pipeline can be retried. A failed
@@ -54,9 +55,19 @@ must use distinct hook names and remove only hooks they own.
## Discovery and topology
`resolve_offload_plan(pipeline, config)` is the single topology entry point.
Every backend calls it once when `enable()` starts and then consumes the
resolved components, block stacks, resident heads, and staged-residency flags
it returns. Backends do not re-read model declarations.
The generic model-level and ordinary layerwise paths call it once when
`enable()` starts and consume its component selection and block stacks.
Ordinary layerwise offload also respects a resolved `skip_reason`, preserving
the legacy missing-DiT no-op without placing auxiliary components. Non-block
DiT and encoder state is placed by resolved block tensor identity, so the
backend does not rediscover block containers by attribute name.
Two compatibility paths remain outside this completed generic cutover:
model-level pipelines implementing `SupportsModelCpuOffload` still own their
custom phase lifecycle (RFC #6648 J4), and distributed layerwise offload still
reads declarations for nested-submodule staging and loader-owned mmap paths
(J3). The custom model-level protocol is delegated to before generic plan
resolution; it must not be subjected to the generic encoder/DiT swap checks.
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
@@ -102,3 +113,17 @@ remains responsible for transfer, synchronization, and storage ownership.
The diffusion [Offloader module design](../../module/diffusion/offloader.md)
describes how these feature contracts fit into the larger diffusion runtime.
## Backend acceptance tests
`tests/diffusion/offloader/test_backend_plan_contract.py` covers the two
behaviors the generic cutover changed. A DiT attribute aliasing a streamed
block must keep the ring's host residency, so placement cannot follow attribute
names; the CUDA case fails against the previous name-based placement. Both
backends must then run with their selector helpers and the model declaration
made unreadable, proving execution consumes the resolved plan alone.
Selection errors, rollback, residency, and enable/disable cycles stay in
`test_plan_resolver.py`, `test_layerwise_backend.py` and
`test_sequential_backend.py`. Model-specific phases and executed tied aliases
remain the J4 contract suite.