From d77a35a32b25d1460ac5c72f66f4d9fd554399a2 Mon Sep 17 00:00:00 2001 From: NATURE Date: Thu, 13 Aug 2026 11:23:47 +0800 Subject: [PATCH] [Bugfix] Scope Bagel FP8 config to diffusion stage (#6085) Signed-off-by: natureofnature Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com> --- .../diffusion/quantization/test_quantization_fp8.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 581a7498f..617083398 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -189,7 +189,7 @@ def _generate_single_stage_video( def _generate_bagel_image( - quantization_config: str | None = None, + diffusion_quantization_config: str | None = None, num_inference_steps: int = 15, ) -> tuple[Any, float]: """Generate an image with BAGEL (multi-stage: LLM + Diffusion). @@ -202,8 +202,8 @@ def _generate_bagel_image( "stage_configs_path": config_path, "stage_init_timeout": 300, } - if quantization_config: - omni_kwargs["quantization_config"] = quantization_config + if diffusion_quantization_config: + omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config model_name = omni_kwargs.pop("model") with OmniRunner(model_name, **omni_kwargs) as runner: @@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory(): @hardware_test(res={"cuda": "H100"}) def test_bagel_fp8_generates_image(): - """BAGEL with FP8 quantization_config generates a valid image. + """BAGEL with diffusion-stage FP8 generates a valid image. FP8 should only apply to the diffusion stage (Stage-1), not the LLM stage (Stage-0). We verify this by checking: 1. Image is generated successfully 2. LLM stage finish_reason is 'stop' (not 'length' from garbled output) """ - image, _ = _generate_bagel_image(quantization_config="fp8") + image, _ = _generate_bagel_image(diffusion_quantization_config="fp8") image.save("test_bagel_fp8.png") @hardware_test(res={"cuda": "H100"}) def test_bagel_bf16_generates_image(): """BAGEL without quantization generates a valid image (baseline).""" - image, _ = _generate_bagel_image(quantization_config=None) + image, _ = _generate_bagel_image(diffusion_quantization_config=None) image.save("test_bagel_bf16.png")