[diffusion] Keep the Cosmos3 Super DiT resident on high-memory GPUs (#36375)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Yihao Wang
2026-08-25 21:10:39 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent 3ec22948c1
commit d7baad0116
4 changed files with 25 additions and 10 deletions
+7 -5
View File
@@ -62,11 +62,13 @@ sglang serve \
--num-gpus 1
```
With `--performance-mode auto`, a Cosmos3 Nano checkpoint keeps its DiT and
VAE resident when every selected GPU has at least 120 GiB available at
startup. Below that threshold, auto mode retains the conservative DiT
component-offload policy. This high-memory override is limited to Nano;
Cosmos3 Super checkpoints keep their existing multi-GPU placement defaults.
With `--performance-mode auto`, a Cosmos3 checkpoint keeps its DiT and VAE
resident when every selected GPU has at least 120 GiB available at startup.
Below that threshold, auto mode retains the conservative DiT
component-offload policy. Cosmos3 runs one DiT per pipeline, so component
offload above the threshold only pays to copy the weights out to host memory
and back on every request. Serve `Cosmos3-Super` across multiple GPUs as
shown below so each rank holds a shard of the weights.
For `Cosmos3-Super`, split the model across multiple GPUs: