[diffusion] fix: reject cache-dit with fsdp (#28834)
This commit is contained in:
@@ -1945,6 +1945,20 @@ class ServerArgs(DisaggServerArgsMixin):
|
|||||||
)
|
)
|
||||||
|
|
||||||
# validate layerwise offload conflicts
|
# validate layerwise offload conflicts
|
||||||
|
if envs.SGLANG_CACHE_DIT_ENABLED and self.use_fsdp_inference:
|
||||||
|
if self.is_arg_explicitly_set("use_fsdp_inference"):
|
||||||
|
raise ValueError(
|
||||||
|
"FSDP inference cannot be enabled together with cache-dit. "
|
||||||
|
"cache-dit wraps known DiT block structures, while FSDP wraps "
|
||||||
|
"and shards modules before cache-dit can inspect them. "
|
||||||
|
"Please disable --use-fsdp-inference or disable "
|
||||||
|
"SGLANG_CACHE_DIT_ENABLED."
|
||||||
|
)
|
||||||
|
logger.warning(
|
||||||
|
"cache-dit is enabled, automatically disabling use_fsdp_inference."
|
||||||
|
)
|
||||||
|
self.use_fsdp_inference = False
|
||||||
|
|
||||||
if self.layerwise_offload_components:
|
if self.layerwise_offload_components:
|
||||||
if self.dit_offload_prefetch_size < 0.0:
|
if self.dit_offload_prefetch_size < 0.0:
|
||||||
raise ValueError("dit_offload_prefetch_size must be non-negative")
|
raise ValueError("dit_offload_prefetch_size must be non-negative")
|
||||||
|
|||||||
@@ -833,6 +833,32 @@ class TestOffloadDefaults(unittest.TestCase):
|
|||||||
self.assertTrue(args.use_fsdp_inference)
|
self.assertTrue(args.use_fsdp_inference)
|
||||||
self.assertTrue(args.enable_cfg_parallel)
|
self.assertTrue(args.enable_cfg_parallel)
|
||||||
|
|
||||||
|
def test_cache_dit_rejects_explicit_fsdp(self):
|
||||||
|
with patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}):
|
||||||
|
with self.assertRaisesRegex(ValueError, "FSDP inference"):
|
||||||
|
self._from_dict_with_pipeline_config(
|
||||||
|
SanaWMPipelineConfig(),
|
||||||
|
kwargs={
|
||||||
|
"model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
|
||||||
|
"num_gpus": 2,
|
||||||
|
"use_fsdp_inference": True,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_cache_dit_auto_disables_implicit_fsdp(self):
|
||||||
|
with patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}):
|
||||||
|
args = self._from_dict_with_pipeline_config(
|
||||||
|
SanaWMPipelineConfig(),
|
||||||
|
kwargs={
|
||||||
|
"model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
|
||||||
|
"num_gpus": 2,
|
||||||
|
"performance_mode": "auto",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(args.use_fsdp_inference)
|
||||||
|
self.assertTrue(args.enable_cfg_parallel)
|
||||||
|
|
||||||
def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self):
|
def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self):
|
||||||
args = self._from_dict_with_pipeline_config(
|
args = self._from_dict_with_pipeline_config(
|
||||||
SanaWMRealtimeConfig(),
|
SanaWMRealtimeConfig(),
|
||||||
|
|||||||
Reference in New Issue
Block a user