minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion][CI]: Add individual component accuracy CI for diffusion models (#18709)

Browse Source 
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Ratish P
2026-04-01 21:51:36 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent e67b95d66b
commit 4f5b55e379
13 changed files with 2670 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+31 -12
View File
@@ -156,7 +156,9 @@ class TextEncoderLoader(ComponentLoader):
return hf_folder, hf_weights_files, use_safetensors
def _get_weights_iterator(
self, source: "Source", to_cpu: bool
self,
source: "Source",
to_cpu: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""get an iterator for the model weights based on the load format."""
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
@@ -166,7 +168,8 @@ class TextEncoderLoader(ComponentLoader):
)
if use_safetensors:
weights_iterator = safetensors_weights_iterator(
hf_weights_files, to_cpu=to_cpu
hf_weights_files,
to_cpu=to_cpu,
)
else:
weights_iterator = pt_weights_iterator(hf_weights_files, to_cpu=to_cpu)
@@ -186,17 +189,27 @@ class TextEncoderLoader(ComponentLoader):
fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
)
yield from self._get_weights_iterator(primary_weights, to_cpu)
yield from self._get_weights_iterator(
primary_weights,
to_cpu,
)
secondary_weights = cast(
Iterable[TextEncoderLoader.Source],
getattr(model, "secondary_weights", ()),
)
for source in secondary_weights:
yield from self._get_weights_iterator(source, to_cpu)
yield from self._get_weights_iterator(
source,
to_cpu,
)
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str,
cpu_offload_flag: bool | None = None,
):
"""Load the text encoders based on the model path, and inference args."""
diffusers_pretrained_config = get_config(
@@ -227,6 +240,7 @@ class TextEncoderLoader(ComponentLoader):
encoder_config,
server_args,
encoder_dtype,
cpu_offload_flag=cpu_offload_flag,
)
def load_model(
@@ -240,7 +254,10 @@ class TextEncoderLoader(ComponentLoader):
# Determine CPU offload behavior and target device
local_torch_device = get_local_torch_device()
should_offload = self.should_offload(server_args, model_config)
fsdp_cpu_offload = self.should_offload(server_args, model_config)
should_offload = (
cpu_offload_flag if cpu_offload_flag is not None else fsdp_cpu_offload
)
if should_offload and not current_platform.is_mps():
model_device = torch.device("cpu")
@@ -263,13 +280,13 @@ class TextEncoderLoader(ComponentLoader):
weights_to_load = {name for name, _ in model.named_parameters()}
loaded_weights = model.load_weights(
self._get_all_weights(model, model_path, to_cpu=should_offload)
self._get_all_weights(
model,
model_path,
to_cpu=should_offload,
)
)
# Explicitly move model to target device after loading weights
if not should_offload:
model = model.to(local_torch_device)
if should_offload:
# Disable FSDP for MPS as it's not compatible
if current_platform.is_mps():
@@ -277,7 +294,7 @@ class TextEncoderLoader(ComponentLoader):
"Disabling FSDP sharding for MPS platform as it's not compatible"
)
model = model.to(local_torch_device)
else:
elif fsdp_cpu_offload:
mesh = init_device_mesh(
current_platform.device_type,
mesh_shape=(1, dist.get_world_size()),
@@ -292,6 +309,8 @@ class TextEncoderLoader(ComponentLoader):
or getattr(model, "_fsdp_shard_conditions", None),
pin_cpu_memory=server_args.pin_cpu_memory,
)
else:
model = model.to("cpu")
else:
model = model.to(local_torch_device)
# We only enable strict check for non-quantized models
python/sglang/multimodal_gen/runtime/loader/weight_utils.py
+6 -1
View File
@@ -24,6 +24,7 @@ try:
except ImportError:
HAS_RUNAI_MODEL_STREAMER = False
from sglang.multimodal_gen import envs
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -135,13 +136,17 @@ def _validate_safetensors_file(file_path: str) -> bool:
def safetensors_weights_iterator(
hf_weights_files: list[str],
to_cpu: bool = True,
use_runai_model_streamer: bool = HAS_RUNAI_MODEL_STREAMER,
use_runai_model_streamer: bool | None = None,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Iterate over the weights in the model safetensor files."""
enable_tqdm = (
not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
)
device = "cpu" if to_cpu else str(get_local_torch_device())
if use_runai_model_streamer is None:
use_runai_model_streamer = (
HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER
)
# Validate files before loading
corrupted_files = [
python/sglang/multimodal_gen/test/server/accuracy_config.py
+386
View File
@@ -0,0 +1,386 @@
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from typing import Dict, Optional
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
class ComponentType(str, Enum):
VAE = "vae"
TRANSFORMER = "transformer"
TEXT_ENCODER = "text_encoder"
@dataclass(frozen=True)
class ComponentSkip:
reason: str
DEFAULT_TIMESTEP = 500.0
TIMESTEP_NORMALIZATION_FACTOR = 1000.0
I2V_IMAGE_DIM = 1280
I2V_TEXT_ENCODER_DIM = 5120
DEFAULT_TEXT_ENCODER_VOCAB_SIZE = 32000
TEXT_ENCODER_INPUT_SEED = 42
TEXT_ENCODER_TOKEN_MIN = 100
TEXT_ENCODER_TOKEN_MAX = 30000
TEXT_ENCODER_TOKEN_LENGTH = 32
# Default thresholds by component. Override per component/case if needed.
DEFAULT_THRESHOLDS = {
ComponentType.VAE: 0.999,
ComponentType.TRANSFORMER: 0.995,
ComponentType.TEXT_ENCODER: 0.98,
}
# Optional per-case overrides: {case_id: {ComponentType: threshold}}
CASE_THRESHOLDS: Dict[str, Dict[ComponentType, float]] = {
# Add overrides here when a specific model/component needs a different threshold.
"flux_2_image_t2i": {ComponentType.TRANSFORMER: 0.99},
"flux_2_image_t2i_layerwise_offload": {ComponentType.TRANSFORMER: 0.99},
"flux_2_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.99},
"flux_2_klein_ti2i_2_gpus": {ComponentType.TRANSFORMER: 0.975},
"flux_2_ti2i": {ComponentType.TRANSFORMER: 0.99},
"flux_2_t2i_customized_vae_path": {ComponentType.TRANSFORMER: 0.99},
"fast_hunyuan_video": {ComponentType.TRANSFORMER: 0.99},
"fsdp-inference": {ComponentType.TRANSFORMER: 0.9935},
"wan2_2_i2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99},
"wan2_2_t2v_a14b_2gpu": {ComponentType.TRANSFORMER: 0.99},
"wan2_2_t2v_a14b_teacache_2gpu": {ComponentType.TRANSFORMER: 0.99},
"wan2_2_t2v_a14b_lora_2gpu": {ComponentType.TRANSFORMER: 0.99},
"zimage_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.9935},
"zimage_image_t2i_2_gpus_non_square": {ComponentType.TRANSFORMER: 0.9935},
}
# Active skip policy. Keep this limited to cases with current, concrete evidence
# of real divergence or unsupported reference loading in the harness.
SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
"flux_image_t2i": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline despite 100% matched weights (CosSim ~0.47)"
)
},
"sana_image_t2i": {
ComponentType.VAE: ComponentSkip(
"HF AutoencoderDC checkpoint leaves required to_qkv_multiscale weights missing, so VAE transfer would compare against partially initialized reference weights"
)
},
"mova_360p_1gpu": {
ComponentType.TRANSFORMER: ComponentSkip(
"HF reference transformer cannot be materialized from the video_dit repo layout"
)
},
"qwen_image_t2i_cache_dit_enabled": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by qwen_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by qwen_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by qwen_image_t2i for the same source component and topology"
),
},
"flux_2_image_t2i_upscaling_4x": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
},
"layerwise_offload": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
},
"zimage_image_t2i_fp8": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
},
"zimage_image_t2i_multi_lora": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by zimage_image_t2i for the same source component and topology"
),
},
"flux_2_ti2i": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
},
"flux_2_t2i_customized_vae_path": {
ComponentType.VAE: ComponentSkip(
"Customized VAE override points to FLUX.2 Tiny AutoEncoder, but the HF reference loader does not yet materialize a trustworthy matching VAE baseline"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_teacache_enabled": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_frame_interp_2x": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_upscaling_4x": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1_3b_lora_1gpu": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"flux_2_ti2i_multi_image_cache_dit": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
},
"wan2_2_ti2v_5b": {
ComponentType.TRANSFORMER: ComponentSkip(
"SGLang transformer loader rejects new parameters in HF checkpoint"
)
},
"fastwan2_2_ti2v_5b": {
ComponentType.TRANSFORMER: ComponentSkip(
"SGLang transformer loader rejects new parameters in HF checkpoint"
)
},
"turbo_wan2_1_t2v_1.3b": {
ComponentType.TRANSFORMER: ComponentSkip(
"Weight transfer match ratio too low for reliable comparison"
)
},
"wan2_1_i2v_14b_480P_2gpu": {
ComponentType.TRANSFORMER: ComponentSkip(
"Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.71) after full weight transfer and matching output shape"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"wan2_1_i2v_14b_lora_2gpu": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_i2v_14b_720P_2gpu for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.68) after full weight transfer and matching output shape"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"wan2_1_i2v_14b_720P_2gpu": {
ComponentType.TRANSFORMER: ComponentSkip(
"Transformer diverges from Diffusers baseline in 2-GPU accuracy run (CosSim ~0.68) after full weight transfer and matching output shape"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"wan2_2_i2v_a14b_2gpu": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
)
},
"wan2_2_t2v_a14b_2gpu": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
)
},
"wan2_2_t2v_a14b_teacache_2gpu": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"wan2_2_t2v_a14b_lora_2gpu": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_2_t2v_a14b_2gpu for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"wan2_1_t2v_14b_2gpu": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU SP-folded accuracy run (CosSim ~0.31) after 100% matched weight transfer"
)
},
"wan2_1_t2v_1.3b_cfg_parallel": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
)
},
"mova_360p_tp2": {
ComponentType.TRANSFORMER: ComponentSkip(
"HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"mova_360p_ring1_uly2": {
ComponentType.TRANSFORMER: ComponentSkip(
"HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"mova_360p_ring2_uly1": {
ComponentType.TRANSFORMER: ComponentSkip(
"HF reference transformer cannot be materialized from the MOVA video_dit repo layout"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.31) after 100% matched weight transfer"
),
},
"flux_image_t2i_2_gpus": {
ComponentType.TEXT_ENCODER: ComponentSkip(
"Text encoder diverges from HF baseline in 2-GPU accuracy run (CosSim ~0.47) after 100% matched weight transfer"
)
},
"zimage_image_t2i_2_gpus_non_square": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by zimage_image_t2i_2_gpus for the same source component and topology"
),
},
"flux_2_image_t2i_2_gpus": {
ComponentType.TRANSFORMER: ComponentSkip(
"2-GPU FLUX.2 transformer diverges strongly from Diffusers baseline (CosSim ~0.54) despite full weight transfer"
)
},
"hunyuan3d_shape_gen": {
ComponentType.VAE: ComponentSkip(
"HF config cannot be parsed as valid JSON for component reference loading"
),
ComponentType.TRANSFORMER: ComponentSkip(
"HF config cannot be parsed as valid JSON for component reference loading"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"HF config cannot be parsed as valid JSON for component reference loading"
),
},
}
# TODO: If a model needs extra compatibility logic, prefer adding a skip or an
# explicit override here instead of adding more ad-hoc hacks in the engine.
def get_threshold(case_id: str, component: ComponentType) -> float:
overrides = CASE_THRESHOLDS.get(case_id, {})
return overrides.get(component, DEFAULT_THRESHOLDS[component])
def get_skip_reason(case: DiffusionTestCase, component: ComponentType) -> Optional[str]:
skip_entry = SKIP_COMPONENTS.get(case.id, {}).get(component)
if skip_entry is None:
return None
return skip_entry.reason
def should_skip_component(case: DiffusionTestCase, component: ComponentType) -> bool:
return get_skip_reason(case, component) is not None
python/sglang/multimodal_gen/test/server/accuracy_hooks.py
+579
View File
@@ -0,0 +1,579 @@
from __future__ import annotations
import inspect
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Optional
import torch
import torch.nn as nn
from sglang.multimodal_gen.test.server.accuracy_config import (
DEFAULT_TIMESTEP,
I2V_IMAGE_DIM,
TIMESTEP_NORMALIZATION_FACTOR,
ComponentType,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
extract_output_tensor,
seed_and_broadcast,
)
Inputs = Dict[str, Any]
BuildInputsFn = Callable[[Any, nn.Module, str, Optional[nn.Module]], Inputs]
PrepareCallFn = Callable[[nn.Module, Inputs], "HookCall"]
NormalizeFn = Callable[[Any], torch.Tensor]
# These are harness defaults for synthetic accuracy inputs.
# They are not checkpoint truth. We use them only when the model config or
# forward signature does not expose a more specific shape or channel count.
DEFAULT_TEXT_SEQ_LEN = 64
DEFAULT_TOKEN_LAYOUT_SIZE = 32
REDUCED_TOKEN_LAYOUT_SIZE = 16
DEFAULT_VIDEO_FRAME_COUNT = 4
DEFAULT_IMAGE_TOKEN_COUNT = 257
ALIAS_ROTARY_TEXT_PAD_MULTIPLE = 32
DEFAULT_TRANSFORMER_IN_CHANNELS = 16
DEFAULT_TRANSFORMER_TEXT_CHANNELS = 4096
DEFAULT_TRANSFORMER_POOLED_CHANNELS = 768
DEFAULT_VAE_LATENT_CHANNELS = 16
DEFAULT_VAE_LATENT_SPATIAL_SIZE = 32
LARGE_CHANNEL_LAYOUT_THRESHOLD = 128
@dataclass(frozen=True)
class TransformerHookCompat:
normalize_reference_timestep: bool = False
negate_reference_output: bool = False
omit_reference_guidance: bool = False
use_2d_hidden_states: bool = False
def _resolve_transformer_hook_compat(case: Any) -> TransformerHookCompat:
model_path = case.server_args.model_path.lower()
if "z-image" in model_path:
return TransformerHookCompat(
normalize_reference_timestep=True,
negate_reference_output=True,
)
if "qwen" in model_path:
return TransformerHookCompat(
normalize_reference_timestep=True,
omit_reference_guidance=True,
)
if "sana" in model_path:
return TransformerHookCompat(
omit_reference_guidance=True,
use_2d_hidden_states=True,
)
if "flux" in model_path:
return TransformerHookCompat(normalize_reference_timestep=True)
return TransformerHookCompat()
@dataclass
class HookCall:
module: nn.Module
args: tuple[Any, ...] = ()
kwargs: Dict[str, Any] = field(default_factory=dict)
negate_output: bool = False
@dataclass(frozen=True)
class NativeHookProfile:
build_inputs: BuildInputsFn
prepare_sglang_call: PrepareCallFn
prepare_reference_call: PrepareCallFn
normalize_sglang_output: NormalizeFn = extract_output_tensor
normalize_reference_output: NormalizeFn = extract_output_tensor
class _DeterministicRNG:
def __init__(self, seed: int = 42) -> None:
self._seed = seed
def randn(
self, shape: tuple[int, ...], device: str, dtype: torch.dtype
) -> torch.Tensor:
torch.manual_seed(self._seed)
tensor = torch.randn(shape, device="cpu", dtype=dtype).to(device)
seed_and_broadcast(self._seed, tensor)
self._seed += 1
return tensor
def _resolve_nested_attr(obj: Any, path: str) -> Any:
current = obj
for name in path.split("."):
if current is None or not hasattr(current, name):
return None
current = getattr(current, name)
return current
def _read_config_value(model: nn.Module, keys: list[str], default: int) -> int:
config = getattr(model, "config", None)
for key in keys:
for root in (model, config):
value = _resolve_nested_attr(root, key) if root is not None else None
if isinstance(value, int) and value > 0:
return value
return default
def _forward_parameter_names(module: nn.Module) -> set[str]:
return set(inspect.signature(module.forward).parameters.keys())
def _infer_transformer_layout(param_names: set[str]) -> str:
if "img_shapes" in param_names or "txt_seq_lens" in param_names:
return "token_shapes"
if "img_ids" in param_names or "txt_ids" in param_names:
return "token_ids"
if "x" in param_names or "cap_feats" in param_names:
return "alias"
return "video"
def _build_position_ids(
height: int, width: int, dims: int, device: str
) -> tuple[torch.Tensor, torch.Tensor]:
img_len = height * width
txt_len = DEFAULT_TEXT_SEQ_LEN
if dims == 4:
img_ids = torch.zeros(img_len, 4, device=device, dtype=torch.bfloat16)
img_ids[:, 1] = torch.arange(height).repeat_interleave(width)
img_ids[:, 2] = torch.arange(width).repeat(height)
txt_ids = torch.zeros(txt_len, 4, device=device, dtype=torch.bfloat16)
else:
img_ids = torch.zeros(img_len, 3, device=device, dtype=torch.bfloat16)
img_ids[:, 0] = torch.arange(height).repeat_interleave(width)
img_ids[:, 1] = torch.arange(width).repeat(height)
txt_ids = torch.zeros(txt_len, 3, device=device, dtype=torch.bfloat16)
return img_ids, txt_ids
def _build_alias_rotary_freqs(
model: nn.Module, device: str, height: int, width: int
) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
cap_len = DEFAULT_TEXT_SEQ_LEN
cap_pad_len = (-cap_len) % ALIAS_ROTARY_TEXT_PAD_MULTIPLE
cap_ids = (
torch.stack(
torch.meshgrid(
torch.arange(cap_len + cap_pad_len),
torch.arange(1),
torch.arange(1),
indexing="ij",
),
dim=-1,
)
.flatten(0, 2)
.to(device)
)
img_ids = (
torch.stack(
torch.meshgrid(
torch.arange(1),
torch.arange(height // 2),
torch.arange(width // 2),
indexing="ij",
),
dim=-1,
)
.flatten(0, 2)
.to(device)
)
cos_cap, sin_cap = model.rotary_emb(cap_ids)
cos_img, sin_img = model.rotary_emb(img_ids)
return ((cos_cap, sin_cap), (cos_img, sin_img))
def _supports_image_conditioning(module: nn.Module) -> bool:
image_embedder = _resolve_nested_attr(module, "condition_embedder.image_embedder")
if image_embedder is not None:
return True
image_dim = _read_config_value(
module, ["arch_config.image_dim", "image_dim"], default=0
)
return image_dim > 0
def _build_transformer_hook_inputs(
case: Any, model: nn.Module, device: str, ref_model: Optional[nn.Module] = None
) -> Inputs:
"""Build one synthetic input bundle that both transformer variants can consume."""
compat = _resolve_transformer_hook_compat(case)
param_names = _forward_parameter_names(model)
if ref_model is not None:
# The input bundle has to satisfy both call signatures.
param_names.update(_forward_parameter_names(ref_model))
rng = _DeterministicRNG()
layout = _infer_transformer_layout(param_names)
in_channels = _read_config_value(
model,
[
"arch_config.in_channels",
"in_channels",
"transformer_config.in_channels",
],
default=DEFAULT_TRANSFORMER_IN_CHANNELS,
)
text_channels = _read_config_value(
model,
[
"text_states_dim",
"arch_config.cap_feat_dim",
"cap_feat_dim",
"caption_channels",
"arch_config.text_dim",
"text_dim",
"arch_config.text_embed_dim",
"text_embed_dim",
"arch_config.joint_attention_dim",
"joint_attention_dim",
"cross_attention_dim",
"hidden_size",
"dim",
],
default=DEFAULT_TRANSFORMER_TEXT_CHANNELS,
)
pooled_channels = _read_config_value(
model,
[
"text_states_dim_2",
"arch_config.pooled_projection_dim",
"pooled_projection_dim",
"pooled_embed_dim",
"text_embed_dim",
"projection_dim",
],
default=DEFAULT_TRANSFORMER_POOLED_CHANNELS,
)
image_channels = _read_config_value(
model,
["arch_config.image_dim", "image_dim", "cross_attention_dim"],
default=I2V_IMAGE_DIM,
)
if layout == "token_shapes":
height, width = DEFAULT_TOKEN_LAYOUT_SIZE, DEFAULT_TOKEN_LAYOUT_SIZE
seq_len = (height // 2) * (width // 2)
hidden_states = rng.randn((1, seq_len, in_channels), device, torch.bfloat16)
elif layout == "token_ids":
height, width = REDUCED_TOKEN_LAYOUT_SIZE, REDUCED_TOKEN_LAYOUT_SIZE
seq_len = height * width
hidden_states = rng.randn((1, seq_len, in_channels), device, torch.bfloat16)
elif layout == "alias":
height, width = DEFAULT_TOKEN_LAYOUT_SIZE, DEFAULT_TOKEN_LAYOUT_SIZE
hidden_states = rng.randn(
(1, in_channels, 1, height, width), device, torch.bfloat16
)
elif compat.use_2d_hidden_states:
spatial_size = (
REDUCED_TOKEN_LAYOUT_SIZE
if "encoder_attention_mask" in param_names
or "encoder_hidden_states_mask" in param_names
else DEFAULT_TOKEN_LAYOUT_SIZE
)
height, width = spatial_size, spatial_size
hidden_states = rng.randn(
(1, in_channels, height, width),
device,
torch.bfloat16,
)
else:
spatial_size = (
REDUCED_TOKEN_LAYOUT_SIZE
if "encoder_attention_mask" in param_names
or "encoder_hidden_states_mask" in param_names
else DEFAULT_TOKEN_LAYOUT_SIZE
)
height, width = spatial_size, spatial_size
hidden_states = rng.randn(
(1, in_channels, DEFAULT_VIDEO_FRAME_COUNT, height, width),
device,
torch.bfloat16,
)
inputs: Inputs = {
"hidden_states": hidden_states,
"encoder_hidden_states": rng.randn(
(1, DEFAULT_TEXT_SEQ_LEN, text_channels), device, torch.bfloat16
),
"timestep": torch.tensor(
[DEFAULT_TIMESTEP], device=device, dtype=torch.bfloat16
),
"guidance": torch.tensor([1.0], device=device, dtype=torch.bfloat16),
}
if "pooled_projections" in param_names:
inputs["pooled_projections"] = rng.randn(
(1, pooled_channels), device, torch.bfloat16
)
if (
"encoder_attention_mask" in param_names
or "encoder_hidden_states_mask" in param_names
):
attention_mask = torch.ones(
1, DEFAULT_TEXT_SEQ_LEN, device=device, dtype=torch.bool
)
inputs["encoder_attention_mask"] = attention_mask
inputs["encoder_hidden_states_mask"] = attention_mask
if "encoder_hidden_states_image" in param_names and _supports_image_conditioning(
model
):
inputs["encoder_hidden_states_image"] = rng.randn(
(1, DEFAULT_IMAGE_TOKEN_COUNT, image_channels), device, torch.bfloat16
)
if "additional_t_cond" in param_names:
inputs["additional_t_cond"] = torch.zeros((1,), device=device, dtype=torch.long)
if "img_shapes" in param_names:
inputs["img_shapes"] = [[(1, height // 2, width // 2)]]
if "txt_seq_lens" in param_names:
inputs["txt_seq_lens"] = [DEFAULT_TEXT_SEQ_LEN]
if "img_ids" in param_names or "txt_ids" in param_names:
id_dims = 4 if in_channels >= LARGE_CHANNEL_LAYOUT_THRESHOLD else 3
img_ids, txt_ids = _build_position_ids(height, width, id_dims, device)
inputs["img_ids"] = img_ids
inputs["txt_ids"] = txt_ids
if "freqs_cis" in param_names and hasattr(model, "rotary_emb"):
if "img_shapes" in inputs and "txt_seq_lens" in inputs:
img_freqs, txt_freqs = model.rotary_emb(
inputs["img_shapes"],
inputs["txt_seq_lens"],
device=hidden_states.device,
)
if torch.is_complex(img_freqs) and torch.is_complex(txt_freqs):
inputs["freqs_cis"] = (
torch.cat([img_freqs.real.float(), img_freqs.imag.float()], dim=-1),
torch.cat([txt_freqs.real.float(), txt_freqs.imag.float()], dim=-1),
)
else:
inputs["freqs_cis"] = (img_freqs, txt_freqs)
elif "img_ids" in inputs and "txt_ids" in inputs:
ids = torch.cat([inputs["txt_ids"], inputs["img_ids"]], dim=0)
inputs["freqs_cis"] = model.rotary_emb(ids)
elif inputs["hidden_states"].ndim == 5:
inputs["freqs_cis"] = _build_alias_rotary_freqs(
model, device, height, width
)
inputs["hook_compat"] = compat
return inputs
def _get_transformer_hook_compat(inputs: Inputs) -> TransformerHookCompat:
compat = inputs.get("hook_compat")
assert isinstance(compat, TransformerHookCompat)
return compat
def _supports_guidance_embedding(module: nn.Module) -> bool:
time_text_embed = getattr(module, "time_text_embed", None)
if time_text_embed is None:
return True
parameters = list(inspect.signature(time_text_embed.forward).parameters.values())
if any(param.kind is inspect.Parameter.VAR_POSITIONAL for param in parameters):
return True
accepted_args = [
param
for param in parameters
if param.name != "self"
and param.kind
in (
inspect.Parameter.POSITIONAL_ONLY,
inspect.Parameter.POSITIONAL_OR_KEYWORD,
inspect.Parameter.KEYWORD_ONLY,
)
]
return len(accepted_args) >= 3
def _prepare_transformer_hook_call(
module: nn.Module, inputs: Inputs, side: str
) -> HookCall:
param_names = _forward_parameter_names(module)
signature = inspect.signature(module.forward)
compat = _get_transformer_hook_compat(inputs)
kwargs: Dict[str, Any] = {}
negate_output = side == "reference" and compat.negate_reference_output
if "hidden_states" in param_names:
kwargs["hidden_states"] = inputs["hidden_states"]
if "x" in param_names:
kwargs["x"] = [inputs["hidden_states"].squeeze(0)]
if "encoder_hidden_states" in param_names:
encoder_value: Any = inputs["encoder_hidden_states"]
if (
side == "sglang"
and "pooled_projections" in inputs
and "pooled_projections" not in param_names
and "encoder_attention_mask" not in param_names
):
encoder_value = [
inputs["encoder_hidden_states"],
inputs["pooled_projections"],
]
kwargs["encoder_hidden_states"] = encoder_value
if "cap_feats" in param_names:
kwargs["cap_feats"] = [inputs["encoder_hidden_states"].squeeze(0)]
if "timestep" in param_names:
timestep = inputs["timestep"]
if side == "reference" and compat.normalize_reference_timestep:
timestep = timestep / TIMESTEP_NORMALIZATION_FACTOR
kwargs["timestep"] = timestep
if "t" in param_names:
timestep = inputs["timestep"]
if side == "reference" and compat.normalize_reference_timestep:
timestep = timestep / TIMESTEP_NORMALIZATION_FACTOR
kwargs["t"] = timestep
if "guidance" in param_names and "guidance" in inputs:
if side == "reference" and compat.omit_reference_guidance:
pass
else:
skip_guidance_for_image_context = (
"encoder_hidden_states_image" in param_names
and "img_ids" not in param_names
and "img_shapes" not in param_names
)
supports_guidance_embedding = _supports_guidance_embedding(module)
requires_guidance_arg = (
signature.parameters["guidance"].default is inspect._empty
)
should_include_guidance = (
not skip_guidance_for_image_context and supports_guidance_embedding
)
if should_include_guidance or requires_guidance_arg:
guidance_value = inputs["guidance"]
if side == "sglang":
guidance_value = guidance_value * TIMESTEP_NORMALIZATION_FACTOR
kwargs["guidance"] = guidance_value
if (
"encoder_hidden_states_image" in param_names
and "encoder_hidden_states_image" in inputs
):
value = inputs["encoder_hidden_states_image"]
kwargs["encoder_hidden_states_image"] = [value] if side == "sglang" else value
for key in (
"pooled_projections",
"img_ids",
"txt_ids",
"img_shapes",
"txt_seq_lens",
"freqs_cis",
"additional_t_cond",
"encoder_attention_mask",
"encoder_hidden_states_mask",
):
if key in param_names and key in inputs:
kwargs[key] = inputs[key]
if "return_dict" in param_names:
kwargs["return_dict"] = True
return HookCall(module=module, kwargs=kwargs, negate_output=negate_output)
def _prepare_transformer_sglang_call(module: nn.Module, inputs: Inputs) -> HookCall:
return _prepare_transformer_hook_call(module, inputs, side="sglang")
def _prepare_transformer_reference_call(module: nn.Module, inputs: Inputs) -> HookCall:
return _prepare_transformer_hook_call(module, inputs, side="reference")
class _VAEDecodeModule(nn.Module):
def __init__(self, vae: nn.Module):
super().__init__()
self.vae = vae
def forward(self, z: torch.Tensor) -> torch.Tensor:
if (
any(
isinstance(module, (nn.Conv3d, nn.ConvTranspose3d))
for module in self.vae.modules()
)
and z.ndim == 4
):
z = z.unsqueeze(2)
output = self.vae.decode(z)
tensor = output.sample if hasattr(output, "sample") else output
if isinstance(tensor, (list, tuple)):
tensor = tensor[0]
return tensor.squeeze(2) if tensor.ndim == 5 else tensor
def _infer_vae_latent_channels(model: nn.Module) -> int:
for path in ("post_quant_conv.in_channels", "post_quant_conv.conv.in_channels"):
value = _resolve_nested_attr(model, path)
if isinstance(value, int) and value > 0:
return value
return _read_config_value(
model,
[
"z_dim",
"arch_config.z_dim",
"latent_channels",
"arch_config.latent_channels",
"num_channels_latents",
"arch_config.num_channels_latents",
"latent_dim",
"z_channels",
"arch_config.z_channels",
],
default=DEFAULT_VAE_LATENT_CHANNELS,
)
def _build_vae_hook_inputs(
case: Any, model: nn.Module, device: str, ref_model: Optional[nn.Module] = None
) -> Inputs:
del case, ref_model
latent_channels = _infer_vae_latent_channels(model)
rng = _DeterministicRNG()
return {
"z": rng.randn(
(
1,
latent_channels,
DEFAULT_VAE_LATENT_SPATIAL_SIZE,
DEFAULT_VAE_LATENT_SPATIAL_SIZE,
),
device,
torch.bfloat16,
)
}
def _prepare_vae_decode_call(module: nn.Module, inputs: Inputs) -> HookCall:
return HookCall(module=_VAEDecodeModule(module), args=(inputs["z"],))
TRANSFORMER_NATIVE_PROFILE = NativeHookProfile(
build_inputs=_build_transformer_hook_inputs,
prepare_sglang_call=_prepare_transformer_sglang_call,
prepare_reference_call=_prepare_transformer_reference_call,
)
VAE_NATIVE_PROFILE = NativeHookProfile(
build_inputs=_build_vae_hook_inputs,
prepare_sglang_call=_prepare_vae_decode_call,
prepare_reference_call=_prepare_vae_decode_call,
)
def resolve_component_native_profile(component: ComponentType) -> NativeHookProfile:
if component == ComponentType.TRANSFORMER:
return TRANSFORMER_NATIVE_PROFILE
if component == ComponentType.VAE:
return VAE_NATIVE_PROFILE
raise KeyError(f"Unsupported native accuracy component: {component.value}")
python/sglang/multimodal_gen/test/server/accuracy_utils.py
+880
View File
@@ -0,0 +1,880 @@
from __future__ import annotations
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
import torch
import torch.nn as nn
from safetensors.torch import load_file as safetensors_load_file
from torch.distributed.tensor import distribute_tensor
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
destroy_model_parallel,
get_data_parallel_world_size,
get_sequence_parallel_world_size,
get_tensor_model_parallel_world_size,
maybe_init_distributed_environment_and_model_parallel,
model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.test.server.accuracy_config import (
DEFAULT_TEXT_ENCODER_VOCAB_SIZE,
I2V_TEXT_ENCODER_DIM,
TEXT_ENCODER_INPUT_SEED,
TEXT_ENCODER_TOKEN_LENGTH,
TEXT_ENCODER_TOKEN_MAX,
TEXT_ENCODER_TOKEN_MIN,
ComponentType,
get_threshold,
)
STAGED_1GPU_NATIVE_CASE_IDS = {
"flux_2_image_t2i",
"qwen_image_layered_i2i",
"flux_2_image_t2i_upscaling_4x",
"flux_2_ti2i",
"flux_2_t2i_customized_vae_path",
"flux_2_ti2i_multi_image_cache_dit",
}
# These case allowlists are accuracy-runner policy. They select the few 1-GPU
# cases that need sequential SGLang/reference execution to stay within memory
# limits during CI and local correctness runs.
STAGED_1GPU_TEXT_ENCODER_CASE_IDS = {
"flux_2_image_t2i",
"flux_2_image_t2i_upscaling_4x",
"mova_360p_1gpu",
"flux_2_ti2i",
"flux_2_t2i_customized_vae_path",
"flux_2_ti2i_multi_image_cache_dit",
}
SOURCE_PREFIXES = (
"module.",
"model.",
"transformer.",
"text_encoder.",
"image_encoder.",
"encoder.",
"decoder.",
"model.language_model.",
"model.visual.",
)
TARGET_PREFIXES = (
"module.",
"model.",
"transformer.",
"text_encoder.",
"image_encoder.",
"encoder.",
"decoder.",
)
@dataclass(frozen=True)
class ComponentSelection:
base_model_id: str
base_model_root: str
component_paths: Dict[str, str]
source_root: str
source_path: str
source_subfolder: str
@dataclass(frozen=True)
class ParameterShardContext:
world_size: int
rank: int
def seed_and_broadcast(seed: int, tensor: torch.Tensor) -> torch.Tensor:
"""Seed and broadcast tensor across ranks for determinism."""
torch.manual_seed(seed)
if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
torch.distributed.broadcast(tensor, src=0)
return tensor
def read_json_file(path: str) -> Dict[str, Any]:
if not os.path.exists(path):
return {}
with open(path) as f:
return json.load(f)
def has_component_files(path: str) -> bool:
if not os.path.isdir(path):
return False
if os.path.exists(os.path.join(path, "config.json")):
return True
for ext in (".safetensors", ".bin", ".pth"):
if any(name.endswith(ext) for name in os.listdir(path)):
return True
return False
def list_safetensor_files(path: str) -> List[str]:
if not os.path.isdir(path):
return []
return sorted(
os.path.join(path, name)
for name in os.listdir(path)
if name.endswith(".safetensors")
)
def is_text_encoder_config(path: str) -> bool:
cfg_path = os.path.join(path, "config.json")
if not os.path.exists(cfg_path):
return False
cfg = read_json_file(cfg_path)
if cfg.get("model_type") == "i2v" or cfg.get("dim") == I2V_TEXT_ENCODER_DIM:
return False
return True
def _resolve_component_subfolder(
model_index: Dict[str, Any], key: str
) -> Optional[str]:
entry = model_index.get(key)
if isinstance(entry, dict):
return entry.get("path") or entry.get("subfolder")
if isinstance(entry, str):
return entry
if entry is not None:
return key
return None
def resolve_component_path(
local_root: str, component: ComponentType, model_index_keys: Tuple[str, ...]
) -> Tuple[str, str]:
model_index_path = os.path.join(local_root, "model_index.json")
model_index = read_json_file(model_index_path)
if model_index:
for key in model_index_keys:
subfolder = _resolve_component_subfolder(model_index, key)
if not subfolder:
continue
candidate = os.path.join(local_root, subfolder)
if not has_component_files(candidate):
continue
if component == ComponentType.TEXT_ENCODER and not is_text_encoder_config(
candidate
):
continue
return candidate, subfolder
if has_component_files(local_root):
if component != ComponentType.TEXT_ENCODER or is_text_encoder_config(
local_root
):
return local_root, ""
raise FileNotFoundError(
f"Could not resolve {component.value} from model_index.json under {local_root}"
)
def extract_component_path_overrides(extra_args: List[str]) -> Dict[str, str]:
component_paths: Dict[str, str] = {}
index = 0
while index < len(extra_args):
arg = extra_args[index]
key_part = arg.split("=", 1)[0] if "=" in arg else arg
if key_part.startswith("--") and key_part.endswith("-path"):
component = key_part[2:-5].replace("-", "_")
if "=" in arg:
component_paths[component] = arg.split("=", 1)[1]
elif index + 1 < len(extra_args) and not extra_args[index + 1].startswith(
"-"
):
index += 1
component_paths[component] = extra_args[index]
index += 1
for component, path in component_paths.items():
component_paths[component] = os.path.expanduser(path)
return component_paths
def load_checkpoint_weights(
module: nn.Module, model_path: str
) -> tuple[list[str], list[str]]:
safetensors_files = list_safetensor_files(model_path)
assert safetensors_files, f"Found no safetensors files in {model_path}"
loaded_state: Dict[str, torch.Tensor] = {}
for safetensor_path in safetensors_files:
loaded_state.update(safetensors_load_file(safetensor_path))
module.load_state_dict(loaded_state, strict=False)
state_keys = set(module.state_dict().keys())
loaded_keys = set(loaded_state.keys())
missing_keys = sorted(state_keys - loaded_keys)
unexpected_keys = sorted(loaded_keys - state_keys)
return missing_keys, unexpected_keys
def select_component_source(
model_id: str,
extra_args: List[str],
component: ComponentType,
model_index_keys: Tuple[str, ...],
) -> ComponentSelection:
component_paths = extract_component_path_overrides(extra_args)
base_model_root = maybe_download_model(model_id)
search_keys = [component.value]
for key in model_index_keys:
if key not in search_keys:
search_keys.append(key)
source_root = base_model_root
component_key = component.value
for key in search_keys:
override_path = component_paths.get(key)
if override_path:
source_root = maybe_download_model(override_path)
component_key = key
break
ordered_keys = [component_key]
for key in search_keys:
if key not in ordered_keys:
ordered_keys.append(key)
source_path, source_subfolder = resolve_component_path(
source_root,
component,
tuple(ordered_keys),
)
return ComponentSelection(
base_model_id=model_id,
base_model_root=base_model_root,
component_paths=component_paths,
source_root=source_root,
source_path=source_path,
source_subfolder=source_subfolder,
)
def ensure_distributed_env_defaults() -> None:
if "WORLD_SIZE" in os.environ:
return
os.environ.update(
{
"MASTER_ADDR": os.getenv("MASTER_ADDR", "127.0.0.1"),
"MASTER_PORT": os.getenv("MASTER_PORT", "29505"),
"RANK": "0",
"LOCAL_RANK": "0",
"WORLD_SIZE": "1",
}
)
def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
tp_size = sgl_args.tp_size
sp_degree = sgl_args.sp_degree
ulysses_degree = sgl_args.ulysses_degree
ring_degree = sgl_args.ring_degree
dp_size = sgl_args.dp_size
enable_cfg_parallel = bool(sgl_args.enable_cfg_parallel)
if (
tp_size is None
or sp_degree is None
or ulysses_degree is None
or ring_degree is None
):
raise RuntimeError(
"ServerArgs must have tp_size, sp_degree, ulysses_degree, and ring_degree before init"
)
if not model_parallel_is_initialized() and torch.distributed.is_initialized():
# A prior case may have failed while distributed groups were only partially
# initialized. Clear any stale group objects before re-initializing.
destroy_model_parallel()
if model_parallel_is_initialized():
current_tp = get_tensor_model_parallel_world_size()
current_sp = get_sequence_parallel_world_size()
current_dp = get_data_parallel_world_size()
if current_tp == tp_size and current_sp == sp_degree and current_dp == dp_size:
return
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_model_parallel()
ensure_distributed_env_defaults()
maybe_init_distributed_environment_and_model_parallel(
tp_size=tp_size,
sp_size=sp_degree,
enable_cfg_parallel=enable_cfg_parallel,
ulysses_degree=ulysses_degree,
ring_degree=ring_degree,
dp_size=dp_size,
)
if torch.distributed.is_initialized():
torch.distributed.barrier()
def build_accuracy_server_args(
base_model_id: str,
base_model_root: str,
case: Any,
component: ComponentType,
num_gpus: int,
component_paths: Dict[str, str],
) -> ServerArgs:
cfg_parallel = bool(case.server_args.cfg_parallel)
kwargs = {
"model_path": base_model_root,
"model_id": base_model_id,
"num_gpus": num_gpus,
"trust_remote_code": True,
"component_paths": component_paths,
"enable_cfg_parallel": cfg_parallel,
}
if case.server_args.tp_size is not None:
kwargs["tp_size"] = case.server_args.tp_size
if case.server_args.ulysses_degree is not None:
kwargs["ulysses_degree"] = case.server_args.ulysses_degree
if case.server_args.ring_degree is not None:
kwargs["ring_degree"] = case.server_args.ring_degree
if component == ComponentType.TEXT_ENCODER:
kwargs["enable_cfg_parallel"] = False
sgl_args = ServerArgs.from_kwargs(**kwargs)
sgl_args.text_encoder_cpu_offload = False
sgl_args.dit_cpu_offload = False
sgl_args.vae_cpu_offload = False
sgl_args.image_encoder_cpu_offload = False
sgl_args.enable_cache_dit = case.server_args.enable_cache_dit
sgl_args.dit_layerwise_offload = case.server_args.dit_layerwise_offload
sgl_args.dit_offload_prefetch_size = case.server_args.dit_offload_prefetch_size
return sgl_args
def set_module_attr(module: nn.Module, name: str, value: Any) -> None:
"""Assign to a nested parameter/buffer path such as `blocks.0.attn.to_q.weight`."""
attrs = name.split(".")
parent = module
for attr in attrs[:-1]:
if hasattr(parent, attr):
parent = getattr(parent, attr)
elif isinstance(parent, (nn.ModuleList, nn.Sequential)):
parent = parent[int(attr)]
elif isinstance(parent, nn.ModuleDict):
parent = parent[attr]
else:
raise AttributeError(
f"Cannot resolve {name} on {module.__class__.__name__}"
)
setattr(parent, attrs[-1], value)
def materialize_module(
module: nn.Module, device: torch.device, dtype: torch.dtype
) -> None:
"""Materialize meta tensors and cast floating tensors onto one target device/dtype."""
for name, param in module.named_parameters():
if param.device.type == "meta":
new_data = torch.zeros(param.shape, device=device, dtype=dtype)
if hasattr(param, "device_mesh") and param.device_mesh is not None:
new_data = distribute_tensor(
new_data, param.device_mesh, param.placements
)
set_module_attr(
module, name, nn.Parameter(new_data, requires_grad=param.requires_grad)
)
elif torch.is_floating_point(param):
param.data = param.data.to(device=device, dtype=dtype)
for name, buf in module.named_buffers():
if buf.device.type == "meta":
new_buf = torch.zeros(buf.shape, device=device, dtype=buf.dtype)
if hasattr(buf, "device_mesh") and buf.device_mesh is not None:
new_buf = distribute_tensor(new_buf, buf.device_mesh, buf.placements)
set_module_attr(module, name, new_buf)
elif torch.is_floating_point(buf):
buf.data = buf.data.to(device=device, dtype=dtype)
def build_parameter_shard_contexts(
module: nn.Module,
) -> Dict[str, ParameterShardContext]:
"""Record TP shard world/rank for each parameter owned by a TP-aware submodule."""
shard_contexts: Dict[str, ParameterShardContext] = {}
for module_name, submodule in module.named_modules():
tp_group = getattr(submodule, "tp_group", None)
if tp_group is None:
continue
context = ParameterShardContext(
world_size=get_group_size(tp_group),
rank=get_group_rank(tp_group),
)
if context.world_size <= 1:
continue
for name, _ in submodule.named_parameters(recurse=False):
qualified_name = f"{module_name}.{name}" if module_name else name
shard_contexts[qualified_name] = context
for name, _ in submodule.named_buffers(recurse=False):
qualified_name = f"{module_name}.{name}" if module_name else name
shard_contexts[qualified_name] = context
return shard_contexts
def build_state_lookup(state: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
"""Index a source state dict under both original and prefix-stripped names."""
lookup: Dict[str, torch.Tensor] = {}
for key, val in state.items():
lookup[key] = val
for prefix in SOURCE_PREFIXES:
if key.startswith(prefix):
lookup[key[len(prefix) :]] = val
return lookup
def normalize_state_key(name: str) -> str:
"""Normalize common naming differences between source and target state dicts."""
return (
name.replace("_fsdp_wrapped_module.", "")
.replace("_orig_mod.", "")
.replace("gamma", "weight")
.replace("beta", "bias")
.replace("scale", "weight")
.replace("shift", "bias")
)
def fuse_qkv(lookup: Dict[str, torch.Tensor], name: str) -> Optional[torch.Tensor]:
if "qkv_proj" not in name:
return None
variants = ["q_proj", "q"]
for repl in variants:
q_name = name.replace("qkv_proj", repl)
k_name = q_name.replace(".q_proj", ".k_proj").replace(".q", ".k")
v_name = q_name.replace(".q_proj", ".v_proj").replace(".q", ".v")
if q_name in lookup and k_name in lookup and v_name in lookup:
return torch.cat([lookup[q_name], lookup[k_name], lookup[v_name]], dim=0)
return None
def fuse_gate_up_proj(
lookup: Dict[str, torch.Tensor], name: str
) -> Optional[torch.Tensor]:
if "gate_up_proj" not in name:
return None
for gate_token, up_token in (("gate_proj", "up_proj"), ("wi_0", "wi_1")):
gate_name = name.replace("gate_up_proj", gate_token)
up_name = name.replace("gate_up_proj", up_token)
if gate_name in lookup and up_name in lookup:
return torch.cat([lookup[gate_name], lookup[up_name]], dim=0)
return None
def generate_name_candidates(
name: str, reverse_mapping: Optional[Dict[str, Tuple[str, Any, Any]]]
) -> List[str]:
candidates: List[str] = []
clean = normalize_state_key(name)
for cand in (name, clean):
if cand not in candidates:
candidates.append(cand)
if reverse_mapping:
for key in (name, clean):
entry = reverse_mapping.get(key)
if entry and entry[0] not in candidates:
candidates.append(entry[0])
for prefix in TARGET_PREFIXES:
if clean.startswith(prefix):
stripped = clean[len(prefix) :]
if stripped and stripped not in candidates:
candidates.append(stripped)
parts = clean.split(".")
for i in range(1, len(parts)):
cand = ".".join(parts[i:])
if cand not in candidates:
candidates.append(cand)
return candidates
def copy_tensor(
dest: torch.Tensor,
src: torch.Tensor,
tp_world: int,
rank: int,
) -> bool:
if src.numel() == 0:
return False
src = src.to(device=dest.device, dtype=dest.dtype)
if hasattr(dest, "device_mesh") and dest.device_mesh is not None:
if src.numel() == dest.numel():
with torch.no_grad():
dt = distribute_tensor(
src.view(dest.shape), dest.device_mesh, dest.placements
)
dest.copy_(dt)
return True
if src.numel() == dest.numel():
with torch.no_grad():
dest.copy_(src.view(dest.shape))
return True
if tp_world > 1 and src.numel() == dest.numel() * tp_world:
if (
src.ndim == dest.ndim
and src.shape[0] == dest.shape[0] * tp_world
and src.shape[1:] == dest.shape[1:]
):
with torch.no_grad():
dest.copy_(src[rank * dest.shape[0] : (rank + 1) * dest.shape[0], ...])
return True
if (
src.ndim >= 2
and dest.ndim >= 2
and src.ndim == dest.ndim
and src.shape[0] == dest.shape[0]
and src.shape[1] == dest.shape[1] * tp_world
and src.shape[2:] == dest.shape[2:]
):
with torch.no_grad():
dest.copy_(src[:, rank * dest.shape[1] : (rank + 1) * dest.shape[1]])
return True
if src.ndim == 4 and dest.ndim == 5 and dest.numel() == src.numel() * dest.shape[2]:
with torch.no_grad():
dest.copy_(src.unsqueeze(2).repeat(1, 1, dest.shape[2], 1, 1))
return True
return False
def _config_to_dict(config: Any) -> Dict[str, Any]:
to_dict = getattr(config, "to_dict", None)
if not callable(to_dict):
return {}
config_dict = to_dict()
return config_dict if isinstance(config_dict, dict) else {}
def resolve_text_encoder_vocab_size(config: Any) -> int:
config_dict = _config_to_dict(config)
vocab_size = config_dict.get("vocab_size")
if isinstance(vocab_size, int) and vocab_size > 0:
return vocab_size
text_config = config_dict.get("text_config")
if isinstance(text_config, dict):
nested_vocab_size = text_config.get("vocab_size")
if isinstance(nested_vocab_size, int) and nested_vocab_size > 0:
return nested_vocab_size
return DEFAULT_TEXT_ENCODER_VOCAB_SIZE
def build_deterministic_text_encoder_inputs(
config: Any, device: str
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build one stable token batch that works across text-encoder implementations."""
vocab_size = resolve_text_encoder_vocab_size(config)
max_token_id = max(
TEXT_ENCODER_TOKEN_MIN + 1, min(vocab_size, TEXT_ENCODER_TOKEN_MAX)
)
torch.manual_seed(TEXT_ENCODER_INPUT_SEED)
input_ids = torch.randint(
TEXT_ENCODER_TOKEN_MIN,
max_token_id,
(1, TEXT_ENCODER_TOKEN_LENGTH),
device="cpu",
dtype=torch.long,
).to(device)
attention_mask = torch.ones_like(input_ids)
return input_ids, attention_mask
def resolve_text_encoder_forward_module(model: nn.Module) -> nn.Module:
get_encoder = getattr(model, "get_encoder", None)
return get_encoder() if callable(get_encoder) else model
def _module_device(module: nn.Module) -> torch.device:
param = next(module.parameters(), None)
if param is not None:
return param.device
buf = next(module.buffers(), None)
if buf is not None:
return buf.device
return torch.device("cpu")
def extract_output_tensor(output: Any) -> torch.Tensor:
"""Best-effort extraction of a tensor from model outputs."""
if isinstance(output, torch.Tensor):
return output
sample = getattr(output, "sample", None)
if sample is not None:
if isinstance(sample, (list, tuple)):
sample = sample[0]
if isinstance(sample, torch.Tensor):
return sample
last_hidden_state = getattr(output, "last_hidden_state", None)
if last_hidden_state is not None:
return last_hidden_state
hidden_states = getattr(output, "hidden_states", None)
if hidden_states:
return hidden_states[-1]
pooler_output = getattr(output, "pooler_output", None)
if pooler_output is not None:
return pooler_output
logits = getattr(output, "logits", None)
if logits is not None:
return logits
if (
isinstance(output, (list, tuple))
and output
and isinstance(output[0], torch.Tensor)
):
return output[0]
raise ValueError(f"Could not extract tensor from output of type {type(output)}")
def run_text_encoder_accuracy_pair(
sgl: nn.Module, ref: nn.Module
) -> tuple[torch.Tensor, torch.Tensor]:
input_ids, attention_mask = build_deterministic_text_encoder_inputs(
ref.config, "cpu"
)
return (
_run_single_text_encoder_forward(sgl, input_ids, attention_mask),
_run_single_text_encoder_forward(ref, input_ids, attention_mask),
)
def _should_stage_case(case: Any, component: ComponentType, num_gpus: int) -> bool:
if num_gpus == 2:
return True
if num_gpus != 1:
return False
if component == ComponentType.TEXT_ENCODER:
return case.id in STAGED_1GPU_TEXT_ENCODER_CASE_IDS
return case.id in STAGED_1GPU_NATIVE_CASE_IDS
def _run_single_text_encoder_forward(
model: nn.Module, input_ids: torch.Tensor, attention_mask: torch.Tensor
) -> torch.Tensor:
"""Run one encoder forward and normalize its output into a tensor."""
with torch.no_grad():
forward_model = resolve_text_encoder_forward_module(model)
model_device = _module_device(forward_model)
output = forward_model(
input_ids.to(device=model_device),
attention_mask=attention_mask.to(device=model_device),
output_hidden_states=True,
)
return extract_output_tensor(output)
def _run_staged_native_component_accuracy_case(
engine_cls: Any,
case: Any,
component: ComponentType,
library: str,
num_gpus: int,
) -> None:
from sglang.multimodal_gen.test.server.accuracy_hooks import (
resolve_component_native_profile,
)
sgl = None
ref = None
try:
sgl, ref, device = engine_cls.load_component_pair(
case,
component,
library,
num_gpus,
materialize_ref_on_device=False,
)
profile = resolve_component_native_profile(component)
inputs = profile.build_inputs(case, sgl, device, ref)
sgl_call = profile.prepare_sglang_call(sgl, inputs)
with torch.no_grad():
sgl_raw = engine_cls._execute_with_native_hook(sgl_call)
sgl_out = profile.normalize_sglang_output(sgl_raw)
sgl_out = engine_cls._apply_output_transforms(sgl_out, sgl_call).detach().cpu()
del sgl_call
del sgl_raw
del sgl
sgl = None
engine_cls.clear_memory()
ref = ref.to(device=device, dtype=torch.bfloat16).eval()
ref_call = profile.prepare_reference_call(ref, inputs)
with torch.no_grad():
ref_raw = engine_cls._execute_with_native_hook(ref_call)
ref_out = profile.normalize_reference_output(ref_raw)
ref_out = engine_cls._apply_output_transforms(ref_out, ref_call).detach().cpu()
del ref_call
del ref_raw
engine_cls.check_accuracy(
sgl_out,
ref_out,
f"{case.id}_{component.value}",
get_threshold(case.id, component),
)
finally:
if sgl is not None:
del sgl
if ref is not None:
del ref
engine_cls.reset_parallel_runtime()
engine_cls.clear_memory()
def _run_staged_text_encoder_accuracy_case(
engine_cls: Any, case: Any, num_gpus: int
) -> None:
sgl = None
ref = None
try:
sgl, ref, device = engine_cls.load_component_pair(
case,
ComponentType.TEXT_ENCODER,
"transformers",
num_gpus,
materialize_sgl_on_device=False,
materialize_ref_on_device=False,
)
input_ids, attention_mask = build_deterministic_text_encoder_inputs(
ref.config, "cpu"
)
sgl = sgl.to(device=device, dtype=torch.bfloat16).eval()
sgl_out = (
_run_single_text_encoder_forward(sgl, input_ids, attention_mask)
.detach()
.cpu()
)
del sgl
sgl = None
engine_cls.clear_memory()
ref = ref.to(device=device, dtype=torch.bfloat16).eval()
ref_out = (
_run_single_text_encoder_forward(ref, input_ids, attention_mask)
.detach()
.cpu()
)
engine_cls.check_accuracy(
sgl_out,
ref_out,
f"{case.id}_encoder",
get_threshold(case.id, ComponentType.TEXT_ENCODER),
)
finally:
if sgl is not None:
del sgl
if ref is not None:
del ref
engine_cls.reset_parallel_runtime()
engine_cls.clear_memory()
def run_native_component_accuracy_case(
engine_cls: Any,
case: Any,
component: ComponentType,
library: str,
num_gpus: int,
) -> None:
if _should_stage_case(case, component, num_gpus):
_run_staged_native_component_accuracy_case(
engine_cls, case, component, library, num_gpus
)
return
engine_cls.clear_memory()
sgl = None
ref = None
try:
sgl, ref, device = engine_cls.load_component_pair(
case, component, library, num_gpus
)
sgl_out, ref_out = engine_cls.run_component_pair_native(
case, component, sgl, ref, device
)
engine_cls.check_accuracy(
sgl_out,
ref_out,
f"{case.id}_{component.value}",
get_threshold(case.id, component),
)
finally:
if sgl is not None:
del sgl
if ref is not None:
del ref
engine_cls.reset_parallel_runtime()
engine_cls.clear_memory()
def run_text_encoder_accuracy_case(engine_cls: Any, case: Any, num_gpus: int) -> None:
if _should_stage_case(case, ComponentType.TEXT_ENCODER, num_gpus):
_run_staged_text_encoder_accuracy_case(engine_cls, case, num_gpus)
return
engine_cls.clear_memory()
sgl = None
ref = None
try:
sgl, ref, _device = engine_cls.load_component_pair(
case, ComponentType.TEXT_ENCODER, "transformers", num_gpus
)
sgl_out, ref_out = run_text_encoder_accuracy_pair(sgl, ref)
engine_cls.check_accuracy(
sgl_out,
ref_out,
f"{case.id}_encoder",
get_threshold(case.id, ComponentType.TEXT_ENCODER),
)
finally:
if sgl is not None:
del sgl
if ref is not None:
del ref
engine_cls.reset_parallel_runtime()
engine_cls.clear_memory()
python/sglang/multimodal_gen/test/server/component_accuracy.py
+534
View File
@@ -0,0 +1,534 @@
from __future__ import annotations
import gc
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple
import diffusers
import torch
import torch.nn as nn
from transformers import (
AutoConfig,
AutoModel,
AutoModelForCausalLM,
T5EncoderModel,
UMT5EncoderModel,
)
try:
from transformers import AutoModelForImageTextToText as AutoVisionTextModel
except ImportError:
try:
from transformers import AutoModelForVision2Seq as AutoVisionTextModel
except ImportError:
AutoVisionTextModel = None
import sglang.multimodal_gen.runtime.managers.forward_context as fc_mod
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
destroy_model_parallel,
get_local_torch_device,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
ComponentLoader,
)
from sglang.multimodal_gen.runtime.loader.utils import (
get_param_names_mapping,
hf_to_custom_state_dict,
)
from sglang.multimodal_gen.runtime.managers.forward_context import ForwardContext
from sglang.multimodal_gen.runtime.models.vaes.wanvae import AutoencoderKLWan
from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
get_diffusers_component_config,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.accuracy_config import (
DEFAULT_TIMESTEP,
ComponentType,
)
from sglang.multimodal_gen.test.server.accuracy_hooks import (
resolve_component_native_profile,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
build_accuracy_server_args,
build_parameter_shard_contexts,
build_state_lookup,
copy_tensor,
fuse_gate_up_proj,
fuse_qkv,
generate_name_candidates,
initialize_parallel_runtime,
load_checkpoint_weights,
materialize_module,
read_json_file,
resolve_text_encoder_forward_module,
select_component_source,
)
from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
logger = init_logger(__name__)
MIN_MATCH_RATIO = float(os.getenv("SGLANG_DIFFUSION_WEIGHT_MATCH_RATIO", "0.98"))
@dataclass(frozen=True)
class ComponentSpec:
model_index_keys: Tuple[str, ...]
reference_library: str
COMPONENT_SPECS: Dict[ComponentType, ComponentSpec] = {
ComponentType.VAE: ComponentSpec(
model_index_keys=(
"vae",
"vae_model",
"autoencoder",
"autoencoder_kl",
"video_vae",
"audio_vae",
),
reference_library="diffusers",
),
ComponentType.TRANSFORMER: ComponentSpec(
model_index_keys=("transformer", "unet", "dit", "video_dit", "audio_dit"),
reference_library="diffusers",
),
ComponentType.TEXT_ENCODER: ComponentSpec(
model_index_keys=(
"text_encoder",
"text_encoder_2",
"text_encoder_3",
"image_encoder",
),
reference_library="transformers",
),
}
# Component loading helpers
def _load_sglang_component(
comp_path: str,
sgl_args: ServerArgs,
component: ComponentType,
library: str,
text_encoder_cpu_offload: bool | None = None,
) -> nn.Module:
loader = ComponentLoader.for_component_type(component.value, library)
if component == ComponentType.TEXT_ENCODER:
component_model = loader.load_customized(
comp_path,
sgl_args,
component.value,
cpu_offload_flag=text_encoder_cpu_offload,
)
else:
component_model = loader.load_customized(comp_path, sgl_args, component.value)
if component_model is None:
raise RuntimeError(f"Failed to load customized {component.value}")
return component_model
def _load_wan_reference_vae(comp_path: str, pipeline_config) -> nn.Module:
vae_config = pipeline_config.vae_config
vae_config.update_model_arch(
get_diffusers_component_config(component_path=comp_path)
)
if hasattr(vae_config, "post_init"):
vae_config.post_init()
vae = AutoencoderKLWan(vae_config)
missing_keys, unexpected_keys = load_checkpoint_weights(vae, comp_path)
if missing_keys:
logger.warning("WAN VAE missing keys: %s", missing_keys)
if unexpected_keys:
logger.warning("WAN VAE unexpected keys: %s", unexpected_keys)
return vae
def _load_reference_component(
comp_path: str,
source_root: str,
component: ComponentType,
hub_id: str,
pipeline_config,
subfolder: str,
) -> nn.Module:
# WAN VAE does not have a clean generic diffusers auto-load path here, and we
# explicitly need checkpoint-loaded weights for reference-side transfer/parity.
if component == ComponentType.VAE and "wan" in hub_id.lower():
return _load_wan_reference_vae(comp_path, pipeline_config)
if component == ComponentType.VAE:
cfg = read_json_file(os.path.join(comp_path, "config.json"))
class_name = cfg.get("_class_name") if cfg else None
cls = getattr(diffusers, str(class_name), None) if class_name else None
if cls is None:
cls = diffusers.AutoencoderKL
return cls.from_pretrained(
source_root,
subfolder=subfolder,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
)
if component == ComponentType.TRANSFORMER:
cfg = read_json_file(os.path.join(comp_path, "config.json"))
class_name = cfg.get("_class_name") if cfg else None
load_kwargs: Dict[str, Any] = {
"torch_dtype": torch.bfloat16,
"trust_remote_code": True,
}
if cfg:
for k, out_k in [
("in_dim", "in_channels"),
("dim", "hidden_size"),
("num_heads", "num_attention_heads"),
("out_dim", "out_channels"),
]:
if k in cfg:
load_kwargs[out_k] = cfg[k]
candidates = [diffusers.AutoModel]
if class_name:
maybe_cls = getattr(diffusers, str(class_name), None)
if maybe_cls is not None:
candidates.insert(0, maybe_cls)
last_error: Optional[Exception] = None
for cls in candidates:
try:
return cls.from_pretrained(comp_path, **load_kwargs)
except Exception as exc:
last_error = exc
raise RuntimeError(f"Failed to load transformer from {comp_path}: {last_error}")
if component == ComponentType.TEXT_ENCODER:
config = AutoConfig.from_pretrained(comp_path, trust_remote_code=True)
kwargs = {
"torch_dtype": torch.bfloat16,
"trust_remote_code": True,
"config": config,
}
architectures = tuple(getattr(config, "architectures", ()) or ())
if (
"UMT5EncoderModel" in architectures
or getattr(config, "model_type", None) == "umt5"
):
class_order = [UMT5EncoderModel, AutoModel, AutoModelForCausalLM]
else:
class_order = [
AutoModel,
AutoModelForCausalLM,
UMT5EncoderModel,
T5EncoderModel,
]
if AutoVisionTextModel is not None:
class_order.append(AutoVisionTextModel)
last_error: Optional[Exception] = None
for cls in class_order:
try:
return cls.from_pretrained(comp_path, **kwargs)
except Exception as exc:
last_error = exc
raise RuntimeError(
f"Failed to load text encoder from {comp_path}: {last_error}"
)
raise RuntimeError(f"Unsupported component {component.value}")
# Public accuracy engine
class AccuracyEngine:
@staticmethod
def reset_parallel_runtime() -> None:
if torch.distributed.is_initialized():
torch.distributed.barrier()
if model_parallel_is_initialized():
destroy_model_parallel()
gc.collect()
if torch.cuda.is_available():
torch.cuda.synchronize()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
@staticmethod
def clear_memory() -> None:
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
@staticmethod
def _execute_with_native_hook(call) -> Any:
output: Any = None
def _hook(_: nn.Module, __: tuple[Any, ...], captured: Any) -> None:
nonlocal output
output = captured
handle = call.module.register_forward_hook(_hook)
try:
call.module(*call.args, **call.kwargs)
finally:
handle.remove()
assert output is not None
return output
@staticmethod
def _apply_output_transforms(tensor: torch.Tensor, call) -> torch.Tensor:
if call.negate_output:
return -tensor
return tensor
@staticmethod
def check_accuracy(
target: torch.Tensor, reference: torch.Tensor, name: str, threshold: float
) -> None:
full_tensor = getattr(target, "full_tensor", None)
if callable(full_tensor):
target = full_tensor()
t, r = target.detach().cpu().float(), reference.detach().cpu().float()
logger.info(
"[%s] Shape: SGL=%s, REF=%s | NaNs: SGL=%s, REF=%s",
name,
list(t.shape),
list(r.shape),
torch.isnan(t).sum(),
torch.isnan(r).sum(),
)
if t.shape != r.shape:
if t.ndim == 5 and t.shape[2] == 1:
t = t.squeeze(2)
if r.ndim == 5 and r.shape[2] == 1:
r = r.squeeze(2)
if t.shape != r.shape:
raise RuntimeError(
f"Accuracy shape mismatch for {name}: {list(t.shape)} vs {list(r.shape)}"
)
cos_sim = torch.nn.functional.cosine_similarity(
t.reshape(-1), r.reshape(-1), dim=0
).item()
rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
logger.info("[%s] Rank %s CosSim=%.6f", name, rank, cos_sim)
assert (
cos_sim > threshold
), f"Accuracy failure in {name}: CosSim {cos_sim:.4f} < {threshold}"
@staticmethod
def transfer_weights(
source: nn.Module,
target: nn.Module,
min_match_ratio: float = MIN_MATCH_RATIO,
target_device: Optional[torch.device] = None,
) -> None:
device = target_device or get_local_torch_device()
dtype = torch.bfloat16
materialize_module(target, device, dtype)
source_state = source.state_dict()
mapping = getattr(target, "param_names_mapping", None) or getattr(
getattr(target, "module", None), "param_names_mapping", None
)
if mapping:
source_state, _ = hf_to_custom_state_dict(
source_state, get_param_names_mapping(mapping)
)
lookup = build_state_lookup(source_state)
reverse_mapping = getattr(
target, "reverse_param_names_mapping", None
) or getattr(
getattr(target, "module", None), "reverse_param_names_mapping", None
)
tp_world = (
get_tensor_model_parallel_world_size()
if model_parallel_is_initialized()
else 1
)
rank = (
get_tensor_model_parallel_rank() if model_parallel_is_initialized() else 0
)
shard_contexts = build_parameter_shard_contexts(target)
matched = 0
total = 0
unmatched_details: List[str] = []
for name, tensor in target.named_parameters():
total += 1
src_tensor = None
for cand in generate_name_candidates(name, reverse_mapping):
if cand in lookup:
src_tensor = lookup[cand]
break
if src_tensor is None:
for cand in generate_name_candidates(name, reverse_mapping):
src_tensor = fuse_qkv(lookup, cand)
if src_tensor is not None:
break
if src_tensor is None:
for cand in generate_name_candidates(name, reverse_mapping):
src_tensor = fuse_gate_up_proj(lookup, cand)
if src_tensor is not None:
break
if src_tensor is None:
unmatched_details.append(f"{name}: no matching source tensor")
continue
shard_context = shard_contexts.get(name)
shard_world_size = (
shard_context.world_size if shard_context is not None else tp_world
)
shard_rank = shard_context.rank if shard_context is not None else rank
if copy_tensor(tensor, src_tensor, shard_world_size, shard_rank):
matched += 1
else:
unmatched_details.append(
f"{name}: source {list(src_tensor.shape)} -> target {list(tensor.shape)} unsupported for shard_world_size={shard_world_size}"
)
for name, tensor in target.named_buffers():
src_tensor = None
for cand in generate_name_candidates(name, reverse_mapping):
if cand in lookup:
src_tensor = lookup[cand]
break
if src_tensor is None:
continue
shard_context = shard_contexts.get(name)
shard_world_size = (
shard_context.world_size if shard_context is not None else tp_world
)
shard_rank = shard_context.rank if shard_context is not None else rank
copy_tensor(tensor, src_tensor, shard_world_size, shard_rank)
ratio = matched / max(total, 1)
logger.info(
"Weight transfer: %s/%s matched (%.2f%%).", matched, total, ratio * 100
)
if ratio < min_match_ratio:
if rank == 0 and unmatched_details:
logger.error(
"Unmatched parameter details:\n%s", "\n".join(unmatched_details)
)
raise RuntimeError(
f"Weight transfer matched {matched}/{total} ({ratio:.2%}); below threshold {min_match_ratio:.2%}."
)
@staticmethod
def run_component_pair_native(
case: DiffusionTestCase,
component: ComponentType,
sgl_model: nn.Module,
ref_model: nn.Module,
device: str,
) -> Tuple[torch.Tensor, torch.Tensor]:
if component == ComponentType.TEXT_ENCODER:
raise ValueError("Text encoder path is not migrated to native hooks yet")
profile = resolve_component_native_profile(component)
inputs = profile.build_inputs(case, sgl_model, device, ref_model)
sgl_call = profile.prepare_sglang_call(sgl_model, inputs)
ref_call = profile.prepare_reference_call(ref_model, inputs)
with torch.no_grad():
sgl_raw = AccuracyEngine._execute_with_native_hook(sgl_call)
ref_raw = AccuracyEngine._execute_with_native_hook(ref_call)
sgl_out = profile.normalize_sglang_output(sgl_raw)
ref_out = profile.normalize_reference_output(ref_raw)
sgl_out = AccuracyEngine._apply_output_transforms(sgl_out, sgl_call)
ref_out = AccuracyEngine._apply_output_transforms(ref_out, ref_call)
return sgl_out, ref_out
@staticmethod
def load_component_pair(
case: DiffusionTestCase,
component: ComponentType,
library: str,
num_gpus: int,
materialize_sgl_on_device: bool = True,
materialize_ref_on_device: bool = True,
) -> Tuple[nn.Module, nn.Module, str]:
spec = COMPONENT_SPECS[component]
if library != spec.reference_library:
logger.warning(
"Overriding library '%s' with '%s' for component '%s'.",
library,
spec.reference_library,
component.value,
)
library = spec.reference_library
hub_id = case.server_args.model_path
component_selection = select_component_source(
hub_id,
case.server_args.extras,
component,
spec.model_index_keys,
)
sgl_args = build_accuracy_server_args(
component_selection.base_model_id,
component_selection.base_model_root,
case,
component,
num_gpus,
component_selection.component_paths,
)
initialize_parallel_runtime(sgl_args)
set_global_server_args(sgl_args)
device = get_local_torch_device()
sgl_component = _load_sglang_component(
component_selection.source_path,
sgl_args,
component,
library,
text_encoder_cpu_offload=(
False
if component != ComponentType.TEXT_ENCODER or materialize_sgl_on_device
else True
),
)
if materialize_sgl_on_device:
sgl_component = sgl_component.to(device=device, dtype=torch.bfloat16)
ref_component = _load_reference_component(
component_selection.source_path,
component_selection.source_root,
component,
hub_id,
sgl_args.pipeline_config,
component_selection.source_subfolder,
)
if materialize_ref_on_device:
ref_component = ref_component.to(device=device, dtype=torch.bfloat16)
if component == ComponentType.TRANSFORMER and "wan" in hub_id.lower():
fc_mod._forward_context = ForwardContext(
current_timestep=0, attn_metadata=None
)
ref_for_transfer = ref_component
if (
component == ComponentType.TEXT_ENCODER
and getattr(ref_component, "shared", None) is None
):
ref_for_transfer = resolve_text_encoder_forward_module(ref_component)
AccuracyEngine.transfer_weights(
ref_for_transfer,
sgl_component,
target_device=(
device if materialize_sgl_on_device else torch.device("cpu")
),
)
if component != ComponentType.VAE:
if not hasattr(fc_mod._forward_context, "attn_metadata"):
fc_mod._forward_context = ForwardContext(
current_timestep=int(DEFAULT_TIMESTEP), attn_metadata=None
)
return sgl_component.eval(), ref_component.eval(), str(device)
python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py
+37
View File
@@ -0,0 +1,37 @@
import pytest
from sglang.multimodal_gen.test.server.accuracy_config import (
ComponentType,
get_skip_reason,
should_skip_component,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
run_native_component_accuracy_case,
run_text_encoder_accuracy_case,
)
from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
from sglang.multimodal_gen.test.server.testcase_configs import ONE_GPU_CASES_A
@pytest.mark.parametrize("case", ONE_GPU_CASES_A, ids=lambda x: x.id)
class TestAccuracy1GPU_A:
"""1-GPU Component Accuracy Suite (Set A)."""
def test_vae_accuracy(self, case):
if should_skip_component(case, ComponentType.VAE):
pytest.skip(get_skip_reason(case, ComponentType.VAE))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.VAE, "diffusers", 1
)
def test_transformer_accuracy(self, case):
if should_skip_component(case, ComponentType.TRANSFORMER):
pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 1
)
def test_encoder_accuracy(self, case):
if should_skip_component(case, ComponentType.TEXT_ENCODER):
pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
run_text_encoder_accuracy_case(AccuracyEngine, case, 1)
python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py
+37
View File
@@ -0,0 +1,37 @@
import pytest
from sglang.multimodal_gen.test.server.accuracy_config import (
ComponentType,
get_skip_reason,
should_skip_component,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
run_native_component_accuracy_case,
run_text_encoder_accuracy_case,
)
from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
from sglang.multimodal_gen.test.server.testcase_configs import ONE_GPU_CASES_B
@pytest.mark.parametrize("case", ONE_GPU_CASES_B, ids=lambda x: x.id)
class TestAccuracy1GPU_B:
"""1-GPU Component Accuracy Suite (Set B)."""
def test_vae_accuracy(self, case):
if should_skip_component(case, ComponentType.VAE):
pytest.skip(get_skip_reason(case, ComponentType.VAE))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.VAE, "diffusers", 1
)
def test_transformer_accuracy(self, case):
if should_skip_component(case, ComponentType.TRANSFORMER):
pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 1
)
def test_encoder_accuracy(self, case):
if should_skip_component(case, ComponentType.TEXT_ENCODER):
pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
run_text_encoder_accuracy_case(AccuracyEngine, case, 1)
python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py
+37
View File
@@ -0,0 +1,37 @@
import pytest
from sglang.multimodal_gen.test.server.accuracy_config import (
ComponentType,
get_skip_reason,
should_skip_component,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
run_native_component_accuracy_case,
run_text_encoder_accuracy_case,
)
from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
from sglang.multimodal_gen.test.server.testcase_configs import TWO_GPU_CASES_A
@pytest.mark.parametrize("case", TWO_GPU_CASES_A, ids=lambda x: x.id)
class TestAccuracy2GPU_A:
"""2-GPU Component Accuracy Suite (Set A)."""
def test_vae_accuracy(self, case):
if should_skip_component(case, ComponentType.VAE):
pytest.skip(get_skip_reason(case, ComponentType.VAE))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.VAE, "diffusers", 2
)
def test_transformer_accuracy(self, case):
if should_skip_component(case, ComponentType.TRANSFORMER):
pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 2
)
def test_encoder_accuracy(self, case):
if should_skip_component(case, ComponentType.TEXT_ENCODER):
pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
run_text_encoder_accuracy_case(AccuracyEngine, case, 2)
python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py
+37
View File
@@ -0,0 +1,37 @@
import pytest
from sglang.multimodal_gen.test.server.accuracy_config import (
ComponentType,
get_skip_reason,
should_skip_component,
)
from sglang.multimodal_gen.test.server.accuracy_utils import (
run_native_component_accuracy_case,
run_text_encoder_accuracy_case,
)
from sglang.multimodal_gen.test.server.component_accuracy import AccuracyEngine
from sglang.multimodal_gen.test.server.testcase_configs import TWO_GPU_CASES_B
@pytest.mark.parametrize("case", TWO_GPU_CASES_B, ids=lambda x: x.id)
class TestAccuracy2GPU_B:
"""2-GPU Component Accuracy Suite (Set B)."""
def test_vae_accuracy(self, case):
if should_skip_component(case, ComponentType.VAE):
pytest.skip(get_skip_reason(case, ComponentType.VAE))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.VAE, "diffusers", 2
)
def test_transformer_accuracy(self, case):
if should_skip_component(case, ComponentType.TRANSFORMER):
pytest.skip(get_skip_reason(case, ComponentType.TRANSFORMER))
run_native_component_accuracy_case(
AccuracyEngine, case, ComponentType.TRANSFORMER, "diffusers", 2
)
def test_encoder_accuracy(self, case):
if should_skip_component(case, ComponentType.TEXT_ENCODER):
pytest.skip(get_skip_reason(case, ComponentType.TEXT_ENCODER))
run_text_encoder_accuracy_case(AccuracyEngine, case, 2)