turn off dit_layerwise_offload for wan on rocm (#17569)
This commit is contained in:
@@ -343,6 +343,11 @@ class Platform:
|
|||||||
"""Get the CPU architecture of the current platform."""
|
"""Get the CPU architecture of the current platform."""
|
||||||
return CpuArchEnum.UNSPECIFIED
|
return CpuArchEnum.UNSPECIFIED
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
|
||||||
|
"""Whether to enable DIT layerwise offload by default on the current platform."""
|
||||||
|
return True
|
||||||
|
|
||||||
def get_attn_backend(self, *args, **kwargs) -> AttentionImpl:
|
def get_attn_backend(self, *args, **kwargs) -> AttentionImpl:
|
||||||
attention_cls_str = self.get_attn_backend_cls_str(*args, **kwargs)
|
attention_cls_str = self.get_attn_backend_cls_str(*args, **kwargs)
|
||||||
return resolve_obj_by_qualname(attention_cls_str)
|
return resolve_obj_by_qualname(attention_cls_str)
|
||||||
|
|||||||
@@ -165,3 +165,8 @@ class RocmPlatform(Platform):
|
|||||||
@classmethod
|
@classmethod
|
||||||
def get_device_communicator_cls(cls) -> str:
|
def get_device_communicator_cls(cls) -> str:
|
||||||
return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too
|
return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # works for ROCm too
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
|
||||||
|
"""ROCm performs better without DIT layerwise offload on Wan."""
|
||||||
|
return False
|
||||||
|
|||||||
@@ -955,6 +955,7 @@ class ServerArgs:
|
|||||||
if (
|
if (
|
||||||
"wan" in self.pipeline_config.__class__.__name__.lower()
|
"wan" in self.pipeline_config.__class__.__name__.lower()
|
||||||
and self.dit_layerwise_offload is None
|
and self.dit_layerwise_offload is None
|
||||||
|
and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
|
||||||
):
|
):
|
||||||
logger.info(
|
logger.info(
|
||||||
"Automatically enable dit_layerwise_offload for Wan for best performance"
|
"Automatically enable dit_layerwise_offload for Wan for best performance"
|
||||||
|
|||||||
Reference in New Issue
Block a user