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

[diffusion] feat: allow --dit-cpu-offload with --dit-layerwise-offload (#26925)

Browse Source 
Co-authored-by: Yiqi Yang <yiqi.yang@kiwiar.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Yiqi Yang
2026-06-01 23:17:53 +08:00
committed by GitHub
co-authored by Yiqi Yang Claude Opus 4.7
parent 693adabff7
commit 89feb18eb9
2 changed files with 44 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/server_args.py
+9 -16
View File
@@ -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, "
python/sglang/multimodal_gen/test/unit/test_server_args.py
+35 -7
View File
@@ -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(