diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 0ec81f3b5..52797b68c 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -84,10 +84,22 @@ 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). +`encoder_dlo_weight_replication` marks the encoders whose loader-produced +weights are safe for AllGather, and `offload_submodules` maps a DiT submodule +that owns its residency to its own block container. DiTs absent from the plan +fall back to `_layerwise_offload_blocks_attrs` (including the deprecated +singular-name compatibility path). + +An undeclared DiT submodule joins the plan only when it is large enough that +keeping it resident would defeat streaming; its block container is then found +by scanning well-known attribute names, which warns once per submodule class. +Size is computed from parameter shapes and dtypes, including meta parameters, +so resolving before mmap loading preserves the same residency decision. +Block ownership validation includes nested components: a block shared with a +parent, sibling, or encoder is rejected before any hooks or storage are changed. +A submodule with no block container must implement the `load_to_device` / +`offload_to_cpu` lifecycle, and the resolver rejects it before any component of +the pipeline is placed or hooked. 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 diff --git a/tests/diffusion/models/magi2/test_native_dlo_allgather.py b/tests/diffusion/models/magi2/test_native_dlo_allgather.py index 9bea5dc30..5cfaabf3f 100644 --- a/tests/diffusion/models/magi2/test_native_dlo_allgather.py +++ b/tests/diffusion/models/magi2/test_native_dlo_allgather.py @@ -36,6 +36,7 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, DistributedLayerwiseOffloadHook, ) +from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent @dataclass(frozen=True) @@ -46,12 +47,6 @@ class _WeightSourceStub: prefix: str -@dataclass(frozen=True) -class _PipelineModulesStub: - dits: list[nn.Module] - dit_names: list[str] - - _WORLD_SIZE = 4 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -189,7 +184,7 @@ def _load_mmap_transform_and_reconstruct( backend._mmap_transforms_by_tensor_id = {} backend._load_weights_via_mmap( pipeline, - _PipelineModulesStub(dits=[target], dit_names=["transformer"]), + (ResolvedComponent(path="transformer", module=target, selected=True),), plan_result.plan, ) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index c24639434..fed60be1f 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -32,6 +32,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( TensorBinding, build_checkpoint_mmap_plan, ) +from vllm_omni.diffusion.offloader import plan_resolver from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy from vllm_omni.diffusion.offloader.block_discovery import ( get_blocks_attr_names, @@ -52,7 +53,7 @@ from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, ) -from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan +from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform @@ -1235,6 +1236,78 @@ class _MmapPostLoadPipeline(nn.Module): class TestMmapWeightLoading: + @pytest.mark.parametrize("initial_device", ["cpu", "meta"]) + @pytest.mark.parametrize("use_allgather", [False, True]) + def test_mmap_keeps_undeclared_nested_blocks_streamed( + self, tmp_path, patched_offload_runtime, monkeypatch, initial_device, use_allgather + ): + class Refiner(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + + def forward(self, x): + for block in self.layers: + x = block(x) + return x + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + self.refiner = Refiner() + + def forward(self, x): + x = self.refiner(x) + for block in self.blocks: + x = block(x) + return x + + pipeline = nn.Module() + pipeline.transformer = Transformer() + weights = {name: param.detach().clone() for name, param in pipeline.named_parameters()} + x = torch.randn(1, 2) + with torch.no_grad(): + expected = pipeline.transformer(x) + checkpoint_file = tmp_path / "model.safetensors" + save_file(weights, str(checkpoint_file)) + # Exercise the real mmap path with a small model instead of allocating + # the 1 GiB needed to trigger the production compatibility threshold. + monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0) + if initial_device == "meta": + pipeline.transformer.to_empty(device="meta") + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={name: TensorBinding(checkpoint_key=name, file_path=str(checkpoint_file)) for name in weights}, + ) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=use_allgather, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + + for _ in range(2): + try: + backend.enable(pipeline) + assert len(backend._all_hook_groups) == 2 + with torch.no_grad(): + torch.testing.assert_close(pipeline.transformer(x), expected) + finally: + backend.disable() + for name, parameter in pipeline.named_parameters(): + assert not parameter.is_meta + torch.testing.assert_close(parameter, weights[name]) + for module in pipeline.modules(): + registry = getattr(module, "_hook_registry", None) + if registry is not None: + assert registry.get_hook("distributed_layerwise_offload") is None + def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime): pipeline = _MmapPostLoadPipeline() weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()} @@ -1250,10 +1323,7 @@ class TestMmapWeightLoading: ), torch.device("cpu"), ) - modules = SimpleNamespace( - dits=[pipeline.transformer], - dit_names=["transformer"], - ) + dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),) plan = HostWeightPlan( backing_kind="checkpoint_mmap", bindings={ @@ -1265,7 +1335,7 @@ class TestMmapWeightLoading: }, ) - backend._load_weights_via_mmap(pipeline, modules, plan) + backend._load_weights_via_mmap(pipeline, dits, plan) assert pipeline.transformer.post_load_calls == 1 assert pipeline.transformer.time_embedder.weight.dtype == torch.float32 @@ -1717,6 +1787,128 @@ class TestOffloadPlan: for name, parameter in pipeline.named_parameters(): torch.testing.assert_close(parameter, expected_parameters[name]) + @pytest.mark.parametrize("shared_owner", ["parent", "encoder", "sibling"]) + def test_nested_duplicate_ownership_is_rejected_before_any_hook(self, patched_offload_runtime, shared_owner): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + offload_submodules={"refiner": "blocks", "other_refiner": "blocks"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.transformer.refiner = _SingleBlockModel(num_blocks=2) + self.transformer.other_refiner = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + owners = { + "parent": self.transformer.blocks, + "encoder": self.text_encoder.encoder.block, + "sibling": self.transformer.other_refiner.blocks, + } + self.transformer.refiner.blocks = owners[shared_owner] + + pipeline = Pipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False), + torch.device("cpu"), + ) + storage = {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()} + try: + with pytest.raises(ValueError, match="claimed by both"): + backend.enable(pipeline) + assert not backend.enabled + assert not backend._all_hook_groups + assert {name: parameter.data_ptr() for name, parameter in pipeline.named_parameters()} == storage + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) + for module in pipeline.modules(): + assert getattr(module, "_hook_registry", None) is None + finally: + backend.disable() + + def test_unstageable_submodule_is_rejected_before_any_hook(self, patched_offload_runtime): + """Plan-dependent rejection happens before earlier components are hooked.""" + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + offload_submodules={"refiner": "missing"}, + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + # Declared for submodule offload but without a block list or the + # load_to_device/offload_to_cpu lifecycle. + self.transformer.refiner = nn.Linear(2, 2) + self.text_encoder = _PlainEncoder() + + pipeline = Pipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False), + torch.device("cpu"), + ) + + # Storage identity proves the encoder was never staged and rolled back. + encoder_storage = [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block] + + with pytest.raises(ValueError, match="must implement load_to_device"): + backend.enable(pipeline) + + assert not backend.enabled + assert [block.weight.data_ptr() for block in pipeline.text_encoder.encoder.block] == encoder_storage + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) + for block in pipeline.transformer.blocks: + assert getattr(block, "_hook_registry", None) is None + + def test_undeclared_submodule_block_scan_warns(self, patched_offload_runtime, monkeypatch): + """The size-triggered attribute scan still works and is deprecated.""" + + class Refiner(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + self.refiner = Refiner() + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + warnings: list[str] = [] + + class _Recorder: + def warning(self, message, *args): + warnings.append(message % args if args else message) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + monkeypatch.setattr(plan_resolver, "logger", _Recorder()) + monkeypatch.setattr(plan_resolver, "_NESTED_OFFLOAD_THRESHOLD_MB", 0) + plan_resolver._warn_nested_block_scan.cache_clear() + plan_resolver._warn_legacy_discovery.cache_clear() + pipeline = Pipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert len(backend._all_hook_groups) == 2 + assert sum("OffloadPlan.offload_submodules" in message for message in warnings) == 1 + + backend.disable() + class TestMmapValidation: """Tests for loader preflight and backend plan realization.""" @@ -2101,7 +2293,7 @@ class TestMmapValidation: ), torch.device("cpu"), ) - modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + dits = (ResolvedComponent(path="transformer", module=pipeline.transformer, selected=True),) result = build_checkpoint_mmap_plan( pipeline, dit_modules=(("transformer", pipeline.transformer),), @@ -2113,7 +2305,7 @@ class TestMmapValidation: ) assert result.plan is not None - backend._load_weights_via_mmap(pipeline, modules, result.plan) + backend._load_weights_via_mmap(pipeline, dits, result.plan) assert pipeline.transformer.validate_called, "validate_loaded_weights should be called" diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 299ec3d3a..e79d1f197 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -53,8 +53,7 @@ from .host_registration import ( HostRegistrationError, register_host_mappings, ) -from .offload_plan import OffloadPlan -from .plan_resolver import ResolvedComponent, resolve_offload_plan +from .plan_resolver import BlockStack, ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -111,7 +110,6 @@ def _forget_active_hwr_registration( # use layerwise offload (streaming hooks) or be moved to GPU as a resident # module. Submodules larger than this are offloaded to save HBM; smaller # ones stay resident for lower latency. -_ON_DEMAND_THRESHOLD_MB = 1024 class DistributedLayerwiseOffloadHook(ModelHook): @@ -1182,7 +1180,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _load_weights_via_mmap( self, pipeline: nn.Module, - modules, + dits: tuple[ResolvedComponent, ...], plan: HostWeightPlan, ) -> None: """Load DiT checkpoint tensors as file-backed safetensors views. @@ -1226,7 +1224,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # converts the child's buffers too. So we must save ALL buffers # from ALL DiT modules BEFORE any to_empty call. saved_buffers: dict[int, dict[str, torch.Tensor]] = {} - for dit_module in modules.dits: + for dit_module in (component.module for component in dits): bufs: dict[str, torch.Tensor] = {} for name, buf in dit_module.named_buffers(): # Check if this buffer is non-persistent on its OWNING module @@ -1246,7 +1244,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Skip modules that are already meta (happens when a parent DiT # module contains a child DiT module — to_empty on the parent # already converted the child). - for dit_module in modules.dits: + for dit_module in (component.module for component in dits): if any(p.is_meta for p in dit_module.parameters()): logger.info( "%s already on meta device (skipping to_empty, %d buffers saved)", @@ -1307,7 +1305,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # assigning checkpoint tensors. The mmap path bypasses that loader, so # preserve the same lifecycle for transforms such as Cosmos3's fp32 # timestep embedder. - for dit_name, dit_module in zip(modules.dit_names, modules.dits): + for dit_name, dit_module in ((component.path, component.module) for component in dits): # Restore non-persistent buffers before post-load hooks and strict # validation. They are constructor-derived and intentionally have # no checkpoint binding. @@ -1352,7 +1350,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._mmap_transforms_by_tensor_id[id(target)] = binding.transform remaining_meta: list[str] = [] - for dit_name, dit_module in zip(modules.dit_names, modules.dits): + for dit_name, dit_module in ((component.path, component.module) for component in dits): remaining_meta.extend( f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta ) @@ -1517,73 +1515,26 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) return True - def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool: - """Try to apply layerwise offload to a large submodule's blocks. - Resolution order: - 1. OffloadPlan.offload_submodules (declarative, if plan is provided) - 2. Heuristic search for common block-list attributes - - Returns True if layerwise offload was applied, False otherwise. - """ - from operator import attrgetter - - blocks = None - blocks_attr = None - - # 1. Check OffloadPlan first (declarative — no guessing) - if plan is not None and name in plan.offload_submodules: - attr_name = plan.offload_submodules[name] - try: - candidate = attrgetter(attr_name)(module) - if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: - blocks = candidate - blocks_attr = attr_name - except AttributeError: - logger.warning( - "OffloadPlan declared block attr '%s' for submodule '%s' " - "but attribute not found — falling back to heuristic", - attr_name, - name, - ) - - # 2. Fallback: heuristic search - if blocks is None: - for attr_name in ("layers", "blocks", "h", "model.layers"): - try: - candidate = attrgetter(attr_name)(module) - except AttributeError: - continue - if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: - blocks = candidate - blocks_attr = attr_name - break - - if blocks is None: - return False - + def _offload_nested_submodule(self, component: ResolvedComponent) -> None: + """Stream one resolved DiT submodule through its own hook ring.""" + stack = component.stacks[0] logger.info( "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, group_size=%d)", - name, - blocks_attr, - len(blocks), - sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576, + component.path, + stack.attrs[0], + len(stack.blocks), + sum(p.nelement() * p.element_size() for p in component.module.parameters()) / 1048576, self._component_transport(DIT_COMPONENT)[1], ) - - # Move non-block parts of the submodule to GPU (small: embeddings, norms) - for child_name, child in module.named_children(): - if child_name != blocks_attr: - child.to(self.device) - - self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) - return True + # Keep the submodule's own non-block state resident by block identity, + # so an attribute aliasing a streamed block is not placed with it. + move_non_block_state_to_device(component.module, (stack.blocks,), self.device) + self._install_hook_group(stack.blocks, DIT_COMPONENT, use_dit_mmap=True) def _prepare_dit_non_block_modules( self, - dit_module: nn.Module, - blocks_attr_names: list[str], - all_dit_modules: set[int], - plan: OffloadPlan | None, + component: ResolvedComponent, + stack: BlockStack, ) -> None: """Place or hook the DiT parts that are outside its repeated blocks. @@ -1591,28 +1542,23 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): an all-resident stage skips placement for modules such as H3's token refiner and enters the forward pass with CPU or meta tensors. """ + dit_module = component.module + nested = {child.path.rsplit(".", 1)[-1]: child for child in component.children} for name, module in dit_module.named_children(): - if name in blocks_attr_names: + if name in stack.attrs: logger.debug("Skipped blocks module %s", name) continue - module_mb = ( - sum( - param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0 - for param in module.parameters() - ) - / 1048576 - ) - explicitly_planned = plan is not None and name in plan.offload_submodules - if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD_MB: - if id(module) in all_dit_modules: + child = nested.get(name) + if child is not None: + if not child.selected: logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) - elif self._try_layerwise_offload_submodule(module, name, plan): - pass + elif child.stacks: + self._offload_nested_submodule(child) else: prepare_component( - module, - name, + child.module, + child.path, device=self.device, stage_on_demand=True, blockwise=False, @@ -1715,7 +1661,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): elif host_weight_plan.backing_kind == "checkpoint_mmap": self._load_weights_via_mmap( pipeline, - resolved.modules, + resolved.dits, host_weight_plan, ) else: @@ -1750,16 +1696,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): enable_encoder_blocks=self._try_layerwise_offload_encoder, ) - if self.config.offloads(DIT_COMPONENT): + if any(component.selected for component in resolved.dits): logger.info( "Applying distributed layer-wise offloading on %s", [component.path for component in resolved.dits], ) - # Collect all DiT module objects to detect submodules that are - # already handled as a separate DiT module (avoids duplicate hooks). - all_dit_modules = set(id(component.module) for component in resolved.dits) - # Apply hooks for each DiT module for component, stack in iter_streamable_dits(resolved, self.device): streaming = list(stack.streaming) @@ -1772,12 +1714,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): len(streaming), ) - self._prepare_dit_non_block_modules( - component.module, - list(stack.attrs), - all_dit_modules, - resolved.declaration, - ) + self._prepare_dit_non_block_modules(component, stack) if not streaming: logger.info("All blocks for %s are resident; no streaming hooks required", component.path) @@ -1804,11 +1741,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.enabled = bool(self._resident_blocks or self._encoder_modules or self._staged_components) if self._using_mmap and not self.enabled: self._release_mmap_handles() - if not self.enabled and not self.config.offloads(DIT_COMPONENT): - raise ValueError( - "None of the selected distributed layerwise offload components have " - "a model-declared streamable or on-demand plan" - ) return # Unified allocation: 2 shared output buffers + 2 shared shard buffers diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py index 08527318f..ca32f6d3f 100644 --- a/vllm_omni/diffusion/offloader/plan_resolver.py +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -17,6 +17,7 @@ from __future__ import annotations from dataclasses import dataclass from functools import cache +from operator import attrgetter from typing import TYPE_CHECKING from torch import nn @@ -27,7 +28,7 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from .block_discovery import get_blocks_from_dit from .component_utils import get_encoder_block_groups, validate_on_demand_component from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy -from .module_collector import ModuleDiscovery, PipelineModules +from .module_collector import ModuleDiscovery from .offload_plan import OffloadPlan, get_offload_plan if TYPE_CHECKING: @@ -76,6 +77,9 @@ class ResolvedComponent: selected: bool stacks: tuple[BlockStack, ...] = () on_demand: bool = False + # Submodules of a DiT that are large enough to own their residency instead + # of staying resident with the rest of its non-block state. + children: tuple[ResolvedComponent, ...] = () @dataclass(frozen=True) @@ -86,11 +90,6 @@ class ResolvedOffloadPlan: encoders: tuple[ResolvedComponent, ...] vaes: tuple[ResolvedComponent, ...] residents: tuple[ResolvedComponent, ...] - # Raw discovery output and the model declaration remain reachable for the - # loader-owned mmap path and for DLO's nested-submodule staging, which - # still interpret topology themselves. - modules: PipelineModules - declaration: OffloadPlan | None = None # Ordinary layerwise offload turns a legacy selection it cannot serve into # a complete no-op, placement included. Distributed layerwise offload warns # and keeps preparing its other components instead, so this stays scoped to @@ -102,6 +101,22 @@ class ResolvedOffloadPlan: return self.dits + self.encoders + self.vaes + self.residents +# A DiT submodule this large owns its residency even when no plan declares it. +_NESTED_OFFLOAD_THRESHOLD_MB = 1024 +_NESTED_BLOCK_ATTRS = ("layers", "blocks", "h", "model.layers") + + +@cache +def _warn_nested_block_scan(module_class: str, attr: str) -> None: + """Warn once per class that a nested block list was found by guessing.""" + logger.warning( + "%s block list was discovered as %r by scanning well-known attribute names. " + "Declare it in OffloadPlan.offload_submodules; the attribute scan is deprecated.", + module_class, + attr, + ) + + @cache def _warn_legacy_discovery(pipeline_class: str) -> None: """Warn once per pipeline class that component discovery guessed. @@ -117,6 +132,77 @@ def _warn_legacy_discovery(pipeline_class: str) -> None: ) +def _nested_block_stack(module: nn.Module, declared_attr: str | None) -> BlockStack | None: + """Resolve one DiT submodule's own block ring, if it has one.""" + candidates = (declared_attr,) if declared_attr is not None else _NESTED_BLOCK_ATTRS + for attr in candidates: + try: + blocks = attrgetter(attr)(module) + except AttributeError: + continue + if isinstance(blocks, nn.ModuleList) and len(blocks) > 1: + if declared_attr is None: + _warn_nested_block_scan(type(module).__name__, attr) + return BlockStack(attrs=(attr,), blocks=tuple(blocks)) + if declared_attr is not None: + logger.warning( + "OffloadPlan declared block attr %r for submodule %s but it is not a streamable block list", + declared_attr, + type(module).__name__, + ) + return None + + +def _resolve_dit_children( + dit_module: nn.Module, + dit_path: str, + stack: BlockStack, + declaration: OffloadPlan | None, + dit_module_ids: set[int], +) -> tuple[ResolvedComponent, ...]: + """Resolve the DiT submodules that own their residency. + + A submodule qualifies when the model declares it or when it is large enough + that keeping it resident would defeat streaming. Submodules that are + themselves discovered DiTs are left to their own component. + """ + children: list[ResolvedComponent] = [] + for name, module in dit_module.named_children(): + if name in stack.attrs: + continue + declared_attr = None if declaration is None else declaration.offload_submodules.get(name) + if declared_attr is None and _module_size_mb(module) <= _NESTED_OFFLOAD_THRESHOLD_MB: + continue + path = f"{dit_path}.{name}" + if id(module) in dit_module_ids: + # Its own component owns this submodule; the parent must not place it. + children.append(ResolvedComponent(path=path, module=module, selected=False)) + continue + nested_stack = _nested_block_stack(module, declared_attr) + if nested_stack is None: + # Without its own ring the submodule needs the pipeline lifecycle. + validate_on_demand_component(module, path) + children.append( + ResolvedComponent( + path=path, + module=module, + selected=True, + stacks=() if nested_stack is None else (nested_stack,), + on_demand=nested_stack is None, + ) + ) + return tuple(children) + + +def _module_size_mb(module: nn.Module) -> float: + """Expected parameter bytes in MiB, including not-yet-loaded meta tensors. + + Resolution precedes mmap materialization, so residency must depend on the + tensor metadata rather than whether the loader has allocated storage yet. + """ + return sum(parameter.nelement() * parameter.element_size() for parameter in module.parameters()) / 1048576 + + def _resolve_dit_stacks( module: nn.Module, path: str, @@ -253,11 +339,25 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved vaes: list[ResolvedComponent] = [] residents: list[ResolvedComponent] = [] + dit_module_ids = {id(module) for module in modules.dits} for path, module in zip(modules.dit_names, modules.dits): stacks: tuple[BlockStack, ...] = () + children: tuple[ResolvedComponent, ...] = () if dit_selected and layerwise: stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit) - dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks)) + # Only distributed layerwise offload streams or stages submodules; + # the ordinary backend keeps all non-block state resident. + if stacks and config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE: + children = _resolve_dit_children(module, path, stacks[0], declaration, dit_module_ids) + dits.append( + ResolvedComponent( + path=path, + module=module, + selected=dit_selected, + stacks=stacks, + children=children, + ) + ) for path, module in zip(modules.encoder_names, modules.encoders): selected = config.should_offload_encoder(path, declaration) @@ -307,8 +407,6 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved encoders=tuple(encoders), vaes=tuple(vaes), residents=tuple(residents), - modules=modules, - declaration=declaration, skip_reason=( "No DiT/transformer modules found for selected DiT layerwise offload" if config.strategy is OffloadStrategy.LAYER_WISE and dit_selected and not dits @@ -322,7 +420,10 @@ def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> Resolved def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None: """Reject topologies where two components would hook the same blocks.""" owner_by_block: dict[int, str] = {} - for component in resolved.components: + pending = list(reversed(resolved.components)) + while pending: + component = pending.pop() + pending.extend(reversed(component.children)) for block in (block for stack in component.stacks for block in stack.blocks): owner = owner_by_block.setdefault(id(block), component.path) if owner != component.path: