From 4f5b55e37947177e77b1ca754089edcd2b879150 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Wed, 1 Apr 2026 19:21:36 +0530 Subject: [PATCH] [diffusion][CI]: Add individual component accuracy CI for diffusion models (#18709) Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> --- .github/workflows/pr-test-multimodal-gen.yml | 102 ++ .github/workflows/pr-test.yml | 2 + .../component_loaders/text_encoder_loader.py | 43 +- .../runtime/loader/weight_utils.py | 7 +- .../test/server/accuracy_config.py | 386 ++++++++ .../test/server/accuracy_hooks.py | 579 ++++++++++++ .../test/server/accuracy_utils.py | 880 ++++++++++++++++++ .../test/server/component_accuracy.py | 534 +++++++++++ .../test/server/test_accuracy_1_gpu_a.py | 37 + .../test/server/test_accuracy_1_gpu_b.py | 37 + .../test/server/test_accuracy_2_gpu_a.py | 37 + .../test/server/test_accuracy_2_gpu_b.py | 37 + scripts/ci/utils/slash_command_handler.py | 2 + 13 files changed, 2670 insertions(+), 13 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/server/accuracy_config.py create mode 100644 python/sglang/multimodal_gen/test/server/accuracy_hooks.py create mode 100644 python/sglang/multimodal_gen/test/server/accuracy_utils.py create mode 100644 python/sglang/multimodal_gen/test/server/component_accuracy.py create mode 100644 python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py create mode 100644 python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py create mode 100644 python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py create mode 100644 python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index 5cdc72275..cfe207f8a 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -156,6 +156,108 @@ jobs: with: artifact-suffix: ${{ matrix.part }} + multimodal-gen-component-accuracy-1-gpu: + if: | + (inputs.target_stage == 'multimodal-gen-component-accuracy-1-gpu') || + ( + !inputs.target_stage && + ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && + inputs.multimodal_gen == 'true' + ) + runs-on: 1-gpu-h100 + timeout-minutes: 240 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} + + - uses: ./.github/actions/check-stage-health + + - uses: ./.github/actions/check-maintenance + + - name: Download artifacts + if: inputs.sgl_kernel == 'true' + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-python3.10-cuda12.9 + + - name: Install dependencies + timeout-minutes: 20 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion + + - name: Run diffusion component accuracy tests (1-GPU) + timeout-minutes: 240 + env: + RUNAI_STREAMER_MEMORY_LIMIT: 0 + run: | + cd python + if [ "${{ inputs.continue_on_error }}" = "true" ]; then + exit_code=0 + python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py || exit_code=$? + python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py || exit_code=$? + exit $exit_code + fi + python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py + python3 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py + + - uses: ./.github/actions/upload-cuda-coredumps + if: always() + + multimodal-gen-component-accuracy-2-gpu: + if: | + (inputs.target_stage == 'multimodal-gen-component-accuracy-2-gpu') || + ( + !inputs.target_stage && + ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && + inputs.multimodal_gen == 'true' + ) + runs-on: 2-gpu-h100 + timeout-minutes: 240 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} + + - uses: ./.github/actions/check-stage-health + + - uses: ./.github/actions/check-maintenance + + - name: Download artifacts + if: inputs.sgl_kernel == 'true' + uses: actions/download-artifact@v4 + with: + path: sgl-kernel/dist/ + merge-multiple: true + pattern: wheel-python3.10-cuda12.9 + + - name: Install dependencies + timeout-minutes: 20 + run: | + CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion + + - name: Run diffusion component accuracy tests (2-GPU) + timeout-minutes: 240 + env: + RUNAI_STREAMER_MEMORY_LIMIT: 0 + run: | + cd python + if [ "${{ inputs.continue_on_error }}" = "true" ]; then + exit_code=0 + torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py || exit_code=$? + torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py || exit_code=$? + exit $exit_code + fi + torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py + torchrun --nproc_per_node=2 -m pytest -s -v sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py + + - uses: ./.github/actions/upload-cuda-coredumps + if: always() + multimodal-gen-unit-test: if: | (inputs.target_stage == 'multimodal-gen-unit-test') || diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 9c759b74d..fbdfb7734 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -890,6 +890,8 @@ jobs: ( inputs.target_stage == 'multimodal-gen-test-1-gpu' || inputs.target_stage == 'multimodal-gen-test-2-gpu' || + inputs.target_stage == 'multimodal-gen-component-accuracy-1-gpu' || + inputs.target_stage == 'multimodal-gen-component-accuracy-2-gpu' || inputs.target_stage == 'multimodal-gen-unit-test' || ( !inputs.target_stage && diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 2e304daf0..5bd93acb0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index bdeb2ac45..64ec3b052 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -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 = [ diff --git a/python/sglang/multimodal_gen/test/server/accuracy_config.py b/python/sglang/multimodal_gen/test/server/accuracy_config.py new file mode 100644 index 000000000..d125e289d --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/accuracy_config.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/accuracy_hooks.py b/python/sglang/multimodal_gen/test/server/accuracy_hooks.py new file mode 100644 index 000000000..0ba46437d --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/accuracy_hooks.py @@ -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}") diff --git a/python/sglang/multimodal_gen/test/server/accuracy_utils.py b/python/sglang/multimodal_gen/test/server/accuracy_utils.py new file mode 100644 index 000000000..48ecace6f --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/accuracy_utils.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/server/component_accuracy.py b/python/sglang/multimodal_gen/test/server/component_accuracy.py new file mode 100644 index 000000000..ba59063c0 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/component_accuracy.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py b/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py new file mode 100644 index 000000000..409acab1e --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_a.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py b/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py new file mode 100644 index 000000000..0c0c7cb51 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/test_accuracy_1_gpu_b.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py b/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py new file mode 100644 index 000000000..c40929bf4 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_a.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py b/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py new file mode 100644 index 000000000..9b3cc1190 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/test_accuracy_2_gpu_b.py @@ -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) diff --git a/scripts/ci/utils/slash_command_handler.py b/scripts/ci/utils/slash_command_handler.py index 69d479921..73808fbfe 100644 --- a/scripts/ci/utils/slash_command_handler.py +++ b/scripts/ci/utils/slash_command_handler.py @@ -271,6 +271,8 @@ def handle_rerun_stage( "stage-c-test-deepep-8-gpu-h200", "multimodal-gen-test-1-gpu", "multimodal-gen-test-2-gpu", + "multimodal-gen-component-accuracy-1-gpu", + "multimodal-gen-component-accuracy-2-gpu", ] # Valid AMD stage names that support target_stage