diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index b38f0443d..e254e3a43 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -10,6 +10,30 @@ For user-facing commands, see the [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md). +## Feature compatibility + +Host-storage optimization and runtime compatibility are separate decisions. +When direct checkpoint mmap is unavailable, DLO can still use tensors produced +by the ordinary loader. "Compatibility path" below means that fallback is +implemented but has less end-to-end coverage than the primary path. + +Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported. + +| Feature | DLO + AllGather | DLO without AllGather | +|---|---|---| +| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. | +| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. | +| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | +| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. | +| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. | +| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. | +| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. | +| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. | + +See [Parallelism compatibility](#parallelism-compatibility) and +[Request and loading constraints](#request-and-loading-constraints) for the +detailed contracts and validation boundaries. + ## Status DLO is implemented for multi-device diffusion execution. The default @@ -181,8 +205,11 @@ requirements. Direct checkpoint mmap can back either transfer path. It is currently limited to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the -ordinary loader. Online quantization remains incompatible with DLO AllGather; -use `--dlo-no-use-allgather` or disable online quantization. +ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the +ordinary loader finalizes their runtime weights and scales; DLO then shards and +reconstructs those tensors with their recorded layouts. Other online methods +must use `--dlo-no-use-allgather` or disable online quantization until their +runtime layouts are validated. A normalized runtime mmap cache, built through the ordinary loader, is the proposed general mechanism for sharing transformed TP or quantized layouts. @@ -197,6 +224,8 @@ Current source-level validation includes: - HSDP + DLO without AllGather acceptance at configuration level; - loader preflight fallback for TP, HSDP, online quantization, unknown custom loaders, missing keys, and shape/dtype mismatches; +- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO + sharding of finalized weights and scales; - exact loader-to-backend plan transfer and ordinary-loader fallback; - rank-local mmap source retention, bounded two-slot staging, and adapter transforms without parameter-side flags; diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md index 8d6cd09c3..84c9641b3 100644 --- a/docs/design/feature/quantization.md +++ b/docs/design/feature/quantization.md @@ -190,12 +190,15 @@ configured parallelism strategy: - `WeightsMapper` transformations must be applied before matching ignored-layer lists or component prefixes. -Distributed layerwise offload has an additional boundary: its sharded mmap -AllGather path is not compatible with online quantization because the path -expects already materialized checkpoint weights. Use the rank-local -`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when -combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md) -for the full compatibility matrix. +Distributed layerwise offload has an additional boundary: direct checkpoint +mmap cannot prepare online-quantized tensors because it bypasses the ordinary +loader that creates their runtime weights and scales. Per-tensor online FP8 +linears can still use DLO AllGather through the ordinary loader: each rank first +finalizes the FP8 tensors, then DLO shards those runtime tensors for +reconstruction during prefetch. Other online methods remain limited to +`--dlo-no-use-allgather` until their packing and scale layouts are validated. +See the [DLO design document](offloader/distributed_layerwise_offload.md) for +the full compatibility matrix. ## Adding a quantization backend diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 6157c866b..49139e695 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -138,8 +138,11 @@ must enter each collective. but it provides no shared-mmap host-memory benefit or guarantee. - HSDP plus AllGather is rejected to avoid double sharding. HSDP without AllGather has limited end-to-end validation. -- Online quantization uses the ordinary loader with no-AllGather. It remains - incompatible with DLO AllGather. +- Per-tensor online FP8 linears use the ordinary loader and can run with either + DLO transfer path. With AllGather, every rank temporarily materializes the + complete FP8 model in host memory before DLO retains only its shard. Other + online quantization methods require no-AllGather until their runtime layouts + are validated. - Resident leading layers require `--dlo-no-use-allgather` and a model `OffloadPlan` that declares eligible `resident_dit_paths`. - DP concurrency requires an explicit, identical inference-step count. diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 397a412e8..c57642587 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -179,6 +179,24 @@ python text_to_image.py --model --quantization fp8 --ignored-layers vllm serve --omni --quantization fp8 ``` +### Distributed layerwise offload + +Per-tensor online FP8 linears can be combined with DLO's default AllGather path: + +```bash +vllm serve \ + --omni \ + --quantization fp8 \ + --enable-distributed-layerwise-offload +``` + +Direct checkpoint mmap is not used for this combination. Each rank runs the +ordinary loader to create the final FP8 weights and scales before DLO shards +them, so startup temporarily materializes the complete FP8 model in host memory +on every rank. Runtime DLO host residency is sharded normally. Use +`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract +does not fit the workload. + ## Parameters | Parameter | Type | Default | Description | diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index f6e5dd43b..e8c66ca37 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -416,9 +416,10 @@ For example, keep the first main block's attention projections in BF16 with: ``` The structured option replaces `--quantization fp8`. Online FP8 can be used -with H3 layerwise offload and with distributed layerwise offload's full-weight -per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not -supported for runtime-created FP8 weights. +with H3 layerwise offload and with both DLO transfer paths. The default +AllGather path uses the ordinary loader to finalize FP8 weights and scales +before sharding them across ranks. `--dlo-no-use-allgather` instead retains +complete rank-local tensors and avoids the synchronized request-wave contract. ## AMD ROCm (gfx942 / gfx950) @@ -870,8 +871,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ H3 native `tile` mode only. - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use pure Ulysses. -- Runtime-created FP8 weights do not support the sharded DLO AllGather path; - use `--dlo-no-use-allgather` when combining online FP8 with DLO. +- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model + in host memory on every rank during startup before retaining only each rank's + shard. Size startup host memory for that transient peak. - TeaCache and Cache-DiT cannot be enabled on the same server. - Image+audio Ref2VA accepts exactly one image and one audio reference. - Video Ref2VA accepts one or more video files, but not an additional standalone diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md index 8eeea5cc2..2866e7a0c 100644 --- a/recipes/cosmos3/Cosmos3-DistOffload.md +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -189,8 +189,10 @@ wait - `--dlo-use-allgather` (default) requires all concurrent requests to have the same `num_inference_steps` (AllGather is a collective that needs all ranks synchronized) -- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path; - use `--dlo-no-use-allgather` or disable online quantization +- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader + instead of direct checkpoint mmap, then shard the finalized FP8 weights and + scales. Size host memory for a transient complete FP8 model on every rank at + startup - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with `--dlo-no-use-allgather` retains the standard TP loader and should be treated as an experimental compatibility path diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index aa7651bb6..527942961 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -14,7 +14,7 @@ from huggingface_hub import snapshot_download from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, @@ -91,6 +91,17 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade return loader +def _make_dlo_online_quant_config() -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model="", + dtype=torch.float32, + quantization_config="fp8", + parallel_config=DiffusionParallelConfig(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + + def test_strict_check_only_validates_source_prefix_parameters(): model = _DummyPipelineModel(source_prefix="transformer.") loader = _make_loader_with_weights(["transformer.weight"]) @@ -394,6 +405,78 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): assert loader.take_host_weight_plan() is None +def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod) + model.transformer.quant_method.uses_meta_device = True + calls: list[object] = [] + allowlist_models: list[nn.Module] = [] + + original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods + + def check_allowlist(candidate: nn.Module) -> tuple[str, ...]: + allowlist_models.append(candidate) + return original_allowlist_check(candidate) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist) + loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign] + loader.load_weights = ( # type: ignore[method-assign] + lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu)) + ) + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model + assert allowlist_models == [model] + assert calls == [("load", True), "process"] + assert loader.take_host_weight_plan() is None + + +def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + class UnsupportedOnlineMethod: + uses_meta_device = True + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = UnsupportedOnlineMethod() + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + with pytest.raises(ValueError, match="per-tensor FP8 linears"): + loader.load_model(load_device="cpu") + + def test_hsdp_processes_quantized_weights_before_sharding(mocker): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod from vllm_omni.diffusion.offloader.module_collector import PipelineModules diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 6c8248391..c67a2d45b 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -185,6 +185,79 @@ class TestDistributedLayerwiseOffloadHook: assert next_block.weight.stride() == expected_stride assert torch.equal(next_block.weight, expected) + def test_allgather_reconstructs_online_fp8_weight_and_scale( + self, + patched_offload_runtime, + monkeypatch, + ): + """DLO must reconstruct finalized online-FP8 tensors byte-for-byte.""" + + class OnlineFp8Block(nn.Module): + def __init__(self): + super().__init__() + logical_weight = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + self.weight = nn.Parameter(logical_weight.t(), requires_grad=False) + self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False) + + current_block = OnlineFp8Block() + rank0_block = OnlineFp8Block() + rank1_block = OnlineFp8Block() + expected_weight = rank0_block.weight.detach().clone() + expected_scale = rank0_block.weight_scale.detach().clone() + expected_stride = rank0_block.weight.stride() + + rank0_hook = DistributedLayerwiseOffloadHook( + next_block=rank0_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + rank1_hook = DistributedLayerwiseOffloadHook( + next_block=rank1_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=1, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + rank0_hook.initialize_hook(current_block) + rank1_hook.initialize_hook(OnlineFp8Block()) + rank0_hook.gpu_shard_buffers = [ + {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2) + ] + + def fake_allgather(output, local_shard, *, group): + del group + remote_shard = rank1_hook.cpu_shards[local_shard.dtype] + shard_size = local_shard.numel() + output[:shard_size].copy_(local_shard) + output[shard_size : 2 * shard_size].copy_(remote_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + rank0_hook.prefetch_layer(slot=0, non_blocking=False) + + assert rank0_block.weight.stride() == expected_stride + torch.testing.assert_close( + rank0_block.weight.float(), + expected_weight.float(), + rtol=0, + atol=0, + ) + torch.testing.assert_close(rank0_block.weight_scale, expected_scale) + def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime): """DTensor wrapper should be preserved through prefetch/offload cycle.""" current_block = TinyBlock(_make_values(1.0)) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 5e05bd547..62bc0525f 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -482,12 +482,17 @@ class DiffusersPipelineLoader: ) else: if _dist_offload and _use_ag 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." + unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model) + if unsupported_methods: + raise ValueError( + "DLO+AllGather supports online quantization only for " + "per-tensor FP8 linears; unsupported online methods: " + f"{', '.join(unsupported_methods)}. Please use " + "--dlo-no-use-allgather or disable online quantization." + ) + logger.info( + "Online per-tensor FP8 with DLO+AllGather: using the " + "ordinary loader before sharding finalized weights and scales" ) if _dist_offload and plan_result is not None: logger.info( @@ -548,6 +553,29 @@ class DiffusersPipelineLoader: ``uses_meta_device=True``, e.g. online FP8).""" return has_online_quantization(model) + @staticmethod + def _unsupported_dlo_allgather_online_quant_methods(model: nn.Module) -> tuple[str, ...]: + """Return unsupported online-quant methods for DLO AllGather. + + Per-tensor online FP8 is safe after the ordinary loader has finalized + its weight and scale parameters. DLO shards those runtime tensors by + dtype and reconstructs their recorded shapes and strides before the + kernel consumes them. Other online methods may create different scale, + packing, or aliasing layouts and remain fail-closed until validated. + """ + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + unsupported: set[str] = set() + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if not getattr(quant_method, "uses_meta_device", False): + continue + if not isinstance(quant_method, Fp8PerTensorOnlineLinearMethod): + unsupported.add(type(quant_method).__name__) + return tuple(sorted(unsupported)) + def _apply_skip_softmax_calibration(self, model: nn.Module) -> None: from vllm_omni.diffusion.attention.backends.trtllm_calibration import ( apply_skip_softmax_calibration,