minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] Keep Cosmos3 Nano resident on 96 GB GPUs (#36641)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-08-29 07:40:46 +08:00
committed by GitHub
parent e1b3bba3cc
commit 50bc1a3767
4 changed files with 103 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cookbook/diffusion/Cosmos/Cosmos3.mdx
+8 -7
View File
@@ -62,13 +62,14 @@ sglang serve \
--num-gpus 1
```
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.
With `--performance-mode auto`, Cosmos3 Nano keeps its DiT and VAE resident
when every selected GPU has at least 90 GiB available at startup. Other
Cosmos3 checkpoints use a 120 GiB threshold. Below the applicable 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:
python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py
+40 -8
View File
@@ -24,6 +24,20 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config impo
)
COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
COSMOS3_NANO_ARCH_SIGNATURE = (4096, 36, 32)
COSMOS3_NANO_KEEP_RESIDENT_MIN_AVAILABLE_GB = 90
COSMOS3_DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120
@functools.lru_cache(maxsize=None)
def _transformer_config(model_path: str) -> dict:
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
get_diffusers_component_config,
)
return get_diffusers_component_config(
component_path=os.path.join(model_path, "transformer")
)
@functools.lru_cache(maxsize=None)
@@ -34,19 +48,25 @@ def is_edge_checkpoint(model_path: str) -> bool:
is available before the weights are on device (e.g. when resolving sampling
defaults in the client process).
"""
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
get_diffusers_component_config,
)
config = get_diffusers_component_config(
component_path=os.path.join(model_path, "transformer")
)
config = _transformer_config(model_path)
return (
config.get("backbone_type") == COSMOS3_EDGE_BACKBONE_TYPE
or config.get("hidden_act") == "relu2"
)
@functools.lru_cache(maxsize=None)
def is_nano_checkpoint(model_path: str) -> bool:
"""Whether the checkpoint uses the Nano transformer architecture."""
config = _transformer_config(model_path)
signature = (
config.get("hidden_size"),
config.get("num_hidden_layers"),
config.get("num_attention_heads"),
)
return signature == COSMOS3_NANO_ARCH_SIGNATURE
@functools.lru_cache(maxsize=None)
def _distilled_sampler_config(model_path: str) -> dict | None:
"""The fixed-step sampler config for a distilled checkpoint, else ``None``.
@@ -125,6 +145,7 @@ class Cosmos3Config(PipelineConfig):
# Pre-computed once in update_config_from_dict from the resolved model_path.
# None until that point (e.g. in unit-test mocks that never call update_config_from_dict).
is_edge: bool | None = None
is_nano: bool | None = None
distilled_sigmas: list[float] | None = None
def __post_init__(self):
@@ -145,6 +166,7 @@ class Cosmos3Config(PipelineConfig):
if self.model_path:
self.distilled_sigmas = get_distilled_sigmas(self.model_path)
self.is_edge = is_edge_checkpoint(self.model_path)
self.is_nano = is_nano_checkpoint(self.model_path)
if self.distilled_sigmas is not None:
self.scheduler_class_override = None
@@ -171,7 +193,17 @@ class Cosmos3Config(PipelineConfig):
def get_model_deployment_config(self) -> ModelDeploymentConfig:
# Keep the DiT and VAE resident when the GPUs have the headroom.
is_nano = self.is_nano
if is_nano is None:
# Directly constructed configs in callers/tests have not resolved
# checkpoint metadata yet; registered model IDs remain unambiguous.
is_nano = "cosmos3-nano" in (self.model_path or "").lower()
threshold_gb = (
COSMOS3_NANO_KEEP_RESIDENT_MIN_AVAILABLE_GB
if is_nano
else COSMOS3_DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB
)
return ModelDeploymentConfig(
keep_resident_min_available_gb=120,
keep_resident_min_available_gb=threshold_gb,
keep_resident_components=("dit", "vae"),
)
python/sglang/multimodal_gen/test/unit/test_cosmos3.py
+42 -1
View File
@@ -14,7 +14,10 @@ from PIL import Image
from sglang.multimodal_gen.configs.models.dits.cosmos3video import (
_build_cosmos3_param_names_mapping,
)
from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config
from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import (
Cosmos3Config,
is_nano_checkpoint,
)
from sglang.multimodal_gen.configs.sample.cosmos3 import (
COSMOS3_EDGE_SUPPORTED_RESOLUTIONS,
Cosmos3SamplingParams,
@@ -841,6 +844,44 @@ class TestCosmos3ActionEndpoint(unittest.TestCase):
class TestCosmos3ModelResolution(unittest.TestCase):
"""Verify Cosmos3 checkpoints resolve to the native SGLang pipeline."""
def test_nano_architecture_detection_does_not_depend_on_model_path(self):
cases = (
(
{
"hidden_size": 4096,
"num_hidden_layers": 36,
"num_attention_heads": 32,
},
True,
),
(
{
"hidden_size": 5120,
"num_hidden_layers": 64,
"num_attention_heads": 64,
},
False,
),
(
{
"hidden_size": 2048,
"num_hidden_layers": 28,
"num_attention_heads": 16,
},
False,
),
)
for index, (transformer_config, expected) in enumerate(cases):
with self.subTest(transformer_config=transformer_config):
is_nano_checkpoint.cache_clear()
with mock.patch(
"sglang.multimodal_gen.configs.pipeline_configs.cosmos3._transformer_config",
return_value=transformer_config,
):
self.assertEqual(
is_nano_checkpoint(f"/models/checkpoint-{index}"), expected
)
def test_hf_checkpoint_uses_registered_native_pipeline_config(self):
for model_path in (
"nvidia/Cosmos3-Nano",
python/sglang/multimodal_gen/test/unit/test_server_args.py
+13 -3
View File
@@ -1440,6 +1440,12 @@ class TestOffloadDefaults(unittest.TestCase):
cosmos3_deployment = Cosmos3Config(
model_path="nvidia/Cosmos3-Nano"
).get_model_deployment_config()
cosmos3_super_deployment = Cosmos3Config(
model_path="nvidia/Cosmos3-Super"
).get_model_deployment_config()
local_cosmos3_nano_deployment = Cosmos3Config(
model_path="/models/custom-checkpoint", is_nano=True
).get_model_deployment_config()
wan_deployment = WanT2V480PConfig().get_model_deployment_config()
mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
@@ -1452,8 +1458,12 @@ class TestOffloadDefaults(unittest.TestCase):
self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
self.assertEqual(qwen_deployment.dit_layerwise_offload_modes, ())
self.assertEqual(cosmos3_deployment.keep_resident_min_available_gb, 120)
self.assertEqual(cosmos3_deployment.keep_resident_min_available_gb, 90)
self.assertEqual(cosmos3_deployment.keep_resident_components, ("dit", "vae"))
self.assertEqual(cosmos3_super_deployment.keep_resident_min_available_gb, 120)
self.assertEqual(
local_cosmos3_nano_deployment.keep_resident_min_available_gb, 90
)
self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
self.assertEqual(wan_deployment.dit_layerwise_offload_modes, ("memory",))
@@ -1930,7 +1940,7 @@ class TestOffloadDefaults(unittest.TestCase):
def test_auto_cosmos3_keeps_dit_resident_on_high_memory_gpu(self):
args = self._from_dict_with_pipeline_config(
Cosmos3Config(model_path="nvidia/Cosmos3-Nano"),
available_memory_gb=139,
available_memory_gb=95,
kwargs={
"model_path": "nvidia/Cosmos3-Nano",
"performance_mode": "auto",
@@ -1947,7 +1957,7 @@ class TestOffloadDefaults(unittest.TestCase):
def test_auto_cosmos3_offloads_dit_below_resident_threshold(self):
args = self._from_dict_with_pipeline_config(
Cosmos3Config(model_path="nvidia/Cosmos3-Nano"),
available_memory_gb=100,
available_memory_gb=85,
kwargs={
"model_path": "nvidia/Cosmos3-Nano",
"performance_mode": "auto",