[Refactor][Diffusion] Resolve offload topology through one plan resolver (#7209)

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-09 08:17:38 +08:00
1 parent b3dd45874a
commit b6f4df320d
11 files changed
+976 -291

No files matched your search

+24 -9
View File
@@ -53,6 +53,11 @@ 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.
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
- `_dit_modules`;
@@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
- `_resident_modules`.
Dotted paths are supported. Legacy pipelines may use the fallback scan of
well-known attribute names, but new integrations should declare components
explicitly.
well-known attribute names; it now warns once per pipeline class, and new
integrations must declare components explicitly.
Both layerwise backends first consume pipeline `OffloadPlan` metadata. A
plan's `block_attrs` maps each DiT path to its ordered block containers, while
`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the
plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
singular-name compatibility path). Discovery metadata describes structure
only; the backend remains responsible for transfer, synchronization, and
storage ownership.
Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps
each DiT path to its ordered block containers, `encoder_block_attrs` declares
streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed
residency, `resident_dit_paths` marks the DiTs that may hold resident layers,
and `encoder_dlo_weight_replication` marks the encoders whose loader-produced
weights are safe for AllGather. DiTs absent from the plan fall back to
`_layerwise_offload_blocks_attrs` (including the deprecated singular-name
compatibility path).
The resolver is pure: it moves no tensor, installs no hook, writes no module
attribute, and reads no process group — multi-rank facts come from
`OffloadConfig`. It owns topology validation, so a configuration the model
cannot serve fails before the first component is placed. Explicit component
selection turns every declaration problem into an error, while the
compatibility topology (no component selector) keeps the historical
warn-and-skip behavior. Metadata still describes structure only; the backend
remains responsible for transfer, synchronization, and storage ownership.
## Cross-strategy invariants