[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:
1 parent
b3dd45874a
commit
b6f4df320d
11 files changed
+976
-291
No files matched your search
@@ -20,6 +20,12 @@ from .distributed_layerwise_backend import (
|
||||
from .layerwise_backend import LayerWiseOffloadBackend
|
||||
from .module_residency import BoundedAllocatorCache, PinnedModuleStager
|
||||
from .offload_plan import OffloadPlan, get_offload_plan
|
||||
from .plan_resolver import (
|
||||
BlockStack,
|
||||
ResolvedComponent,
|
||||
ResolvedOffloadPlan,
|
||||
resolve_offload_plan,
|
||||
)
|
||||
from .sequential_backend import (
|
||||
ModelLevelOffloadBackend,
|
||||
apply_sequential_offload,
|
||||
@@ -42,6 +48,10 @@ __all__ = [
|
||||
"OffloadConfig",
|
||||
"OffloadPlan",
|
||||
"OffloadStrategy",
|
||||
"BlockStack",
|
||||
"ResolvedComponent",
|
||||
"ResolvedOffloadPlan",
|
||||
"resolve_offload_plan",
|
||||
"SupportsModelCpuOffload",
|
||||
"LayerWiseOffloadBackend",
|
||||
"DistributedLayerwiseOffloadBackend",
|
||||
|
||||
Reference in new issue
Block a user