[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:
1 parent
5385dda4fe
commit
85da943391
7 files changed
+228
-56
No files matched your search
@@ -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.
|
||||
Reference in new issue
Block a user