[Core] Support online FP8 with DLO AllGather (#6279)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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;

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
 

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.

docs/user_guide/quantization/fp8.md

@@ -179,6 +179,24 @@ python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers
 vllm serve <your-model> --omni --quantization fp8
 ```
 
+### Distributed layerwise offload
+
+Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
+
+```bash
+vllm serve <your-model> \
+  --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 |

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

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

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

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))

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,