vllm-omni/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

Distributed Layerwise Offloading

Distributed layerwise offloading (DLO) extends block streaming to multi-device deployments. With AllGather enabled, each rank stores roughly 1 / dp_size of the host weights and reconstructs each layer at runtime. Without AllGather, each rank streams its standard-loader rank-local weights independently.

See the DLO feature design for the implementation contract and compatibility matrix.

Execution model

DLO overlaps three operations with a fixed two-block device buffer:

Compute stream:  [Layer N]          [Layer N+1]        [Layer N+2]
H2D stream:      [H2D shard N+1]    [H2D shard N+2]
AllGather:       [AG N+1]           [AG N+2]
Slots:           slot 0: Layer N    slot 1: Layer N+1

AllGather communicates only request-independent weight shards, so data- parallel ranks may process different requests concurrently.

Usage

# Four ranks with sharded host weights and AllGather
vllm serve /path/to/model --omni \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 4

# Standard-loader rank-local weights, without DLO AllGather
vllm serve /path/to/model --omni \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 4 \
  --dlo-no-use-allgather

# Sequence parallel deployment
vllm serve /path/to/model --omni \
  --enable-distributed-layerwise-offload \
  --usp 4

from vllm_omni import Omni

omni = Omni(
    model="/path/to/model",
    enable_distributed_layerwise_offload=True,
    dlo_use_allgather=True,
)

Flags

Flag	Meaning	Default
--enable-distributed-layerwise-offload	Enable DLO	false
--data-parallel-size N	DP ranks and AllGather weight-sharding group	1
--dlo-use-allgather	Shard host weights and reconstruct with AllGather	true
--dlo-no-use-allgather	Keep standard-loader rank-local weights	false
--dlo-resident-layers N	Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths	0

mmap weight loading

The DLO plus AllGather path:

1. saves non-persistent buffers such as RoPE frequencies;
2. moves the normally created transformer to the meta device;
3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
4. calls model-specific post_load_weights() conversions; and
5. restores the saved non-persistent buffers.

This avoids one full checkpoint RSS copy per rank and does not require model-specific loading code.

Declarative topology

Models may declare an OffloadPlan instead of embedding offload logic:

from vllm_omni.diffusion.offloader import OffloadPlan


class MyPipeline(nn.Module):
    _dit_modules = ["transformer"]
    _offload_plan = OffloadPlan(
        block_attrs={"transformer": ("blocks",)},
        offload_submodules={"context_encoder": "layers"},
    )

When no plan exists, discovery falls back to _layerwise_offload_blocks_attrs and then heuristic attribute lookup.

Data-parallel concurrency

With data_parallel_size > 1 and AllGather enabled, the scheduler can process up to dp_size requests per denoising step. Every concurrent request must set the same explicit num_inference_steps; None is rejected because every rank must enter each collective.

Limitations

• Online FP8 quantization is rejected with the DLO plus AllGather mmap path. Use --dlo-no-use-allgather or disable online quantization.
• Tensor parallel size greater than one is rejected in the mmap path because it bypasses TP-aware loader callbacks. The no-AllGather path retains the standard loader and remains experimental with TP.
• HSDP plus AllGather is rejected to avoid double sharding. HSDP without AllGather has limited end-to-end validation.
• 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.

See the Cosmos3 DistOffload recipe for an end-to-end example.