vllm-omni/recipes/cosmos3/Cosmos3-DistOffload.md

Cosmos3 — Distributed Layerwise Offload

Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment

When to use this recipe

Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency.

Key benefits:

• Each rank stores only 1/dp_size of the model weights on host memory
• Only 2 transformer blocks reside on each device at any time (fixed double-buffer)
• DP multi-concurrency: N concurrent requests with near-linear throughput scaling
• mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS

Prerequisites

• 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.)
• vLLM-Omni with distributed layerwise offload support
• Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID)

Serve commands

Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput)

vllm serve /path/to/Cosmos3-Nano \
  --omni \
  --host 0.0.0.0 --port 8000 \
  --no-guardrails --init-timeout 3600 \
  --vae-use-tiling --vae-patch-parallel-size 1 \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 4

Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM)

vllm serve /path/to/Cosmos3-Super \
  --omni \
  --host 0.0.0.0 --port 8000 \
  --no-guardrails --init-timeout 3600 \
  --vae-use-tiling --vae-patch-parallel-size 1 \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 2

Cosmos3-Super with DP=4 + AllGather

vllm serve /path/to/Cosmos3-Super \
  --omni \
  --host 0.0.0.0 --port 8000 \
  --no-guardrails --init-timeout 3600 \
  --vae-use-tiling --vae-patch-parallel-size 1 \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 4

Without DLO AllGather (rank-local weights, no DLO sharding)

Use --dlo-no-use-allgather when you want each rank to stream the standard loader's rank-local weights independently — no DLO AllGather synchronization, no concurrent request requirement, but no DLO weight sharding and higher host memory use. In pure DP, this is equivalent to independent full-weight layerwise-offload replicas; with TP or HSDP, each rank retains the standard loader's local tensor layout.

vllm serve /path/to/Cosmos3-Nano \
  --omni \
  --host 0.0.0.0 --port 8000 \
  --no-guardrails --init-timeout 3600 \
  --vae-use-tiling --vae-patch-parallel-size 1 \
  --enable-distributed-layerwise-offload \
  --data-parallel-size 4 \
  --dlo-no-use-allgather

With SP (sequence parallel) instead of DP

For long-sequence workloads where SP parallelism is preferred:

vllm serve /path/to/Cosmos3-Super \
  --omni \
  --host 0.0.0.0 --port 8000 \
  --no-guardrails --init-timeout 3600 \
  --vae-use-tiling --vae-patch-parallel-size 1 \
  --enable-distributed-layerwise-offload \
  --usp 4

Request examples

Text-to-Image (T2I)

curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \
  -H "Content-Type: multipart/form-data" \
  -F "model=/path/to/Cosmos3-Nano" \
  -F "prompt=A robot arm cleaning a plate, cinematic" \
  -F "size=1024x1024" \
  -F "num_inference_steps=50" \
  -F "guidance_scale=6.0" \
  -F "seed=42"

Text-to-Video (T2V)

curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
  -H "Accept: video/mp4" \
  -F "model=/path/to/Cosmos3-Nano" \
  -F "prompt=A robot arm cleaning a plate, cinematic shot" \
  -F "negative_prompt=blurry" \
  -F "size=832x480" \
  -F "num_frames=29" \
  -F "fps=24" \
  -F "num_inference_steps=35" \
  -F "guidance_scale=6.0" \
  -F "max_sequence_length=4096" \
  -F "flow_shift=10.0" \
  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
  -F "seed=17"

Important: When using DP multi-concurrency (AllGather mode), num_inference_steps must be specified explicitly and must be the same for all concurrent requests. num_inference_steps=None is not allowed because it may resolve differently per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather.

Multiple concurrent requests (DP multi-concurrency)

With --data-parallel-size 4, send 4 concurrent requests to maximize throughput:

for i in 1 2 3 4; do
  curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
    -H "Accept: video/mp4" \
    -F "model=/path/to/Cosmos3-Nano" \
    -F "prompt=A robot arm cleaning a plate, shot ${i}" \
    -F "negative_prompt=blurry" \
    -F "size=832x480" \
    -F "num_frames=29" \
    -F "fps=24" \
    -F "num_inference_steps=35" \
    -F "guidance_scale=6.0" \
    -F "max_sequence_length=4096" \
    -F "flow_shift=10.0" \
    -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
    -F "seed=17" &
done
wait

CLI flags

Flag	Description	Default
--enable-distributed-layerwise-offload	Enable DLO with H2D + AllGather overlap	false
--data-parallel-size N	Number of DP ranks; DLO weight sharding + concurrent requests when AllGather is enabled	1
--dlo-use-allgather	Use shard + AllGather for weight reconstruction (recommended)	true
--dlo-no-use-allgather	Stream standard-loader rank-local weights independently (no DLO sharding, no DLO AllGather)	false
--usp N	Use SP instead of DP for weight sharding (long sequences)	1

Mode comparison

Mode	CLI flags	CPU/rank	HBM/card	Throughput	Use case
DLO + AllGather (DP4)	--enable-distributed-layerwise-offload --data-parallel-size 4	1/4 model	2 blocks	3.3× HSDP	Max throughput, short sequences
DLO + AllGather (DP2)	--enable-distributed-layerwise-offload --data-parallel-size 2	1/2 model	2 blocks	2× HSDP	Balanced throughput + memory
DLO + AllGather (SP4)	--enable-distributed-layerwise-offload --usp 4	1/4 model	2 blocks	1× (single req)	Long sequences (720p+)
DLO no-AllGather	--enable-distributed-layerwise-offload --dlo-no-use-allgather	Standard-loader rank-local weights	2 blocks	0.4× HSDP	No DLO AllGather overhead, high CPU
HSDP	--use-hsdp --hsdp-shard-size 4	0	1/4 model	1× (baseline)	Weights fit in HBM

Memory expectations

Model	Config	Host RAM (cgroup)	HBM/card
Nano (33 GB)	DP2 + AllGather	~38 GB	~10 GB
Nano (33 GB)	DP4 + AllGather	~47 GB	~10 GB
Super (124 GB)	DP2 + AllGather	~157 GB	~15 GB
Super (124 GB)	DP4 + AllGather	~172 GB	~10 GB

Host RAM is dominated by page cache (1× model_size, shared across ranks). HBM only holds 2 transformer blocks + framework + activations, independent of model size.

Limitations

• --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 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
• HSDP + AllGather is rejected (would double-shard weights). HSDP with --dlo-no-use-allgather is accepted at configuration level, but is not yet a broadly validated production combination