diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md index da735d60b..9a7a7e00b 100644 --- a/docs/design/feature/distributed_layerwise_offload.md +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -5,6 +5,8 @@ models. DLO keeps only a small number of DiT blocks on the accelerator and streams the remaining blocks from host memory. The distributed backend can either shard those host-side weights across an existing parallel group or keep the standard loader's rank-local weights and avoid an additional collective. +The rank-local mmap storage design is tracked in +[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). For user-facing commands, see the [distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md) @@ -16,8 +18,10 @@ DLO is implemented for multi-device diffusion execution. The default AllGather path is the primary path for DP and SP deployments. The `--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful for standard-loader sharding, workstation bring-up, and systems where an -additional DLO collective is undesirable, but it does not reduce host weight -storage across ranks. +additional DLO collective is undesirable. For explicitly mmap-compatible +models in a TP1, non-HSDP topology, this path shares file-backed DiT checkpoint +pages across processes and keeps only bounded staging buffers private. TP, +HSDP, and other models retain standard-loader rank-local host storage. The compatibility matrix below describes the current implementation. The unit-level guards are covered, but not every parallelism combination has a @@ -68,9 +72,18 @@ requests must take the same execution path at every denoising step. ### Rank-local path without DLO AllGather With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to -one. The regular model loader remains responsible for preparing each rank's -weights, including TP-local tensors or HSDP-managed parameters. DLO then -streams those rank-local tensors block by block using H2D copies only. +one and streams complete runtime blocks using H2D copies only. Host storage is +selected independently from that transfer protocol: + +- For an explicitly mmap-compatible model with TP1 and no HSDP, each worker + retains immutable safetensors views backed by the node's shared OS page + cache. Two pinned host buffers, sized to the largest streamed block, stage + mmap pages for asynchronous H2D. Checkpoint-to-runtime adapters run while a + block is packed, so transformed weights do not become a persistent private + model copy. +- With TP, HSDP, online quantization, or a model without mmap support, the + regular loader remains responsible for preparing rank-local tensors. This + preserves TP/HSDP/quantization loader callbacks. This mode means: @@ -80,19 +93,26 @@ This mode means: not shard weights across SP ranks. - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only DLO's additional weight AllGather is disabled. -- Pure DP deployments keep a full host-side model copy per rank, subject to - shared OS page-cache behavior. +- Pure DP deployments with a compatible mmap checkpoint share the streamed + DiT's file-backed pages. Other components and fallback paths may still keep + rank-private host weights. - The scheduler does not require a synchronized DP request wave for DLO. -This path is intentionally not implemented by reusing the AllGather mmap -loader. It relies on the standard loader so model-specific TP/HSDP transforms -remain intact. +The mmap source remains immutable in both transfer modes. AllGather copies a +persistent `1 / group_size` shard from it and then closes the mappings; +rank-local mode keeps the mappings open and copies a complete block through +bounded staging without adding a collective. + +The current implementation packs mmap views into a staging slot synchronously +on the host; the subsequent pinned H2D transfer overlaps with block compute. +This keeps host memory bounded but makes CPU packing bandwidth part of the +rank-local path's performance profile. ## Parallelism compatibility | Parallelism | DLO + AllGather | DLO without AllGather | |---|---|---| -| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. | +| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight collective. Compatible TP1 models share file-backed DiT pages across replicas. | | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | | **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. | | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | @@ -122,11 +142,12 @@ AllGather DP multi-concurrency requires: The no-AllGather path does not impose these DLO-specific synchronized-wave requirements. -The mmap loader is used only by the supported DLO+AllGather path when the -model has a compatible checkpoint layout. Online quantization is incompatible -with that sharded mmap path; use `--dlo-no-use-allgather` or disable online -quantization. Models that do not meet the mmap requirements use the regular -loader path. +The mmap loader is used when the model explicitly declares a compatible +checkpoint layout. It is independent of AllGather for TP1/non-HSDP execution. +Models can declare a per-parameter checkpoint-to-runtime adapter; the adapter +must preserve dtype and element count so it can be applied while packing one +staging block. MiniMax-H3 uses this mechanism for grouped QKV. Online +quantization and TP/HSDP no-AllGather execution use the regular loader path. ## Validation coverage @@ -136,6 +157,8 @@ Current source-level validation includes: - HSDP + DLO without AllGather acceptance at configuration level; - TP rejection in the DLO+AllGather mmap path; - resident-layer requests requiring no-AllGather; +- rank-local mmap source retention, layout transforms, and bounded host + staging; - DP request-wave validation for denoising-step compatibility; - sharding, double-buffer, AllGather-size, and heterogeneous-block regression tests. @@ -150,7 +173,8 @@ on the target CUDA/NCCL or CANN/HCCL hardware. scaling path. - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is not the goal. -- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with - the expectation of higher host-memory use and lower validation confidence. +- Use **no-AllGather** when replicas must avoid DLO collectives. TP1 models with + mmap support can share DiT checkpoint pages; TP/HSDP and unsupported models + retain the higher-memory standard-loader path. - Prefer **HSDP alone** for production HSDP deployments until the combined HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index 02ab9f074..34b5a24ac 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -189,7 +189,10 @@ multi-device deployments. In the default mode, each rank stores only reconstructed at runtime via **AllGather** on a dedicated communication stream, overlapped with computation via a fixed double-buffer scheme. The `--dlo-no-use-allgather` mode instead streams standard-loader rank-local -weights independently and does not shard weights across ranks. +weights independently and does not shard weights across ranks. For explicitly +mmap-compatible TP1 models, the same mode can retain node-shared file-backed +DiT weights and use bounded per-worker host staging instead of a private full +DiT copy. For the implementation design and compatibility matrix, see [Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md). @@ -201,8 +204,9 @@ For the implementation design and compatibility matrix, see time, regardless of model size - **DP multi-concurrency**: N concurrent requests processed in parallel (AllGather only gathers weight shards, which are request-independent) -- **mmap weight loading**: weights loaded as mmap views pointing to shared OS - page cache, eliminating O(dp_size × model_size) RSS during model creation +- **mmap weight storage**: compatible checkpoints use views into the shared OS + page cache. AllGather copies persistent rank shards; rank-local mode keeps + the views and packs complete blocks through two bounded host staging slots - **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU (CANN/HCCL) via vLLM-Omni's platform abstraction @@ -224,7 +228,7 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 -# Without DLO AllGather (each rank streams standard-loader rank-local weights) +# Without DLO AllGather (compatible TP1 models use node-shared mmap storage) vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --data-parallel-size 4 \ @@ -258,20 +262,26 @@ m = Omni( ### How model weights are loaded (mmap path) -When DLO + AllGather is active, the offloader: +When DLO and an mmap-compatible model are active, the offloader: 1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`) from the normally-created transformer 2. **Converts to meta device** via `to_empty(device="meta")`, releasing random initialization weights 3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`, - which return views into the OS page cache (shared across all ranks, 0 RSS) -4. **Calls `post_load_weights()`** to apply model-specific dtype conversions - (e.g. Cosmos3's `time_embedder` → FP32) + which return views into the OS page cache without eagerly copying the full + tensor into anonymous process memory +4. **Applies declared per-parameter layout adapters while packing**, when the + checkpoint and TP1 runtime layouts differ (for example MiniMax-H3 grouped + QKV) 5. **Restores non-persistent buffers** from saved copies -This approach requires **zero model-specific code changes** — no pipeline or -transformer modifications are needed. +With AllGather, each rank copies only its persistent shard and then releases +the source mappings. Without AllGather, mappings remain open and every worker +owns two pinned staging slots sized to the largest streamed block. The latter +shares file-backed pages, not process-local Python objects, and does not alter +request routing or scheduling. Packing into the host slot is synchronous; the +pinned H2D copy can overlap with computation. ### OffloadPlan (declarative topology metadata) @@ -313,7 +323,7 @@ request while AllGather synchronizes only weight shards (request-independent). `--dlo-no-use-allgather` or disable online quantization. - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware weight-loader callbacks are bypassed. The no-AllGather path retains the - standard TP loader and is experimental. + standard TP loader instead of rank-local mmap and is experimental. - HSDP + AllGather is rejected because it would double-shard parameters. HSDP with no-AllGather is accepted at configuration level, but has limited end-to-end validation. @@ -373,6 +383,7 @@ reasoner/generator components inside the model forward pass. | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` | | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | +| MiniMaxH3Pipeline | `MiniMaxAI/MiniMax-H3` (`FL2VA`) | `MiniMaxH3DiTModel` | - | ✓ | ✓ | `"blocks"` | | SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` | | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` | | Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` | @@ -382,8 +393,9 @@ reasoner/generator components inside the model forward pass. - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks - Distributed Layerwise Offloading reuses the layerwise block topology, but the model and parallelism combination still needs validation. AllGather mode - additionally requires a compatible checkpoint/mmap layout; no-AllGather mode - retains the standard loader's rank-local tensors. See the + additionally requires a compatible checkpoint/mmap layout. For explicitly + compatible TP1 models, no-AllGather uses node-shared mmap sources; other + layouts retain the standard loader's rank-local tensors. See the [DLO design](../../design/feature/distributed_layerwise_offload.md) and [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples. diff --git a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py index f8117b0f9..2cac7a666 100644 --- a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py +++ b/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py @@ -11,6 +11,13 @@ Examples: python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2 + # Rank-local mmap storage, with no DLO collective or DP wave scheduling. + VLLM_WORKER_MULTIPROC_METHOD=spawn \ + VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ + CUDA_VISIBLE_DEVICES=0,1 \ + python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ + --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2-no-allgather + VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ CUDA_VISIBLE_DEVICES=0,1 \ @@ -45,9 +52,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA") parser.add_argument( "--mode", - choices=("dlo-dp2", "request"), + choices=("dlo-dp2", "dlo-dp2-no-allgather", "request"), required=True, - help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode", + help="DLO DP2 (with or without AllGather) or ordinary non-DLO TP2 request mode", ) parser.add_argument("--steps", type=int, default=2) parser.add_argument("--duration", type=float, default=5.0) @@ -74,14 +81,14 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: "stage_init_timeout": args.init_timeout, "init_timeout": args.init_timeout, } - if args.mode == "dlo-dp2": + if args.mode in {"dlo-dp2", "dlo-dp2-no-allgather"}: common.update( tensor_parallel_size=1, data_parallel_size=2, text_encoder_tp_size=1, vae_patch_parallel_size=1, enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + dlo_use_allgather=args.mode == "dlo-dp2", dlo_resident_layers=0, ) else: diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index b9ae3335c..58a304cfc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -191,6 +191,30 @@ def test_model_load_weights_transforms_before_calling_vllm_loader(): ] +def test_unquantized_qkv_declares_equivalent_mmap_transform(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + + arch = h3.MiniMaxH3DiTArchConfig( + hidden_size=1, + num_attention_heads=2, + attention_head_dim=1, + rope_inv_freq_len=1, + ) + attention = h3.MiniMaxH3Attention( + arch, + quant_config=None, + prefix="blocks.0.attn", + ) + transform = attention.qkv_proj.weight.mmap_weight_transform + checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1) + + assert transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5] + + def test_pipeline_resolves_transformer_component_quant_config(): from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( _resolve_component_quant_config, diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index fd004d056..4e303cea7 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -30,7 +30,12 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( PinnedResidentLayerGroup, ) from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager -from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan +from vllm_omni.diffusion.offloader.offload_plan import ( + OffloadPlan, + get_offload_plan, + should_select_dlo_mmap, + supports_mmap_loading, +) from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -48,6 +53,9 @@ class DummyEvent: def record(self, _stream) -> None: return None + def synchronize(self) -> None: + return None + @contextmanager def dummy_stream(_stream): @@ -350,6 +358,79 @@ class TestDistributedLayerwiseOffloadHook: assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0])) + def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime): + current_block = nn.Linear(2, 2, bias=False) + next_block = nn.Linear(2, 2, bias=False) + next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2)) + source_storage = next_block.weight.untyped_storage().data_ptr() + next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) + next_block.weight.mmap_weight_transform_pending = True + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + + source = hook.cpu_sources[torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage + assert hook.cpu_shards == {} + assert next_block.weight.numel() == 0 + + hook.cpu_staging_buffers = [ + {torch.float32: torch.empty(4)}, + {torch.float32: torch.empty(4)}, + ] + hook.prefetch_layer(slot=0, non_blocking=False) + + assert torch.equal( + next_block.weight, + torch.tensor([[2.0, 3.0], [0.0, 1.0]]), + ) + assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2)) + + def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime): + hooks = [] + for size in (4, 9): + current_block = nn.Linear(1, 1, bias=False) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32)) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + shared_buffers=[None, None], + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + hooks.append(hook) + + resident_block = nn.Module() + resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32)) + resident_group = PinnedResidentLayerGroup( + [resident_block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers( + hooks, + resident_group, + ) + + assert len(staging) == 2 + assert staging[0][torch.float32].numel() == 12 + assert staging[1][torch.float32].numel() == 12 + class TestPinnedResidentLayerGroup: def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): @@ -386,6 +467,42 @@ class TestPinnedResidentLayerGroup: group.offload() + def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime): + blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)] + expected = [] + source_storage = [] + for index, block in enumerate(blocks): + block.weight.data.fill_(index + 1) + expected.append(block.weight.detach().clone()) + source_storage.append(block.weight.untyped_storage().data_ptr()) + + group = PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + + for index, state in enumerate(group._states): + source = state["cpu_sources"][torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage[index] + assert state["cpu_shards"] == {} + assert len(group._cpu_staging_buffers) == 2 + + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) + + blocks[0].weight.data.zero_() + group.offload() + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) + + group.offload() + assert all(block.weight.numel() == 0 for block in blocks) + def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): module = nn.Linear(2, 2, bias=False) @@ -779,6 +896,39 @@ class TestOffloadPlan: with pytest.raises(Exception): plan.block_attrs = {"x": ("y",)} # type: ignore + def test_explicit_component_sources_can_opt_in_to_mmap(self): + class Pipeline(nn.Module): + _supports_mmap_loading = True + weights_sources = [object()] + + assert supports_mmap_loading(Pipeline()) + + def test_component_sources_do_not_implicitly_enable_mmap(self): + class Pipeline(nn.Module): + weights_sources = [object()] + + assert not supports_mmap_loading(Pipeline()) + + @pytest.mark.parametrize( + ("kwargs", "expected"), + [ + ({"dlo_use_allgather": True, "tensor_parallel_size": 2}, True), + ({"dlo_use_allgather": False, "tensor_parallel_size": 1}, True), + ({"dlo_use_allgather": False, "tensor_parallel_size": 2}, False), + ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "use_hsdp": True}, False), + ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "has_online_quant": True}, False), + ], + ) + def test_dlo_mmap_policy(self, kwargs, expected): + defaults = { + "mmap_supported": True, + "dlo_use_allgather": False, + "has_online_quant": False, + "tensor_parallel_size": 1, + "use_hsdp": False, + } + assert should_select_dlo_mmap(**(defaults | kwargs)) is expected + def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime): class TokenRefiner(nn.Module): def __init__(self): @@ -827,6 +977,44 @@ class TestOffloadPlan: class TestMmapValidation: """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" + def test_component_source_prefix_builds_model_namespace(self, tmp_path): + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + checkpoint_file = checkpoint_dir / "model.safetensors" + save_file({"blocks.0.weight": torch.arange(4, dtype=torch.float32)}, str(checkpoint_file)) + + class Pipeline(nn.Module): + _supports_mmap_loading = True + + def __init__(self): + super().__init__() + self.weights_sources = [ + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + ] + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + ), + torch.device("cpu"), + ) + + model_map = backend._build_mmap_model_map(Pipeline()) + + assert model_map == { + "transformer.blocks.0.weight": ( + "blocks.0.weight", + str(checkpoint_file), + ) + } + def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch): """DLO+AllGather should reject when TP world size > 1.""" import json diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d1a1ed15e..bf1f1fb29 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -363,22 +363,39 @@ class DiffusersPipelineLoader: else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - # Skip load_weights only for DLO+AllGather when the model - # supports mmap loading and no online quantization is active. - # This condition MUST match the gate in - # DistributedLayerwiseOffloadBackend.enable() so that the - # loader skips ⟺ enable() uses mmap. - from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading + # Skip load_weights when DLO can keep checkpoint tensors as + # mmap-backed host masters. Storage selection is independent + # of whether DLO reconstructs weights with AllGather. The + # no-AllGather path can use mmap only when the checkpoint is + # already the rank-local runtime layout (currently TP=1 and no + # HSDP); TP/HSDP still require the regular loader callbacks. + # The loader and backend share should_select_dlo_mmap() so + # skipping load_weights always selects mmap in enable(). + from vllm_omni.diffusion.offloader.offload_plan import ( + should_select_dlo_mmap, + supports_mmap_loading, + ) _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) _use_ag = getattr(self.od_config, "dlo_use_allgather", True) _has_online_quant = self._has_online_quant(model) _supports_mmap = supports_mmap_loading(model) + _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) + _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) - _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant + _skip_load = _dist_offload and should_select_dlo_mmap( + mmap_supported=_supports_mmap, + dlo_use_allgather=_use_ag, + has_online_quant=_has_online_quant, + tensor_parallel_size=_tp_size, + use_hsdp=_use_hsdp, + ) if _skip_load: - logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())") + logger.info( + "DLO mmap storage active (%s): skipping load_weights (will map checkpoint weights in enable())", + "AllGather" if _use_ag else "rank-local", + ) else: if _dist_offload and _use_ag and _has_online_quant: raise ValueError( diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 659b64bc4..8283fdbd3 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -11,6 +11,7 @@ from __future__ import annotations import math from collections.abc import Iterable, Mapping from dataclasses import dataclass +from functools import partial from typing import TYPE_CHECKING, Any import torch @@ -350,6 +351,18 @@ class MiniMaxH3Attention(nn.Module): prefix=f"{prefix}.qkv_proj", return_bias=True, ) + if quant_config is None: + # The checkpoint interleaves Q/K/V by query group, while vLLM's + # fused QKV linear consumes contiguous Q, then K, then V rows. The + # regular loader performs this conversion before weight_loader(). + # DLO mmap applies the same adapter while packing one bounded host + # staging block; the raw file-backed master remains node-shared. + self.qkv_proj.weight.mmap_weight_transform = partial( + _reorder_grouped_qkv_to_qkv, + num_query_groups=arch.num_attention_heads, + heads_per_group=1, + head_dim=arch.attention_head_dim, + ) self.num_heads = self.qkv_proj.num_heads self.num_kv_heads = self.qkv_proj.num_kv_heads self.rot_dim = 6 * arch.rope_inv_freq_len diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 3b96b8ab7..fc11632d8 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -532,13 +532,10 @@ class MiniMaxH3Pipeline( encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), ) - # H3's regular loader performs checkpoint-layout conversions (grouped QKV - # and fused MLP). Keep it on that path until mmap can run the same loader - # callbacks for every affected parameter. - _supports_mmap_loading: ClassVar[bool] = False - # TODO(offload): Re-enable after the generic rank-local mmap path can run - # this model's grouped-QKV reorder and fused-MLP packing before TP sharding. - # Do not bypass the regular loader until that equivalence is tested. + # TP1 DLO mmap applies the grouped-QKV adapter while packing each block. + # The fused MLP checkpoint is already gate-then-up, matching its TP1 + # runtime layout. TP>1 and quantized layouts stay on the regular loader. + _supports_mmap_loading: ClassVar[bool] = True _PROFILER_TARGETS: ClassVar[list[str]] = [ "_prepare_reference_videos", "encode_prompt", @@ -648,6 +645,10 @@ class MiniMaxH3Pipeline( od_config.quantization_config, "transformer", ) + if transformer_quant_config is not None: + # Quantized parameters may have packed weights and scale tensors + # whose runtime layout is owned by quant-method loader callbacks. + self._supports_mmap_loading = False self.transformer = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 2db80923d..bed80256a 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -32,6 +32,8 @@ class OffloadConfig: dlo_use_allgather: bool = True dlo_resident_layers: int = 0 # leading DiT layers kept on device model_path: str | None = None # checkpoint path for mmap weight loading + # Appended to preserve the positional order of the pre-existing fields. + tensor_parallel_size: int = 1 @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -58,6 +60,7 @@ class OffloadConfig: parallel_config = getattr(od_config, "parallel_config", None) use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False + tensor_parallel_size = int(getattr(parallel_config, "tensor_parallel_size", 1)) if parallel_config else 1 # Derive dp_size from parallel_config — not user-configurable. # The offload adapts to whatever DP/SP is already configured. @@ -116,7 +119,8 @@ class OffloadConfig: dp_size = 1 logger.info( "Distributed layerwise offload: dlo_use_allgather=False, " - "streaming standard-loader rank-local weights (no DP shard or AllGather)" + "streaming complete rank-local blocks (no DLO shard or AllGather); " + "the backend will select mmap or standard-loader host storage" ) # HSDP already shards parameters into DTensors. Running distributed @@ -134,6 +138,7 @@ class OffloadConfig: strategy=strategy, pin_cpu_memory=pin_cpu_memory, use_hsdp=use_hsdp, + tensor_parallel_size=tensor_parallel_size, dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, dlo_resident_layers=dlo_resident_layers, diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 7da7c8d37..2fd3a8ab2 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1,15 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Distributed Layerwise Offload backend with H2D + AllGather overlap. +"""Distributed Layerwise Offload backend with double-buffered H2D. This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: -* Shards model weights across DP ranks and stores only the local shard - (1/DP_size of full model) in host pinned memory per rank. +* Optionally shards model weights across DP ranks and reconstructs each block + with AllGather, or streams a complete rank-local block without a collective. +* Can retain compatible checkpoint tensors as node-shared mmap sources instead + of creating a persistent private host copy in every rank. * Uses a fixed double-buffer scheme that keeps only two layers' worth of weights on each device at any time. -* Asynchronously pipelines both H2D transfers and AllGather communications - on dedicated streams, fully overlapping them with computation. +* Pipelines H2D transfers and, when enabled, AllGather communications on + dedicated streams, overlapping them with computation. * Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms via vLLM-Omni's platform abstraction layer. """ @@ -31,7 +33,12 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig from .block_discovery import get_blocks_from_dit from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .offload_plan import ( + OffloadPlan, + get_offload_plan, + should_select_dlo_mmap, + supports_mmap_loading, +) from .tensor_utils import ( dtype_size as _dtype_size, ) @@ -74,6 +81,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" @@ -83,6 +91,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.dp_size = dp_size self.rank = rank self.pin_memory = pin_memory + self.rank_local_mmap = rank_local_mmap self.copy_stream = copy_stream or current_omni_platform.Stream() self.comm_stream = comm_stream or current_omni_platform.Stream() @@ -98,8 +107,18 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Sharded host weights for the next block, keyed by dtype self.cpu_shards: dict[torch.dtype, torch.Tensor] = {} + # File-backed source tensors for rank-local mmap. Unlike cpu_shards, + # these remain immutable views of the checkpoint and are never pinned + # or flattened into a model-sized private allocation. + self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + # Rank-local mmap uses two host staging slots shared by every hook in + # this worker. They are assigned by the backend after all block sizes + # are known, mirroring the shared device-buffer allocation. + self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + self.cpu_staging_events: list[Any | None] = [None, None] + # Current slot index (0 or 1). Updated dynamically by the previous # hook's prefetch_layer call via _prefetched_slot. This ensures # correct slot tracking for ALL block counts (including odd N). @@ -154,14 +173,20 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.next_block_parameters = dict(self.next_block.named_parameters()) self.next_block_buffers = dict(self.next_block.named_buffers()) - # Shard next block's weights and store local shard in pinned CPU memory - self.cpu_shards, self.metadata = self._shard_and_pin( - self.next_block_parameters, - self.next_block_buffers, - self.dp_size, - self.rank, - self.pin_memory, - ) + if self.rank_local_mmap: + self.cpu_sources, self.metadata = self._collect_mmap_sources( + self.next_block_parameters, + self.next_block_buffers, + ) + else: + # Shard next block's weights and store local shard in pinned CPU memory. + self.cpu_shards, self.metadata = self._shard_and_pin( + self.next_block_parameters, + self.next_block_buffers, + self.dp_size, + self.rank, + self.pin_memory, + ) # Allocate device buffers only if not using shared buffers from backend if self._owns_buffers: @@ -185,11 +210,65 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._ag_output_sizes: dict[torch.dtype, int] = {} for dtype, metas in self.metadata.items(): total_numel = sum(m["numel"] for m in metas) - shard_numel = self.cpu_shards[dtype].numel() - self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel + if self.rank_local_mmap: + self._ag_output_sizes[dtype] = total_numel + else: + shard_numel = self.cpu_shards[dtype].numel() + self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel return module + @staticmethod + def _collect_mmap_sources( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]: + """Retain file-backed tensors and replace module storage with placeholders. + + The returned sources preserve safetensors mmap storage. Runtime-layout + adapters are applied only while packing a bounded staging slot, so no + full-model anonymous CPU copy is retained by a worker. + """ + cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} + metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + offsets: dict[torch.dtype, int] = {} + + for name, target in chain(params.items(), bufs.items()): + source = target.to_local() if hasattr(target, "to_local") else target + if source.device.type != "cpu": + raise ValueError( + f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}." + ) + + dtype = source.dtype + offset = offsets.get(dtype, 0) + transform = getattr(target, "mmap_weight_transform", None) + if not getattr(target, "mmap_weight_transform_pending", False): + transform = None + + cpu_sources.setdefault(dtype, []).append( + { + "name": name, + "tensor": source.detach(), + "transform": transform, + } + ) + metadata.setdefault(dtype, []).append( + { + "name": name, + "offset": offset, + "numel": source.numel(), + "shape": source.shape, + } + ) + offsets[dtype] = offset + source.numel() + + # The detached source above keeps the mmap storage alive while the + # module parameter/buffer is rebound to the rotating device slot. + set_tensor_storage(target, make_offload_placeholder(target)) + + return cpu_sources, metadata + @staticmethod def _shard_and_pin( params: dict[str, nn.Parameter], @@ -306,6 +385,52 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Prefetch: H2D + AllGather (overlapped on dedicated streams) # # ------------------------------------------------------------------ # + @staticmethod + def _pack_mmap_sources( + cpu_sources: dict[torch.dtype, list[dict[str, Any]]], + metadata: dict[torch.dtype, list[dict[str, Any]]], + slot_buffers: dict[torch.dtype, torch.Tensor], + ) -> dict[torch.dtype, torch.Tensor]: + staged: dict[torch.dtype, torch.Tensor] = {} + for dtype, metas in metadata.items(): + total_numel = sum(meta["numel"] for meta in metas) + destination = slot_buffers[dtype][:total_numel] + sources = cpu_sources[dtype] + for source_info, meta in zip(sources, metas, strict=True): + source = source_info["tensor"] + transform = source_info["transform"] + if callable(transform): + source = transform(source) + if source.dtype != dtype or source.numel() != meta["numel"]: + raise ValueError( + "mmap weight transform changed tensor layout for " + f"{source_info['name']!r}: expected dtype={dtype}, " + f"numel={meta['numel']}, got dtype={source.dtype}, " + f"numel={source.numel()}" + ) + start = meta["offset"] + destination[start : start + meta["numel"]].copy_(source.reshape(-1)) + staged[dtype] = destination + return staged + + def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]: + """Pack this block's mmap views into one bounded host staging slot.""" + previous_copy = self.cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + # Platform events normally expose synchronize(). Retain a + # correctness fallback for test and non-CUDA platform shims. + current_omni_platform.synchronize() + self.cpu_staging_events[slot] = None + + slot_buffers = self.cpu_staging_buffers[slot] + if slot_buffers is None: + raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated") + return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers) + @torch.compiler.disable def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None: """Prepare next block's weights into the shared device buffer for *slot*. @@ -321,11 +446,17 @@ class DistributedLayerwiseOffloadHook(ModelHook): assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated" if self.dp_size <= 1 or self.dp_group is None: + cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards with current_omni_platform.stream(self.copy_stream): - for dtype, cpu_shard in self.cpu_shards.items(): + for dtype, cpu_shard in cpu_weights.items(): gw = gpu_weights[dtype] - gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=non_blocking) + async_copy = non_blocking and cpu_shard.is_pinned() + gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy) evt.record(self.copy_stream) + if self.rank_local_mmap: + # The CPU slot may be overwritten only after this H2D copy has + # finished. The shared event protects reuse by another hook. + self.cpu_staging_events[slot] = evt else: gpu_shards: dict[torch.dtype, torch.Tensor] = {} shard_bufs = self.gpu_shard_buffers[slot] @@ -468,6 +599,7 @@ def apply_distributed_block_hook( comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, ) -> DistributedLayerwiseOffloadHook: """Register a DistributedLayerwiseOffloadHook on *module*.""" registry = HookRegistry.get_or_create(module) @@ -481,6 +613,7 @@ def apply_distributed_block_hook( comm_stream=comm_stream, pin_memory=pin_memory, shared_buffers=shared_buffers, + rank_local_mmap=rank_local_mmap, ) registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -495,15 +628,16 @@ def remove_distributed_block_hook(module: nn.Module) -> None: class PinnedResidentLayerGroup: - """Keep selected layers in pinned host memory between requests. + """Keep selected layers available for stage-scoped device residency. TODO(offload): Extract this alongside PinnedModuleStager after the distributed shard-and-pin operation becomes a shared storage primitive. It currently remains here because it depends on DLO's local-shard layout. Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a - pinned CPU master copy and never copies generated device weights back to - host. Entering the denoise stage performs one asynchronous H2D pass; + pinned CPU master copy (or mmap source plus bounded staging) and never + copies generated device weights back to host. Entering the denoise stage + performs one asynchronous H2D pass; leaving it only restores zero-sized placeholders and releases the device buffers. This lets the following VAE stage reuse the same HBM. @@ -518,10 +652,14 @@ class PinnedResidentLayerGroup: device: torch.device, copy_stream: Any, pin_memory: bool, + rank_local_mmap: bool = False, + defer_staging: bool = False, ) -> None: self.device = device self.copy_stream = copy_stream self.loaded = False + self.rank_local_mmap = rank_local_mmap + self.pin_memory = pin_memory self._states: list[dict[str, Any]] = [] self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] @@ -529,21 +667,47 @@ class PinnedResidentLayerGroup: params = dict(block.named_parameters()) bufs = dict(block.named_buffers()) targets: dict[str, torch.Tensor] = {**params, **bufs} - cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( - params, - bufs, - dp_size=1, - rank=0, - pin_memory=pin_memory, - ) + if rank_local_mmap: + cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources( + params, + bufs, + ) + cpu_shards = {} + else: + cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( + params, + bufs, + dp_size=1, + rank=0, + pin_memory=pin_memory, + ) + cpu_sources = {} self._states.append( { "targets": targets, "cpu_shards": cpu_shards, + "cpu_sources": cpu_sources, "metadata": metadata, } ) + self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + self._cpu_staging_events: list[Any | None] = [None, None] + if rank_local_mmap and not defer_staging: + max_sizes: dict[torch.dtype, int] = {} + for state in self._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + for _ in range(2): + buffers = {} + for dtype, total in max_sizes.items(): + buffer = torch.empty(total, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + self._cpu_staging_buffers.append(buffers) + def load(self) -> None: if self.loaded: return @@ -554,20 +718,41 @@ class PinnedResidentLayerGroup: gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] for state in self._states: block_buffers: dict[torch.dtype, torch.Tensor] = {} - for dtype, cpu_shard in state["cpu_shards"].items(): - block_buffers[dtype] = torch.empty( - cpu_shard.shape, - dtype=dtype, - device=self.device, - ) + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + block_buffers[dtype] = torch.empty(total, dtype=dtype, device=self.device) gpu_buffers.append(block_buffers) self.copy_stream.wait_stream(current_omni_platform.current_stream()) ready = current_omni_platform.Event() with current_omni_platform.stream(self.copy_stream): - for state, block_buffers in zip(self._states, gpu_buffers): - for dtype, cpu_shard in state["cpu_shards"].items(): - block_buffers[dtype].copy_(cpu_shard, non_blocking=True) + for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)): + if self.rank_local_mmap: + slot = index % 2 + previous_copy = self._cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + current_omni_platform.synchronize() + cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources( + state["cpu_sources"], + state["metadata"], + self._cpu_staging_buffers[slot], + ) + else: + cpu_weights = state["cpu_shards"] + + for dtype, cpu_weight in cpu_weights.items(): + block_buffers[dtype].copy_( + cpu_weight, + non_blocking=cpu_weight.is_pinned(), + ) + if self.rank_local_mmap: + slot_ready = current_omni_platform.Event() + slot_ready.record(self.copy_stream) + self._cpu_staging_events[slot] = slot_ready ready.record(self.copy_stream) for state, block_buffers in zip(self._states, gpu_buffers): @@ -630,6 +815,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] self._resident_blocks: list[nn.Module] = [] self._resident_layer_group: PinnedResidentLayerGroup | None = None + self._using_mmap = False + self._using_rank_local_mmap = False def load_resident_layers(self) -> None: """Load the model-declared leading blocks for the denoise stage.""" @@ -664,70 +851,146 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if attr == "mmap_weight_transform": replacement.mmap_weight_transform_pending = True - def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: - """Load DiT weights from safetensors via mmap views (no RSS). - - When the transformer is created on meta device, this method - replaces meta params with mmap views of the checkpoint files. - The views point to OS page cache (shared across ranks), so - no private copy is created. _shard_and_pin then copies only - the 1/dp_size shard from the mmap view to a private buffer. - - Non-DiT modules (VAE, encoders) are NOT affected — they were - created on CPU with real weights via from_pretrained. - """ + @staticmethod + def _add_safetensors_source( + checkpoint_dir: str, + prefix: str, + remap_fn: Any | None, + model_to_ckpt: dict[str, tuple[str, str]], + ) -> int: + """Index one checkpoint directory into the pipeline parameter namespace.""" import glob import json from safetensors import safe_open - model_path = self.config.model_path - if not model_path: - logger.warning("No model_path for mmap weight loading, skipping") - return + indexed_keys = 0 + index_files = glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json")) + if index_files: + for index_file in index_files: + with open(index_file, encoding="utf-8") as handle: + index = json.load(handle) + for checkpoint_key, filename in index.get("weight_map", {}).items(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + model_to_ckpt[model_name] = ( + checkpoint_key, + os.path.join(checkpoint_dir, filename), + ) + indexed_keys += 1 + return indexed_keys - # Resolve HF repo ID to local snapshot path. - # If model_path is a local directory, use it directly; otherwise - # download the safetensors index + weight files via HuggingFace Hub. - if not os.path.isdir(model_path): - from vllm.model_executor.model_loader.weight_utils import ( - download_weights_from_hf, + single_files = glob.glob(os.path.join(checkpoint_dir, "*.safetensors")) + for filename in single_files: + with safe_open(filename, framework="pt", device="cpu") as handle: + for checkpoint_key in handle.keys(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + model_to_ckpt[model_name] = (checkpoint_key, filename) + indexed_keys += 1 + return indexed_keys + + def _build_mmap_model_map(self, pipeline: nn.Module) -> dict[str, tuple[str, str]]: + """Resolve safetensors sources to ``model name -> (key, file)``.""" + remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) + model_to_ckpt: dict[str, tuple[str, str]] = {} + indexed_keys = 0 + + # Prefer the same component sources consumed by DiffusersPipelineLoader. + # Besides avoiding unrelated checkpoint subdirectories, the source + # prefix gives models such as MiniMax-H3 an exact pipeline namespace. + sources = tuple(getattr(pipeline, "weights_sources", ())) + for source in sources: + source_root = source.model_or_path + if not os.path.isdir(source_root): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + source_root = download_weights_from_hf( + model_name_or_path=source.model_or_path, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + revision=source.revision, + subfolder=source.subfolder, + ) + checkpoint_dir = ( + os.path.join(source_root, source.subfolder) if source.subfolder is not None else source_root + ) + indexed_keys += self._add_safetensors_source( + checkpoint_dir, + source.prefix, + remap_fn, + model_to_ckpt, ) - logger.info("model_path %s is not local, downloading from HF", model_path) - model_path = download_weights_from_hf( - model_name_or_path=model_path, - cache_dir=None, - allow_patterns=["*.safetensors", "*.safetensors.index.json"], - ) + if not sources: + model_path = self.config.model_path + if not model_path: + raise RuntimeError("No model_path or weights_sources for mmap weight loading") + if not os.path.isdir(model_path): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) - # Build {checkpoint_key: file_path} from safetensors index - weight_map: dict[str, str] = {} - for idx_file in glob.glob(os.path.join(model_path, "**", "*.safetensors.index.json"), recursive=True): - idx_dir = os.path.dirname(idx_file) - with open(idx_file) as f: - idx = json.load(f) - for key, filename in idx.get("weight_map", {}).items(): - weight_map[key] = os.path.join(idx_dir, filename) + logger.info("model_path %s is not local, downloading from HF", model_path) + model_path = download_weights_from_hf( + model_name_or_path=model_path, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + ) - if not weight_map: - single_files = glob.glob(os.path.join(model_path, "**", "model.safetensors"), recursive=True) - for sf in single_files: - from safetensors import safe_open as _safe_open + # Legacy mmap-enabled pipelines may keep the transformer one level + # below model_path, so preserve recursive discovery in this fallback. + import glob - f = _safe_open(sf, framework="pt", device="cpu") - for key in f.keys(): - weight_map[key] = sf - logger.info("Found single-file safetensors: %s (%d keys)", sf, len(f.keys())) + checkpoint_dirs = { + os.path.dirname(path) + for pattern in ("*.safetensors.index.json", "*.safetensors") + for path in glob.glob(os.path.join(model_path, "**", pattern), recursive=True) + } + for checkpoint_dir in sorted(checkpoint_dirs): + indexed_keys += self._add_safetensors_source( + checkpoint_dir, + "", + remap_fn, + model_to_ckpt, + ) - if not weight_map: + if not model_to_ckpt: raise RuntimeError( - "Distributed layerwise offload could not find safetensors " - f"weights at {model_path}. Expected *.safetensors.index.json " - "or model.safetensors. Ensure the checkpoint path is correct." + "Distributed layerwise offload could not find compatible " + "safetensors weights in the configured model sources." ) + logger.info( + "Built mmap source map: %d model tensors from %d checkpoint keys", + len(model_to_ckpt), + indexed_keys, + ) + return model_to_ckpt - logger.info("Loading DiT weights via mmap (meta → page cache, no RSS): %d keys", len(weight_map)) + def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: + """Load DiT checkpoint tensors as file-backed safetensors views. + + When the transformer is created on meta device, this method + replaces meta params with mmap views of the checkpoint files. + The views point to OS page cache shared across ranks. AllGather mode + copies only the rank's 1/dp_size shard to a persistent private buffer; + rank-local mode retains the views and packs one block at a time into + bounded staging storage. + + Non-DiT modules (VAE, encoders) are NOT affected — they were + created on CPU with real weights via from_pretrained. + """ + from safetensors import safe_open + + model_to_ckpt = self._build_mmap_model_map(pipeline) + logger.info( + "Loading DiT weights via mmap (meta -> shared page cache): %d tensors", + len(model_to_ckpt), + ) # --- Convert DiT modules to meta device --- # The transformer was created normally (with random weights and @@ -812,25 +1075,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): len(saved_buffers.get(id(dit_module), {})), ) - # Build reverse mapping {model_param_name: (ckpt_key, file_path)} using - # the pipeline's _remap_ckpt_key (if available). This handles the - # model-specific checkpoint key remapping (e.g., Cosmos3's - # layers.0.mlp_moe_gen.gate_proj → transformer.gen_layers.0.mlp.gate_proj). - remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) - model_to_ckpt: dict[str, tuple[str, str]] = {} - if remap_fn is not None: - for ckpt_key, file_path in weight_map.items(): - model_name = remap_fn(ckpt_key) - if model_name is not None: - model_to_ckpt[model_name] = (ckpt_key, file_path) - logger.info( - "Built reverse remap: %d model params from %d checkpoint keys", len(model_to_ckpt), len(weight_map) - ) - else: - # No remap function — use checkpoint keys directly - for ckpt_key, file_path in weight_map.items(): - model_to_ckpt[ckpt_key] = (ckpt_key, file_path) - # Collect all meta params in DiT modules dit_ids = set() for dit_module in modules.dits: @@ -850,8 +1094,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): for block in blocks: block_module_ids.update(id(m) for m in block.modules()) - for name, param in pipeline.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): + for name, target in chain(pipeline.named_parameters(), pipeline.named_buffers()): + if not (hasattr(target, "is_meta") and target.is_meta): continue # Check if this param is inside a block (will be handled by @@ -895,13 +1139,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) - self._attach_mmap_param_attrs(name, replacement, param) - parent._parameters[parts[-1]] = replacement + leaf_name = parts[-1] + if leaf_name in parent._parameters: + replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad) + self._attach_mmap_param_attrs(name, replacement, target) + parent._parameters[leaf_name] = replacement + else: + parent._buffers[leaf_name] = tensor loaded += 1 loaded_names.add(name) - logger.info("Mmap weight loading: %d non-block params loaded, %d skipped", loaded, skipped) + logger.info("Mmap weight loading: %d non-block tensors loaded, %d skipped", loaded, skipped) # Now load DiT block params from safetensors using mmap views. # Each block param is assigned an mmap view (no RSS). @@ -927,8 +1175,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # e.g. "transformer.language_model.layers.0" or "transformer.gen_layers.0" block_full_prefix = f"{dit_name}.{blocks_attr}.{block_idx}" - for bname, bparam in block.named_parameters(): - if not (hasattr(bparam, "is_meta") and bparam.is_meta): + for bname, block_target in chain(block.named_parameters(), block.named_buffers()): + if not (hasattr(block_target, "is_meta") and block_target.is_meta): continue # Use reverse remap to find checkpoint key @@ -952,27 +1200,33 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) - self._attach_mmap_param_attrs(full_param_name, replacement, bparam) - parent._parameters[bparts[-1]] = replacement + leaf_name = bparts[-1] + if leaf_name in parent._parameters: + replacement = torch.nn.Parameter(tensor, requires_grad=block_target.requires_grad) + self._attach_mmap_param_attrs(full_param_name, replacement, block_target) + parent._parameters[leaf_name] = replacement + else: + parent._buffers[leaf_name] = tensor block_loaded += 1 loaded_names.add(full_param_name) - logger.info("Mmap block loading: %d params loaded, %d skipped", block_loaded, skipped) + logger.info("Mmap block loading: %d tensors loaded, %d skipped", block_loaded, skipped) # Strict validation: check that all meta params were loaded. # This replaces the validation that AutoWeightsLoader.load_weights() # performs in the regular load path. remaining_meta = [ - name for name, param in pipeline.named_parameters() if hasattr(param, "is_meta") and param.is_meta + name + for name, tensor in chain(pipeline.named_parameters(), pipeline.named_buffers()) + if hasattr(tensor, "is_meta") and tensor.is_meta ] if remaining_meta: # Filter out params inside submodules that will be loaded later # by _load_module_weights_from_mmap (non-block submodules) dit_param_names = set() - for dit_module in modules.dits: - for pname, _ in dit_module.named_parameters(): - dit_param_names.add(pname) + for dit_name, dit_module in zip(modules.dit_names, modules.dits): + for pname, _ in chain(dit_module.named_parameters(), dit_module.named_buffers()): + dit_param_names.add(f"{dit_name}.{pname}") truly_missing = [n for n in remaining_meta if n in dit_param_names] if truly_missing: logger.warning( @@ -1021,10 +1275,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - # Place on the same device as the module's parameters - # Place on the offloader's target device (not the module's - # current device, which may still be meta at this point). - parent._buffers[parts[-1]] = buf.to(self.device) + # Rank-local hooks collect every block buffer as a CPU source. + # AllGather preserves the existing behavior of restoring these + # generated buffers on the target device before shard packing. + restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device + parent._buffers[parts[-1]] = buf.to(restore_device) def _load_module_weights_from_mmap(self, module: nn.Module, dit_name: str, child_name: str) -> None: """Load a non-block DiT submodule's weights from safetensors via mmap.""" @@ -1037,8 +1292,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): model_to_ckpt = getattr(self, "_mmap_model_to_ckpt", {}) loaded = 0 - for pname, param in module.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): + for pname, target in chain(module.named_parameters(), module.named_buffers()): + if not (hasattr(target, "is_meta") and target.is_meta): continue # Build full model param name and look up in reverse remap @@ -1065,9 +1320,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) - self._attach_mmap_param_attrs(full_name, replacement, param) - parent._parameters[parts[-1]] = replacement + leaf_name = parts[-1] + if leaf_name in parent._parameters: + replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad) + self._attach_mmap_param_attrs(full_name, replacement, target) + parent._parameters[leaf_name] = replacement + else: + parent._buffers[leaf_name] = tensor loaded += 1 if loaded > 0: @@ -1265,6 +1524,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, ) sub_hooks = [last_hook] for i, block in enumerate(blocks[:-1]): @@ -1280,6 +1540,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, ) sub_hooks.append(hook) @@ -1402,33 +1663,51 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.config.dlo_resident_layers, ) - # Load weights via mmap for DLO+AllGather. - # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather. - # It must apply declarative checkpoint-to-runtime adapters before the - # standard TP loader shards each block, and must cover staged encoders - # and VAEs as well as DiT blocks. The current AllGather mmap path is - # deliberately not reused because it changes the weight layout. - # Gate condition MUST match diffusers_loader.py: - # supports_mmap_loading(pipeline) and not _has_online_quant - # When the gate is False, the loader has already loaded weights via + # Select checkpoint mmap independently from the DLO transfer protocol. + # AllGather consumes rank shards copied from mmap. Without AllGather, + # TP1/no-HSDP workers retain the same file-backed tensors and pack only + # the next block into bounded, rank-private pinned staging buffers. + # TP/HSDP rank-local layouts still come from the standard loader. + # The loader and backend share should_select_dlo_mmap(), so this + # decision stays aligned with whether load_weights() was skipped. + # When selection is False, the loader has already loaded weights via # regular load_weights() + _process_weights_after_loading(). - if self.config.dlo_use_allgather and self.dp_size > 1: - _has_online_quant = any( - getattr(getattr(module, "quant_method", None), "uses_meta_device", False) - for module in pipeline.modules() + _has_online_quant = any( + getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in pipeline.modules() + ) + if self.config.dlo_use_allgather and _has_online_quant: + raise ValueError( + "Online quantization is incompatible with DLO+AllGather: " + "the sharding + AllGather mechanism flattens weights by " + "dtype, which breaks quantized weight/scale layouts. " + "Please use --dlo-no-use-allgather or disable online " + "quantization." ) - if _has_online_quant: - raise ValueError( - "Online quantization is incompatible with DLO+AllGather: " - "the sharding + AllGather mechanism flattens weights by " - "dtype, which breaks quantized weight/scale layouts. " - "Please use --dlo-no-use-allgather or disable online " - "quantization." + + mmap_supported = supports_mmap_loading(pipeline) + self._using_mmap = should_select_dlo_mmap( + mmap_supported=mmap_supported, + dlo_use_allgather=self.config.dlo_use_allgather, + has_online_quant=_has_online_quant, + tensor_parallel_size=self.config.tensor_parallel_size, + use_hsdp=self.config.use_hsdp, + ) + self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1 + if self._using_mmap: + self._load_weights_via_mmap(pipeline, modules) + if self._using_rank_local_mmap: + logger.info( + "DLO rank-local mmap storage enabled: checkpoint pages are " + "node-shared; each worker owns only two bounded host staging slots" ) - if supports_mmap_loading(pipeline): - self._load_weights_via_mmap(pipeline, modules) - else: - logger.info("Weights loaded via regular loader — skipping mmap (model does not support mmap)") + elif not mmap_supported: + logger.info("Weights loaded via regular loader - model does not support mmap") + elif not self.config.dlo_use_allgather and (self.config.tensor_parallel_size != 1 or self.config.use_hsdp): + logger.info( + "Rank-local mmap disabled for TP=%d/HSDP=%s; using standard-loader tensors", + self.config.tensor_parallel_size, + self.config.use_hsdp, + ) # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) @@ -1528,6 +1807,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, ) block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] @@ -1544,6 +1824,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, ) block_hooks.append(hook) @@ -1572,11 +1853,15 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.device, self.copy_stream, self.config.pin_cpu_memory, + rank_local_mmap=self._using_rank_local_mmap, + defer_staging=bool(self._all_hook_groups), ) pipeline._dlo_residency_controller = self if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) + if self._using_mmap and not self.enabled: + self._release_mmap_handles() return # Unified allocation: 2 shared output buffers + 2 shared shard buffers @@ -1590,6 +1875,21 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): unified_shard_buffers = None if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) + unified_cpu_staging = None + cpu_staging_events = None + if self._using_rank_local_mmap: + unified_cpu_staging = self._allocate_shared_cpu_staging_buffers( + all_hooks, + self._resident_layer_group, + ) + cpu_staging_events = [None, None] + if self._resident_layer_group is not None: + # Resident and streamed layers execute in the same stage and + # reuse the same host slots. Events serialize slot reuse. + self._resident_layer_group._cpu_staging_buffers = [ + buffers for buffers in unified_cpu_staging if buffers is not None + ] + self._resident_layer_group._cpu_staging_events = cpu_staging_events # Shared slot-group tracker: _shared_slot_group[slot] = group_id # that last wrote to that slot. Group-first hooks use this to @@ -1602,6 +1902,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): hook._owns_buffers = False if unified_shard_buffers is not None: hook.gpu_shard_buffers = unified_shard_buffers + if unified_cpu_staging is not None and cpu_staging_events is not None: + hook.cpu_staging_buffers = unified_cpu_staging + hook.cpu_staging_events = cpu_staging_events hook._group_id = group_idx hook._shared_slot_group = shared_slot_group @@ -1628,9 +1931,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.enabled = True - # Release mmap file handles — _shard_and_pin has copied all shards, - # params now point to offload placeholders (not mmap views). - self._release_mmap_handles() + if self._using_mmap and not self._using_rank_local_mmap: + # AllGather mode copied each rank's persistent shard, so the source + # mappings are no longer needed. Rank-local mode retains them as + # the node-shared host master until disable(). + self._release_mmap_handles() # Assign GPU buffers to sharded on-demand modules (VAE/encoders). # Each module gets a dedicated input (shard-sized) and output @@ -1674,9 +1979,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): pass def disable(self) -> None: - if not self.enabled: + if not self.enabled and not hasattr(self, "_mmap_file_cache"): return + if self._using_rank_local_mmap: + current_omni_platform.synchronize() + for blocks in self._blocks: for block in blocks: remove_distributed_block_hook(block) @@ -1691,6 +1999,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.clear() self._resident_blocks.clear() self._resident_layer_group = None + self._release_mmap_handles() + self._using_mmap = False + self._using_rank_local_mmap = False self.enabled = False logger.info("Distributed layer-wise offloading disabled") @@ -1730,6 +2041,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) return shared_buffers + @staticmethod + def _allocate_shared_cpu_staging_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + resident_group: PinnedResidentLayerGroup | None = None, + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate two bounded host slots for rank-local mmap -> device copies.""" + max_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + for dtype, metas in hook.metadata.items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + if resident_group is not None: + for state in resident_group._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + + pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory) + shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + buffers: dict[torch.dtype, torch.Tensor] = {} + for dtype, total_numel in max_sizes.items(): + buffer = torch.empty(total_numel, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + shared_staging[slot] = buffers + + logger.info( + "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()}, + pin_memory, + ) + return shared_staging + @staticmethod def _allocate_shared_shard_buffers( hooks: list[DistributedLayerwiseOffloadHook], diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index f8943757e..904fed650 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -60,15 +60,46 @@ def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: def supports_mmap_loading(pipeline: nn.Module) -> bool: """Whether the pipeline supports mmap weight loading. - A pipeline supports mmap if any module defines ``_remap_ckpt_key``, - which maps checkpoint keys to model parameter names. Models without - this method must use the regular ``load_weights()`` path. + A pipeline supports mmap when either: + + * a module defines ``_remap_ckpt_key``, which maps checkpoint keys to + model parameter names; or + * the pipeline explicitly opts in with ``_supports_mmap_loading=True`` + and declares ``weights_sources``. In that case each source prefix is + used as the model namespace. + + The explicit opt-in is important: many pipelines declare + ``weights_sources`` but still require arbitrary loader callbacks and must + stay on the regular ``load_weights()`` path. This check is shared between ``diffusers_loader.py`` (to decide whether to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()`` (to decide whether to call ``_load_weights_via_mmap``), ensuring both paths use the same gate condition. """ - return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any( - callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules() - ) + declared_support = getattr(pipeline, "_supports_mmap_loading", None) + if declared_support is False: + return False + + has_remap = any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()) + if has_remap: + return True + + return declared_support is True and bool(getattr(pipeline, "weights_sources", ())) + + +def should_select_dlo_mmap( + *, + mmap_supported: bool, + dlo_use_allgather: bool, + has_online_quant: bool, + tensor_parallel_size: int, + use_hsdp: bool, +) -> bool: + """Whether loader and backend should enter DLO's mmap path. + + AllGather preserves its existing fail-closed topology validation inside + the mmap loader. The extra TP/HSDP gate applies to rank-local mmap only. + """ + rank_local_compatible = tensor_parallel_size == 1 and not use_hsdp + return mmap_supported and not has_online_quant and (dlo_use_allgather or rank_local_compatible)