[diffusion] multi-platform: support diffusion on amd and fix encoder loading on MI325 (#13760)
Co-authored-by: Sabre Shao <sabre.shao@amd.com> Co-authored-by: Yusheng (Ethan) Su <yushengsu.thu@gmail.com> Co-authored-by: Hubert Lu <Hubert.Lu@amd.com> Co-authored-by: xsun <sunxiao04@gmail.com>
This commit is contained in:
@@ -273,7 +273,13 @@ class TextEncoderLoader(ComponentLoader):
|
||||
|
||||
def should_offload(self, server_args, model_config: ModelConfig | None = None):
|
||||
should_offload = server_args.text_encoder_cpu_offload
|
||||
fsdp_shard_conditions = getattr(model_config, "_fsdp_shard_conditions", [])
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
@@ -477,7 +483,13 @@ class TextEncoderLoader(ComponentLoader):
|
||||
class ImageEncoderLoader(TextEncoderLoader):
|
||||
def should_offload(self, server_args, model_config: ModelConfig | None = None):
|
||||
should_offload = server_args.image_encoder_cpu_offload
|
||||
fsdp_shard_conditions = getattr(model_config, "_fsdp_shard_conditions", [])
|
||||
# _fsdp_shard_conditions is in arch_config, not directly on model_config
|
||||
arch_config = (
|
||||
getattr(model_config, "arch_config", model_config) if model_config else None
|
||||
)
|
||||
fsdp_shard_conditions = (
|
||||
getattr(arch_config, "_fsdp_shard_conditions", []) if arch_config else []
|
||||
)
|
||||
use_cpu_offload = should_offload and len(fsdp_shard_conditions) > 0
|
||||
return use_cpu_offload
|
||||
|
||||
|
||||
@@ -23,6 +23,11 @@ try:
|
||||
except ImportError:
|
||||
HAS_RUNAI_MODEL_STREAMER = False
|
||||
|
||||
# Disable runai_model_streamer on AMD/ROCm due to global state issues
|
||||
# that cause "Streamer is handling previous request" errors
|
||||
if torch.version.hip is not None:
|
||||
HAS_RUNAI_MODEL_STREAMER = False
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
|
||||
Reference in New Issue
Block a user