[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:
Yuzhen Zhou
2025-12-19 02:38:46 -05:00
committed by GitHub
parent f2d64e6782
commit 4bf06635fc
34 changed files with 823 additions and 72 deletions

View File

@@ -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

View File

@@ -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