[diffusion] feat: performance_mode=speed enables torch.compile by default (#30016)
This commit is contained in:
@@ -73,6 +73,17 @@ class ServerArgsAutoTuner:
|
|||||||
|
|
||||||
if args.performance_mode == "speed":
|
if args.performance_mode == "speed":
|
||||||
logger.info("Applying performance_mode=speed")
|
logger.info("Applying performance_mode=speed")
|
||||||
|
if not args.enable_torch_compile and not args.is_arg_explicitly_set(
|
||||||
|
"enable_torch_compile"
|
||||||
|
):
|
||||||
|
# speed means fastest: compile by default. An explicit
|
||||||
|
# --enable-torch-compile false still wins (e.g. models where
|
||||||
|
# compile measures slower, like short-step Z-Image runs).
|
||||||
|
args.enable_torch_compile = True
|
||||||
|
logger.info(
|
||||||
|
"performance_mode=speed enables torch.compile "
|
||||||
|
"(pass --enable-torch-compile false to opt out)"
|
||||||
|
)
|
||||||
if args.num_gpus >= 2 and self._can_apply_fsdp_policy(
|
if args.num_gpus >= 2 and self._can_apply_fsdp_policy(
|
||||||
require_memory_headroom=False
|
require_memory_headroom=False
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -1477,6 +1477,40 @@ class TestOffloadDefaults(unittest.TestCase):
|
|||||||
self.assertFalse(args.text_encoder_cpu_offload)
|
self.assertFalse(args.text_encoder_cpu_offload)
|
||||||
self.assertFalse(args.image_encoder_cpu_offload)
|
self.assertFalse(args.image_encoder_cpu_offload)
|
||||||
|
|
||||||
|
def test_speed_mode_enables_torch_compile_by_default(self):
|
||||||
|
args = self._from_dict_with_pipeline_config(
|
||||||
|
QwenImagePipelineConfig(),
|
||||||
|
kwargs={
|
||||||
|
"model_path": "Qwen/Qwen-Image",
|
||||||
|
"performance_mode": "speed",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(args.enable_torch_compile)
|
||||||
|
|
||||||
|
def test_speed_mode_preserves_explicit_torch_compile_off(self):
|
||||||
|
args = self._from_dict_with_pipeline_config(
|
||||||
|
QwenImagePipelineConfig(),
|
||||||
|
kwargs={
|
||||||
|
"model_path": "Qwen/Qwen-Image",
|
||||||
|
"performance_mode": "speed",
|
||||||
|
"enable_torch_compile": False,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(args.enable_torch_compile)
|
||||||
|
|
||||||
|
def test_auto_mode_leaves_torch_compile_off(self):
|
||||||
|
args = self._from_dict_with_pipeline_config(
|
||||||
|
QwenImagePipelineConfig(),
|
||||||
|
kwargs={
|
||||||
|
"model_path": "Qwen/Qwen-Image",
|
||||||
|
"performance_mode": "auto",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(args.enable_torch_compile)
|
||||||
|
|
||||||
def test_memory_mode_wan_uses_layerwise_offload(self):
|
def test_memory_mode_wan_uses_layerwise_offload(self):
|
||||||
args = self._from_dict_with_pipeline_config(
|
args = self._from_dict_with_pipeline_config(
|
||||||
WanT2V480PConfig(),
|
WanT2V480PConfig(),
|
||||||
|
|||||||
Reference in New Issue
Block a user