diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 22c152794..331c9e9d8 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -893,21 +893,18 @@ class ServerArgs(DisaggArgsMixin): if explicitly_set_component_names is not None: self.layerwise_offload_components = explicitly_set_component_names - self._disable_cpu_offload_for_layerwise_components( + self._disable_non_dit_cpu_offload_for_layerwise_components( explicitly_set_component_names ) return - def _disable_cpu_offload_for_layerwise_components( + def _disable_non_dit_cpu_offload_for_layerwise_components( self, component_names: list[str] ) -> None: - # Layerwise offload owns H2D/D2H for selected component weights. + # non-DiT layerwise offload replaces the corresponding component-level CPU offload flag_names = cpu_offload_flags_for_layerwise_components(component_names) disabled_flag_names: list[str] = [] - if "dit_cpu_offload" in flag_names and self.dit_cpu_offload is not False: - self.dit_cpu_offload = False - disabled_flag_names.append("dit_cpu_offload") if ( "text_encoder_cpu_offload" in flag_names and self.text_encoder_cpu_offload is not False @@ -1254,7 +1251,9 @@ class ServerArgs(DisaggArgsMixin): default=ServerArgs.dit_layerwise_offload, help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. " "It selects only the DiT layerwise group. Cannot be used together with cache-dit " - "(SGLANG_CACHE_DIT_ENABLED), dit_cpu_offload, or use_fsdp_inference.", + "(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. May be combined with " + "--dit-cpu-offload, in which case DiT weights stay on host memory and only the " + "layers needed for the current step are brought on-device (lowest peak GPU memory).", ) parser.add_argument( "--layerwise-offload-components", @@ -1836,20 +1835,14 @@ class ServerArgs(DisaggArgsMixin): if self.dit_offload_prefetch_size < 0.0: raise ValueError("dit_offload_prefetch_size must be non-negative") - should_disable_dit_cpu_offload = self.is_dit_layerwise_offload_selected - if self.use_fsdp_inference and should_disable_dit_cpu_offload: + is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected + if self.use_fsdp_inference and is_dit_layerwise_offload_selected: logger.warning( "layerwise offload is selected for DiT components, automatically disabling use_fsdp_inference." ) self.use_fsdp_inference = False - if should_disable_dit_cpu_offload and self.dit_cpu_offload is not False: - logger.warning( - "layerwise offload is selected for DiT components, automatically disabling dit_cpu_offload." - ) - self.dit_cpu_offload = False - - if envs.SGLANG_CACHE_DIT_ENABLED and should_disable_dit_cpu_offload: + if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected: raise ValueError( "DiT layerwise offload cannot be enabled together with cache-dit. " "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, " 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 d8ddd7616..6071d22d1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -559,7 +559,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.text_encoder_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["image_encoder", "vae"]) - def test_layerwise_components_disable_matching_cpu_offloads(self): + def test_layerwise_components_disable_matching_non_dit_cpu_offloads(self): args = self._from_dict_with_task_type( ModelTaskType.T2V, memory_gb=16, @@ -580,11 +580,34 @@ class TestOffloadDefaults(unittest.TestCase): args._adjust_layerwise_offload_components() self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.dit_cpu_offload) + # dit_cpu_offload is complementary to DiT layerwise offload (keeps + # weights off-device during load), so it must be preserved here. + 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) + def test_dit_layerwise_offload_preserves_dit_cpu_offload(self): + """Combining --dit-cpu-offload with --dit-layerwise-offload must keep both on. + + dit_cpu_offload controls initial residency (host memory), while + dit_layerwise_offload only swaps layers on/off device at inference. + Force-disabling dit_cpu_offload here would push the full DiT to GPU at + load time and OOM low-VRAM cards. + """ + args = self._from_dict_with_task_type( + ModelTaskType.T2I, + memory_gb=32, + kwargs={ + "dit_cpu_offload": True, + "dit_layerwise_offload": True, + }, + ) + + self.assertTrue(args.dit_cpu_offload) + self.assertTrue(args.dit_layerwise_offload) + self.assertEqual(args.layerwise_offload_components, ["dit"]) + def test_pipeline_configs_declare_auto_tune_hints(self): qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config() wan_deployment = WanT2V480PConfig().get_model_deployment_config() @@ -704,7 +727,10 @@ class TestOffloadDefaults(unittest.TestCase): self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) + # dit_cpu_offload is complementary to DiT layerwise offload: + # layerwise only moves layers on/off device at runtime, while + # dit_cpu_offload keeps the initial weights on host memory. + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual(args.dit_offload_prefetch_size, 2) @@ -744,7 +770,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertTrue(args.layerwise_offload_components) self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual( @@ -839,7 +865,9 @@ class TestOffloadDefaults(unittest.TestCase): }, ) - self.assertFalse(args.dit_cpu_offload) + # dit_cpu_offload defaults to True from _adjust_offload and is now + # preserved alongside DiT layerwise offload (the two are complementary). + self.assertTrue(args.dit_cpu_offload) self.assertEqual(args.layerwise_offload_components, ["dit"]) def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self): @@ -875,7 +903,7 @@ class TestOffloadDefaults(unittest.TestCase): ) self.assertFalse(args.use_fsdp_inference) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertTrue(args.layerwise_offload_components) self.assertTrue(args.text_encoder_cpu_offload) self.assertTrue(args.image_encoder_cpu_offload) @@ -1124,7 +1152,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.use_fsdp_inference) self.assertTrue(args.layerwise_offload_components) - self.assertFalse(args.dit_cpu_offload) + self.assertTrue(args.dit_cpu_offload) self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) self.assertEqual(