diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py index 7ca5e4da0..9fd4dcbd4 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/hunyuan.py @@ -164,6 +164,5 @@ class FastHunyuanConfig(HunyuanConfig): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( - auto_disable_component_offload_min_available_memory_gb=150, - auto_disable_component_offload_components=("vae",), + keep_resident_components=("vae",), ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 9c6e9b47e..90380ffbf 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -200,8 +200,8 @@ class LTX2PipelineConfig(PipelineConfig): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( - auto_disable_component_offload_min_available_memory_gb=70, - auto_disable_component_offload_components=("dit",), + keep_resident_min_available_gb=70, + keep_resident_components=("dit",), auto_cfg_parallel_degree_by_num_gpus=((4, 1), (8, 1)), ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index 378bb0529..bbb4d2767 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py @@ -14,13 +14,11 @@ class ModelDeploymentConfig: auto_dit_layerwise_offload: bool = False # if the available memory is bigger than this value, keep dit resident instead of apply layerwise-offload auto_dit_layerwise_offload_high_memory_disable_gb: float | None = None - auto_disable_component_offload_min_available_memory_gb: float | None = None - # keep this explicit because large encoders can OOM even when DiT fits resident - auto_disable_component_offload_components: tuple[OffloadComponentName, ...] = ( - "dit", - "text_encoder", - "image_encoder", - ) + keep_resident_min_available_gb: float | None = None + # only vae -- it is tiny so keeping it resident barely shifts memory; large + # encoders stay offloaded and dit placement stays with the FSDP/dit-layerwise + # policy + keep_resident_components: tuple[OffloadComponentName, ...] = ("vae",) fsdp_auto_min_available_memory_gb: float | None = None fsdp_auto_requires_cfg: bool = True fsdp_auto_requires_default_parallelism: bool = True diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py index 386cc4234..691f13685 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sana_wm.py @@ -329,7 +329,7 @@ class SanaWMRealtimeConfig(SanaWMPipelineConfig): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( auto_dit_layerwise_offload=True, - auto_disable_component_offload_min_available_memory_gb=120, - auto_disable_component_offload_components=("dit",), + keep_resident_min_available_gb=120, + keep_resident_components=("dit",), auto_enable_cfg_parallel=False, ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index d272567e9..bda7ef9d9 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py @@ -115,8 +115,8 @@ class TurboWanT2V1_3B480PConfig(TurboWanT2V480PConfig): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( auto_dit_layerwise_offload=True, - auto_disable_component_offload_min_available_memory_gb=60, - auto_disable_component_offload_components=( + keep_resident_min_available_gb=60, + keep_resident_components=( "text_encoder", "image_encoder", "vae", @@ -202,8 +202,8 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): def get_model_deployment_config(self) -> ModelDeploymentConfig: return ModelDeploymentConfig( auto_dit_layerwise_offload=True, - auto_disable_component_offload_min_available_memory_gb=60, - auto_disable_component_offload_components=( + keep_resident_min_available_gb=60, + keep_resident_components=( "text_encoder", "image_encoder", "vae", diff --git a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py index 7424745df..47477af62 100644 --- a/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args_auto_tune.py @@ -32,6 +32,12 @@ DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES = ( (LAYERWISE_OFFLOAD_VAE_GROUP, "vae_cpu_offload"), ) +# task-type defaults for keep_resident_min_available_gb when a model does not pin +# one: image vae is tiny so any datacenter gpu keeps it resident, video vae is +# larger so it only stays resident on very-high-memory gpus +IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB = 45.0 +DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0 + class ServerArgsAutoTuner: """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures""" @@ -46,6 +52,17 @@ class ServerArgsAutoTuner: def _deployment_config(self) -> ModelDeploymentConfig: return self.server_args.pipeline_config.get_model_deployment_config() + def _resolve_keep_resident_min_available_gb( + self, deployment_config: ModelDeploymentConfig + ) -> float | None: + # explicit per-model > task-type default > global default + explicit = deployment_config.keep_resident_min_available_gb + if explicit is not None: + return explicit + if self.server_args.pipeline_config.task_type.is_image_gen(): + return IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB + return DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB + def adjust_based_on_performance_mode(self) -> None: """Adjust the server args based on the performance mode""" args = self.server_args @@ -96,8 +113,8 @@ class ServerArgsAutoTuner: min_available_gb = self._get_min_available_device_memory_gb() deployment_config = self._deployment_config() - disable_threshold_gb = ( - deployment_config.auto_disable_component_offload_min_available_memory_gb + disable_threshold_gb = self._resolve_keep_resident_min_available_gb( + deployment_config ) if ( min_available_gb is not None @@ -105,7 +122,7 @@ class ServerArgsAutoTuner: and min_available_gb >= disable_threshold_gb ): changed = [] - components = deployment_config.auto_disable_component_offload_components + components = deployment_config.keep_resident_components if ( args.layerwise_offload_components is not None and not args.is_arg_explicitly_set("layerwise_offload_components") @@ -400,9 +417,7 @@ class ServerArgsAutoTuner: return components deployment_config = self._deployment_config() - threshold_gb = ( - deployment_config.auto_disable_component_offload_min_available_memory_gb - ) + threshold_gb = self._resolve_keep_resident_min_available_gb(deployment_config) if threshold_gb is None: return components @@ -410,9 +425,7 @@ class ServerArgsAutoTuner: if min_available_gb is None or min_available_gb < threshold_gb: return components - resident_components = set( - deployment_config.auto_disable_component_offload_components - ) + resident_components = set(deployment_config.keep_resident_components) filtered_components = [ component for component in components 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 c7c72337f..e70b8604a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import FastHunyuanConfig from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, LTX23PipelineConfig, @@ -837,12 +838,8 @@ class TestOffloadDefaults(unittest.TestCase): self.assertTrue(zimage_deployment.fsdp_auto_requires_cfg) self.assertFalse(zimage_deployment.auto_dit_layerwise_offload) - self.assertEqual( - ltx_deployment.auto_disable_component_offload_min_available_memory_gb, 70 - ) - self.assertEqual( - ltx_deployment.auto_disable_component_offload_components, ("dit",) - ) + self.assertEqual(ltx_deployment.keep_resident_min_available_gb, 70) + self.assertEqual(ltx_deployment.keep_resident_components, ("dit",)) self.assertEqual( ltx_deployment.auto_cfg_parallel_degree_by_num_gpus, ((4, 1), (8, 1)) ) @@ -859,6 +856,15 @@ class TestOffloadDefaults(unittest.TestCase): self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60) self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload) + # fasthunyuan no longer pins 150gb -- falls back to the global video default + fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config() + self.assertIsNone(fast_hunyuan_deployment.keep_resident_min_available_gb) + self.assertEqual(fast_hunyuan_deployment.keep_resident_components, ("vae",)) + + # default keeps only vae resident (encoders are large, dit owned by FSDP) + self.assertEqual(qwen_deployment.keep_resident_components, ("vae",)) + self.assertIsNone(qwen_deployment.keep_resident_min_available_gb) + def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self): args = self._from_dict_with_pipeline_config( SanaWMPipelineConfig(), @@ -948,7 +954,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertIsNone(args.image_encoder_cpu_offload) self.assertFalse(args.enable_cfg_parallel) - def test_default_auto_replaces_text_encoder_cpu_offload_with_layerwise(self): + def test_default_auto_keeps_image_vae_resident_when_memory_allows(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), kwargs={"model_path": "Qwen/Qwen-Image"}, @@ -956,10 +962,25 @@ class TestOffloadDefaults(unittest.TestCase): self.assertEqual(args.performance_mode, "auto") self.assertFalse(args.use_fsdp_inference) + # 80gb > image threshold (45gb): only vae kept resident, encoders stay + # offloaded layerwise, dit unchanged + self.assertTrue(args.dit_cpu_offload) + self.assertEqual( + args.layerwise_offload_components, + ["text_encoder", "image_encoder"], + ) + self.assertFalse(args.vae_cpu_offload) + + def test_auto_image_offloads_aux_below_resident_threshold(self): + # 40gb < image threshold (45gb): aux incl. vae still offloaded to save vram + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + memory_gb=40, + kwargs={"model_path": "Qwen/Qwen-Image"}, + ) + + self.assertEqual(args.performance_mode, "auto") self.assertTrue(args.dit_cpu_offload) - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( args.layerwise_offload_components, ["text_encoder", "image_encoder", "vae"], @@ -1296,7 +1317,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertEqual(args.ltx2_two_stage_device_mode, "resident") self.assertEqual(args.layerwise_offload_components, ["text_encoder"]) - def test_auto_multi_gpu_qwen_replaces_text_encoder_offload_with_cfg(self): + def test_auto_multi_gpu_qwen_keeps_vae_resident_with_cfg(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), kwargs={ @@ -1308,14 +1329,14 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) + # 80gb > image threshold (45gb): only vae resident, encoders offloaded; + # cfg/dit unchanged self.assertTrue(args.dit_cpu_offload) - self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) + self.assertFalse(args.vae_cpu_offload) def test_auto_multi_gpu_zimage_base_prefers_fsdp(self): args = self._from_dict_with_pipeline_config( @@ -1357,11 +1378,12 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) + self.assertFalse(args.vae_cpu_offload) + # explicit use_fsdp_inference skips the residency pass, but the layerwise + # filter still drops vae (kept resident); encoders stay offloaded self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self): @@ -1377,13 +1399,14 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) + # 50gb still > image threshold (45gb): vae resident, encoders offloaded; + # fsdp skipped (qwen does not opt into auto fsdp) self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) + self.assertFalse(args.vae_cpu_offload) def test_auto_multi_gpu_qwen_uses_selected_gpu_min_available_memory(self): args = self._from_dict_with_pipeline_config( @@ -1400,7 +1423,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) - def test_auto_multi_gpu_qwen_replaces_text_encoder_offload_with_headroom(self): + def test_auto_multi_gpu_qwen_keeps_vae_resident_with_headroom(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), available_memory_gb={1: 72, 2: 80}, @@ -1414,13 +1437,14 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.enable_cfg_parallel) + # min available across selected gpus is 72gb > image threshold (45gb): + # vae resident, encoders offloaded self.assertTrue(args.dit_cpu_offload) - self.assertFalse(args.text_encoder_cpu_offload) - self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( args.layerwise_offload_components, - ["text_encoder", "image_encoder", "vae"], + ["text_encoder", "image_encoder"], ) + self.assertFalse(args.vae_cpu_offload) def test_speed_mode_single_gpu_disables_offload(self): args = self._from_dict_with_pipeline_config(