diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx index 76ff862d9..3239274e2 100644 --- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx +++ b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx @@ -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: diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 838e71d89..81e97360e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -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"), ) diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 097923c7e..288365c12 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -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", diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 10278031e..1f35c3850 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -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",