diff --git a/docs/diffusion/quantization.md b/docs/diffusion/quantization.md index 04fa798b6..8714e5ed4 100644 --- a/docs/diffusion/quantization.md +++ b/docs/diffusion/quantization.md @@ -43,10 +43,47 @@ backend. | quant_family | checkpoint form | canonical CLI | supported models | extra dependency | platform / notes | |------------------|--------------------------------------------------------------------------------------------|------------------------------------------------------|--------------------------------------------------------------|---------------------------------------|-----------------------------------------------------------------------------------------------------------------------| | `fp8` | Quantized transformer component folder, or safetensors with `quantization_config` metadata | `--transformer-path` or `--transformer-weights-path` | ALL | None | Component-folder and single-file flows are both supported | +| `modelopt-fp8` | Converted ModelOpt FP8 transformer directory or repo with `config.json` | `--transformer-weights-path` | FLUX.2, Wan2.2 | None | Override config is read from the quantized transformer repo; `dit_cpu_offload` and `dit_layerwise_offload` are auto-disabled | | `nvfp4-modelopt` | NVFP4 safetensors file, sharded directory, or repo providing transformer weights | `--transformer-weights-path` | FLUX.2 | `comfy-kitchen` optional on Blackwell | Blackwell can use a best-performance kit when available; otherwise SGLang falls back to the generic ModelOpt FP4 path | | `nunchaku-svdq` | Pre-quantized Nunchaku transformer weights, usually named `svdq-{int4\|fp4}_r{rank}-...` | `--transformer-weights-path` | Model-specific support such as Qwen-Image, FLUX, and Z-Image | `nunchaku` | SGLang can infer precision and rank from the filename and supports both `int4` and `nvfp4` | | `msmodelslim` | Pre-quantized msmodelslim transformer weights | `--model-path` | Wan2.2 family | None | Currently only compatible with the Ascend NPU family and supports both `w8a8` and `w4a4` | +## ModelOpt FP8 + +### Usage Examples + +ModelOpt FP8 checkpoints should be converted into an SGLang-loadable transformer +override first, then loaded with the original base model: + +```bash +sglang generate \ + --model-path black-forest-labs/FLUX.2-dev \ + --transformer-weights-path BBuf/flux2-dev-modelopt-fp8-sglang-transformer \ + --prompt "A Logo With Bold Large Text: SGL Diffusion" \ + --save-output +``` + +```bash +sglang generate \ + --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --transformer-weights-path BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer \ + --prompt "a fox walking through neon rain" \ + --save-output +``` + +### Notes + +- `--transformer-weights-path` is the canonical flag for converted ModelOpt FP8 + diffusion checkpoints. +- If the override repo or local directory contains its own `config.json`, + SGLang reads the quantization config from that override instead of relying on + the base model config. +- `dit_cpu_offload` and `dit_layerwise_offload` are automatically disabled for + ModelOpt FP8 checkpoints because the runtime expects the transformed FP8 + weights to remain GPU-resident in their column-major layout. +- To build the converted checkpoint yourself from a ModelOpt diffusers export, + use `python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint`. + ## NVFP4 ### Usage Examples diff --git a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py new file mode 100644 index 000000000..573f80f8a --- /dev/null +++ b/python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py @@ -0,0 +1,142 @@ +import sys + +import pytest +import torch + +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp8Config, + ModelOptFp8LinearMethod, +) +from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8 +from sglang.srt.layers.quantization.fp8_utils import ( + cutlass_fp8_supported, + input_to_float8, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, suite="stage-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=80, suite="nightly-kernel-1-gpu", nightly=True) + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +MAX_FP8_DIFF = 5e-4 +TEST_CASES = [ + pytest.param(19, 150, 80, id="misaligned_projection_shape"), + pytest.param(512, 3072, 4096, id="flux2_added_kv_projection_shape"), +] + + +def _modelopt_fp8_supported() -> bool: + return torch.cuda.is_available() and cutlass_fp8_supported() + + +def _calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + if denominator == 0: + return 0.0 + sim = 2 * (x * y).sum() / denominator + return (1 - sim).item() + + +def _dequantize_fp8_input(qinput: torch.Tensor, x_scale: torch.Tensor) -> torch.Tensor: + return qinput.to(torch.float32) * x_scale.to(torch.float32) + + +def _dequantize_fp8_weight( + weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + if weight_scale.ndim == 0 or weight_scale.numel() == 1: + scale = weight_scale.to(torch.float32) + else: + scale = weight_scale.to(torch.float32).reshape(-1, 1).t() + return weight.to(torch.float32) * scale + + +def _build_layer( + weight_q: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: torch.Tensor, +) -> tuple[torch.nn.Module, ModelOptFp8LinearMethod]: + output_size, input_size = weight_q.shape + method = ModelOptFp8LinearMethod( + ModelOptFp8Config(is_checkpoint_fp8_serialized=True) + ) + layer = torch.nn.Module() + method.create_weights( + layer=layer, + input_size_per_partition=input_size, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=DTYPE, + weight_loader=lambda *args, **kwargs: None, + ) + layer = layer.to(device=DEVICE) + + layer.weight.data.copy_(weight_q) + layer.weight_scale.data.copy_(weight_scale.reshape_as(layer.weight_scale)) + layer.input_scale.data.copy_(input_scale.reshape_as(layer.input_scale)) + method.process_weights_after_loading(layer) + return layer, method + + +@pytest.mark.skipif( + not _modelopt_fp8_supported(), + reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support", +) +@pytest.mark.parametrize("m,n,k", TEST_CASES) +def test_checkpoint_processing(m: int, n: int, k: int) -> None: + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260410 + m + n + k) + + weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator) + weight_q, weight_scale = input_to_float8(weight) + input_scale = torch.tensor(1.0, device=DEVICE, dtype=torch.float32) + + layer, _ = _build_layer(weight_q, weight_scale, input_scale) + + assert tuple(layer.weight.shape) == (k, n) + assert tuple(layer.weight.stride()) == (1, k) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.input_scale.ndim == 0 + assert tuple(layer.weight_scale.shape) == (n, 1) + + expected_weight = weight_q.t().to(torch.float32) * weight_scale.to(torch.float32) + actual_weight = _dequantize_fp8_weight(layer.weight, layer.weight_scale) + torch.testing.assert_close(actual_weight, expected_weight, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif( + not _modelopt_fp8_supported(), + reason="Diffusion ModelOpt FP8 scaled mm correctness requires CUDA FP8 support", +) +@pytest.mark.parametrize("m,n,k", TEST_CASES) +def test_shape_correctness(m: int, n: int, k: int) -> None: + generator = torch.Generator(device=DEVICE) + generator.manual_seed(20260410 + m + n + k) + + x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator) + weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator) + weight_q, weight_scale = input_to_float8(weight) + _, input_scale = input_to_float8(x) + + layer, method = _build_layer(weight_q, weight_scale, input_scale) + + qinput, x_scale = static_quant_fp8( + x.contiguous(), + layer.input_scale, + repeat_scale=method.cutlass_fp8_supported, + ) + expected = torch.matmul( + _dequantize_fp8_input(qinput, x_scale), + _dequantize_fp8_weight(layer.weight, layer.weight_scale), + ) + + actual = method.apply(layer, x) + diff = _calc_diff(actual, expected.to(dtype=DTYPE)) + assert diff < MAX_FP8_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}" + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md new file mode 100644 index 000000000..3379baaae --- /dev/null +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md @@ -0,0 +1,308 @@ +--- +name: sglang-diffusion-modelopt-quant +description: Use when quantizing a diffusion DiT with NVIDIA ModelOpt and making the resulting FP8 or NVFP4 checkpoint loadable, verifiable, and benchmarkable in SGLang Diffusion. +--- + +# SGLang Diffusion ModelOpt Quant + +## Overview + +Use this skill when the task is to take a diffusion transformer through the full ModelOpt workflow: + +- quantize it with NVIDIA ModelOpt +- adapt the exported checkpoint to SGLang Diffusion +- verify that quality holds up +- benchmark whether the quantized checkpoint is actually faster + +This skill owns the ModelOpt-to-SGLang bridge. It is not a generic kernel-tuning skill. + +## Core Rules + +- Use ModelOpt's official `quantize.py` as the PTQ source of truth. +- Keep the workflow generic. Put model-specific fallback logic in small isolated branches, not in the main conversion path. +- Benchmark only when BF16 and quantized commands are identical except for the checkpoint override being tested. +- For diffusion FP8, pin `dit_cpu_offload=false` and `dit_layerwise_offload=false`. +- For multi-transformer pipelines, use per-component overrides when different components need different checkpoints. +- When a branch is missing the validated helper tools, refresh `python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py` and `python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py` instead of inventing one-off scripts elsewhere. +- After validating a new ModelOpt quant path, update both the FP8 and NVFP4 support tables in this skill before closing the task. + +## Read First + +Read these sources before changing code: + +- NVIDIA ModelOpt diffusers guide: `examples/diffusers/README.md` +- ModelOpt quantization entrypoint: `examples/diffusers/quantization/quantize.py` +- ModelOpt diffusers quant presets: `examples/diffusers/quantization/config.py` +- SGLang diffusion quant runtime: + - `python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py` + - `python/sglang/multimodal_gen/runtime/utils/quantization_utils.py` + - `python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py` +- Helper tools in this repo: + - [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) + - [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) + +If you are working on a new model family, inspect the transformer's config and tensor naming before changing the generic converter. + +## What SGLang Supports Here + +This repo now contains: + +- flat `quant_method=modelopt` plus `quant_algo=FP8/NVFP4` resolution +- diffusion-side ModelOpt FP8 linear loading +- diffusion-side NVFP4 loading from ModelOpt exports +- FLUX.2 packed-QKV detection that distinguishes packed NVFP4 checkpoints from standard diffusers exports +- automatic protection against incompatible FP8 offload modes +- FP8 export conversion: + [`python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py`](../../../tools/convert_modelopt_fp8_checkpoint.py) +- trajectory similarity validation: + [`python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py`](../../../tools/compare_diffusion_trajectory_similarity.py) + +## Documentation Maintenance + +- Keep two separate support tables in this skill: one for FP8 and one for NVFP4. +- After finishing a new quant support path, update both tables in every mirrored copy of this skill. +- Each row must record the validated scope, the Hugging Face repo or path for the quantized DiT weights, and the key caveats. +- If the quantized DiT weights are not published yet, write `unpublished` explicitly instead of leaving the field blank. + +## FP8 Supported Models + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.1-dev` | single-transformer override, deterministic latent/image comparison, H100 benchmark, torch-profiler trace | `BBuf/flux1-dev-modelopt-fp8-sglang-transformer` | SGLang converter keeps a validated BF16 fallback set for modulation and FF projection layers; use `--model-id FLUX.1-dev` for local mirrors | +| `black-forest-labs/FLUX.2-dev` | single-transformer override load and generation path | `BBuf/flux2-dev-modelopt-fp8-sglang-transformer` | published SGLang-ready transformer override | +| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | primary `transformer` quantized, `transformer_2` kept BF16 | `BBuf/wan22-t2v-a14b-modelopt-fp8-sglang-transformer` | do not describe this as dual-transformer full-model FP8 unless that path is validated separately | + +## NVFP4 Supported Models + +| Base Model | Validated Scope | HF DiT Weights | Notes | +| --- | --- | --- | --- | +| `black-forest-labs/FLUX.2-dev` | packed-QKV load path | `black-forest-labs/FLUX.2-dev-NVFP4` | validated packed export detection and runtime layout handling | + +## FP8 Vs NVFP4 + +FP8 and NVFP4 are not wired into SGLang in exactly the same way. + +FP8: + +- the validated ModelOpt diffusers FP8 export still needs an extra SGLang-side conversion step +- SGLang expects explicit `weight_scale` and `input_scale` +- the validated path also materializes SGLang-native `float8_e4m3fn` weights from `backbone.pt` + +NVFP4: + +- the official diffusers export often already contains packed FP4 weights, scale tensors, and enough safetensors metadata for SGLang to rebuild the quant config +- in that case SGLang mainly needs to detect the checkpoint family and rearrange tensors into the runtime layout +- this is why NVFP4 often does not need an extra offline conversion pass like FP8 does + +Important caveat: + +- "often" does not mean "always" +- the exact load path still depends on the checkpoint family, especially whether a model uses a packed-QKV layout + +## Generic Workflow + +### 1. Verify The BF16 Baseline First + +Before quantizing anything: + +- run the original BF16 model in SGLang +- fix the prompt, seed, size, step count, and GPU topology +- save the output and `perf.json` + +Do not start quantization work until the BF16 path is already healthy. + +### 2. Quantize With Official ModelOpt + +Use ModelOpt's official script. Generic template: + +```bash +python quantize.py \ + --model \ + --override-model-path \ + --model-dtype \ + --format \ + --batch-size 1 \ + --calib-size \ + --n-steps \ + --quantize-mha \ + --prompts-file \ + --quantized-torch-ckpt-save-path /ckpt \ + --hf-ckpt-dir /hf +``` + +For multi-transformer models: + +- quantize each backbone deliberately +- keep each output directory separate +- save both `backbone.pt` and the matching `hf/` export + +### 3. Convert FP8 Exports For SGLang + +FP8 requires an extra conversion step: + +```bash +PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ + --modelopt-hf-dir /hf \ + --modelopt-backbone-ckpt /ckpt/backbone.pt \ + --base-transformer-dir \ + --output-dir /sglang_transformer \ + --overwrite +``` + +What the converter does: + +- reads `weight_quantizer._amax` and `input_quantizer._amax` from `backbone.pt` +- writes `weight_scale` and `input_scale` +- materializes eligible FP8 weights as `float8_e4m3fn` +- preserves ModelOpt `ignore` layers as BF16 +- strips stale `_quantizer.*` tensors and fallback-layer scales that should not survive into the SGLang-native checkpoint + +For `FLUX.1-dev`, the validated fallback set currently keeps these modules in BF16: + +- `transformer_blocks.*.norm1.linear` +- `transformer_blocks.*.norm1_context.linear` +- `transformer_blocks.*.ff.net.0.proj` +- `transformer_blocks.*.ff.net.2` +- `transformer_blocks.*.ff_context.net.0.proj` +- `transformer_blocks.*.ff_context.net.2` +- `single_transformer_blocks.*.norm.linear` + +Use `--model-type flux1` to force that profile, or rely on `--model-type auto` when the export config identifies `FluxTransformer2DModel`. + +### 4. Load The Quantized Checkpoint In SGLang + +Single-transformer example: + +```bash +sglang generate \ + --model-path \ + --transformer-path \ + --prompt "" \ + --seed \ + --save-output +``` + +Multi-transformer example: + +```bash +sglang generate \ + --model-path \ + --transformer-path \ + --transformer-2-path \ + --prompt "" \ + --seed \ + --save-output +``` + +Guideline: + +- use the global `--transformer-path` only when the model effectively has one transformer override to apply +- use per-component overrides when different backbones need different checkpoints +- the preferred CLI form is `---path` +- config-expanded forms such as `--component_paths.transformer_2=...` also resolve to the same internal override map + +### 5. Validate Accuracy + +Use two levels of validation. + +Reduced deterministic validation: + +- keep prompt, seed, resolution, and step count fixed +- compare BF16 and quantized runs +- capture denoising trajectories +- inspect per-step latent cosine similarity plus MAE or RMSE +- compare final frames with image metrics such as PSNR or MAE + +Tool: + +```bash +PYTHONPATH=python python3 -m sglang.multimodal_gen.tools.compare_diffusion_trajectory_similarity \ + --model-path \ + --model-id \ + --prompt "" \ + --width \ + --height \ + --num-inference-steps \ + --guidance-scale \ + --seed \ + --candidate-transformer-path \ + --output-json +``` + +Use `--model-id FLUX.1-dev` when `--model-path` points to a local directory but the runtime still needs the native FLUX.1 model registration. + +Full-output validation: + +- run the same user-facing generation config in BF16 and quantized mode +- inspect the output visually +- only claim "quality preserved" for the exact scope you actually checked + +### 6. Benchmark Correctly + +Benchmark only when these match between BF16 and quantized: + +- prompt +- seed +- width and height +- frame count +- inference step count +- GPU count and topology +- offload flags +- compile settings +- profiler settings + +Only the quantized checkpoint path should differ. + +Interpretation rule: + +- the primary expected gain is in denoising +- text-encoding and VAE differences are secondary and should not be over-attributed unless they were quantized too + +### 7. Add Model-Specific Fallbacks Only When Needed + +If the generic FP8 path fails on a new model family: + +- inspect which modules are numerically sensitive or loader-incompatible +- keep fallback patterns small and explicit +- isolate them in the converter instead of scattering ad-hoc exceptions +- re-run deterministic trajectory checks after every fallback change + +Do not turn one validated model quirk into a generic rule unless another family also needs it. + +## FP8 Offload Constraint + +Current diffusion ModelOpt FP8 support requires: + +- `dit_cpu_offload=false` +- `dit_layerwise_offload=false` + +Reason: + +- the FP8 linear path depends on a CUTLASS-compatible weight layout after loading +- the offload and restore path does not preserve that layout +- in particular, layerwise offload can flatten and rebuild FP8 weights in a way that breaks the column-major requirement used by the FP8 GEMM path + +Runtime behavior: + +- SGLang currently force-disables these two flags when it detects `modelopt_fp8` +- benchmark commands should still pin them explicitly so the command line itself makes the comparison rule obvious + +## Claim Discipline + +When documenting results: + +- claim only scopes that were actually validated end to end +- do not collapse "single-transformer FP8 override" into "full-model FP8" +- do not call a practical deployment comparison a benchmark if BF16 and quantized commands used different offload behavior + +## Current Code Areas + +| File | Role | +| --- | --- | +| `runtime/layers/quantization/__init__.py` | registers diffusion quant methods | +| `runtime/layers/quantization/modelopt_quant.py` | ModelOpt FP8 and NVFP4 runtime loading | +| `runtime/utils/quantization_utils.py` | resolves flat ModelOpt configs and reconstructs NVFP4 config from metadata | +| `runtime/loader/transformer_load_utils.py` | guards incompatible FP8 offload modes | +| `runtime/models/dits/flux_2.py` | packed-QKV handling for the packed FLUX.2 NVFP4 family | +| `tools/convert_modelopt_fp8_checkpoint.py` | FP8 offline conversion into SGLang-native layout | +| `tools/compare_diffusion_trajectory_similarity.py` | reduced deterministic BF16-vs-quantized validation | diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 5e3eaf940..2cae2aa4b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -8,15 +8,17 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, + ModelOptFp8Config, ) from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig -QuantizationMethods = Literal["fp8", "modelopt_fp4", "modelslim"] +QuantizationMethods = Literal["fp8", "modelopt_fp8", "modelopt_fp4", "modelslim"] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) # The customized quantization methods which will be added to this dict. _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { + "modelopt_fp8": ModelOptFp8Config, "modelopt_fp4": ModelOptFp4Config, "modelslim": ModelSlimConfig, "fp8": Fp8Config, diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 6a89844a8..627edb93e 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging +import re from functools import lru_cache from typing import Any, Dict, List, Optional @@ -21,12 +22,20 @@ from sglang.multimodal_gen.runtime.models.parameter import ( ) from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.srt.layers.quantization.fp8_utils import ( + apply_fp8_linear, + cutlass_fp8_supported, +) from sglang.srt.layers.quantization.modelopt_quant import ( pad_nvfp4_activation_for_cutlass, pad_nvfp4_weight, slice_nvfp4_output, ) -from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.layers.quantization.utils import ( + convert_to_channelwise, + is_layer_skipped, + requantize_with_max_scale, +) from sglang.srt.layers.utils.common import copy_or_rebind_param from sglang.srt.utils.common import round_up @@ -82,15 +91,84 @@ class ModelOptQuantConfig(QuantizationConfig): def override_quantization_method(cls, hf_quant_config, user_quant) -> Optional[str]: if hf_quant_config is None: return None - quant_algo = hf_quant_config.get("quant_algo", "").upper() - if user_quant == "modelopt": - if not ("NVFP4" in quant_algo or "FP4" in quant_algo): - logger.warning( - f"Unsupported quant_algo '{quant_algo}' for 'modelopt'; defaulting to modelopt_fp4." - ) + + quant_algo = ( + hf_quant_config.get("quant_algo") + or hf_quant_config.get("quantization", {}).get("quant_algo") + or "" + ).upper() + if user_quant in {"modelopt", "modelopt_fp8"} and "FP8" in quant_algo: + return "modelopt_fp8" + if user_quant in {"modelopt", "modelopt_fp4"} and ( + "NVFP4" in quant_algo or "FP4" in quant_algo + ): return "modelopt_fp4" return None + def is_layer_excluded(self, prefix: str) -> bool: + for pattern in self.exclude_modules: + regex_str = re.escape(pattern).replace(r"\*", r".*") + if re.fullmatch(regex_str, prefix): + return True + return False + + +class ModelOptFp8Config(ModelOptQuantConfig): + """Config class for ModelOpt FP8 diffusion checkpoints.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + exclude_modules: Optional[List[str]] = None, + packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + ) -> None: + super().__init__(exclude_modules, packed_modules_mapping) + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + if is_checkpoint_fp8_serialized: + logger.warning( + "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change." + ) + + @classmethod + def get_name(cls) -> str: + return "modelopt_fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 89 + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config": + quant_method = config.get("quant_algo") + exclude_modules = config.get("ignore") + if quant_method is None: + try: + quantization_section = cls.get_from_keys(config, ["quantization"]) + quant_method = quantization_section.get("quant_algo") + exclude_modules = quantization_section.get("exclude_modules") + except ValueError as exc: + raise ValueError( + "Cannot find 'quant_algo' in the model's quantization config." + ) from exc + + if quant_method is None or "FP8" not in quant_method: + raise ValueError( + "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion." + ) + + return cls( + is_checkpoint_fp8_serialized=True, + exclude_modules=exclude_modules, + packed_modules_mapping=config.get("packed_modules_mapping"), + ) + + def get_quant_method(self, layer: torch.nn.Module, prefix: str): + return self._get_quant_method(layer, prefix, Linear=ModelOptFp8LinearMethod) + class ModelOptFp4Config(ModelOptQuantConfig): """Config class for NVFP4.""" @@ -101,6 +179,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): group_size: int = None, exclude_modules: List[str] = None, packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + checkpoint_uses_packed_qkv: bool = False, ) -> None: super().__init__(exclude_modules, packed_modules_mapping) self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized @@ -110,6 +189,7 @@ class ModelOptFp4Config(ModelOptQuantConfig): "format is experimental and subject to change." ) self.group_size = group_size + self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv @classmethod def get_name(cls) -> str: @@ -193,30 +273,100 @@ class ModelOptFp4Config(ModelOptQuantConfig): group_size=group_size, exclude_modules=exclude_modules, packed_modules_mapping=config.get("packed_modules_mapping"), + checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False), ) - def is_layer_excluded(self, prefix: str): - import regex as re - - fused_patterns = ["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"] - prefix_split = prefix.split(".") - for pattern in self.exclude_modules: - regex_str = pattern.replace(".", r"\.").replace("*", r".*") - pattern_split = pattern.split(".") - if re.fullmatch(regex_str, prefix): - return True - elif ( - pattern_split[-1] in fused_patterns - and pattern_split[-1] in prefix_split[-1] - ): - assert len(prefix_split) == 5 and len(pattern_split) == 5 - return True - return False - def get_quant_method(self, layer: torch.nn.Module, prefix: str): return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod) +class ModelOptFp8LinearMethod(LinearMethodBase): + """Linear method for ModelOpt static FP8 checkpoints.""" + + def __init__(self, quant_config: ModelOptFp8Config): + self.quant_config = quant_config + self.cutlass_fp8_supported = cutlass_fp8_supported() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del input_size, output_size + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_fp8_serialized + else params_dtype + ) + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + if self.quant_config.is_checkpoint_fp8_serialized: + for scale_name in ["weight_scale", "input_scale"]: + layer.register_parameter( + scale_name, + PerTensorScaleParameter( + data=torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + max_w_scale, quantized_weight = requantize_with_max_scale( + layer.weight, layer.weight_scale, layer.logical_widths + ) + # Preserve the parameter subclass metadata while rebinding to the + # transposed FP8 view expected by the runtime. + layer.weight.data = quantized_weight.t().detach() + layer.weight.requires_grad_(False) + if self.cutlass_fp8_supported: + max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths) + copy_or_rebind_param(layer, "weight_scale", max_w_scale) + copy_or_rebind_param(layer, "input_scale", layer.input_scale.max()) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_fp8_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + input_scale=layer.input_scale, + bias=bias, + cutlass_fp8_supported=self.cutlass_fp8_supported, + ) + + class ModelOptFp4LinearMethod(LinearMethodBase): """NVFP4 linear method using CUTLASS FP4 GEMM.""" diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 1d5f87921..7559f1fda 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -161,6 +161,59 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter): ) +class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter): + """Adapter for diffusion ModelOpt FP8 checkpoints.""" + + def __init__( + self, + *, + server_args: ServerArgs, + quant_config: Optional[QuantizationConfig], + ) -> None: + self.server_args = server_args + self.quant_config = quant_config + + @staticmethod + def _maybe_disable_incompatible_dit_offload_modes( + server_args: ServerArgs, + quant_config: Optional[QuantizationConfig], + ) -> None: + if quant_config is None: + return + + quant_name_getter = getattr(type(quant_config), "get_name", None) + quant_name = quant_name_getter() if callable(quant_name_getter) else None + if quant_name != "modelopt_fp8": + return + + disabled_args: list[str] = [] + + if server_args.dit_cpu_offload: + server_args.dit_cpu_offload = False + disabled_args.append("dit_cpu_offload") + + if server_args.dit_layerwise_offload: + server_args.dit_layerwise_offload = False + disabled_args.append("dit_layerwise_offload") + + if not disabled_args: + return + + logger.warning( + "ModelOpt FP8 diffusion checkpoints currently require the transformer " + "FP8 weights to stay GPU-resident in their column-major layout; " + "disabling %s for this run. Text encoder / VAE offload settings are " + "left unchanged.", + ", ".join(disabled_args), + ) + + def prepare(self) -> None: + _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( + server_args=self.server_args, + quant_config=self.quant_config, + ) + + def resolve_transformer_safetensors_to_load( server_args: ServerArgs, component_model_path: str ) -> list[str]: @@ -292,7 +345,11 @@ def _build_transformer_quant_adapters( cls_name=cls_name, server_args=server_args, quant_config=quant_config, - ) + ), + _ModelOptFp8OffloadAdapter( + server_args=server_args, + quant_config=quant_config, + ), ] if nunchaku_config is not None: adapters.append( @@ -305,6 +362,27 @@ def _build_transformer_quant_adapters( return adapters +def _resolve_quant_config_from_transformer_override( + transformer_weights_path: str, +) -> Optional[QuantizationConfig]: + """Resolve quant config from an override transformer repo or directory.""" + override_quantized_path = maybe_download_model(transformer_weights_path) + if not os.path.isdir(override_quantized_path): + return None + + override_config_path = os.path.join(override_quantized_path, "config.json") + if not os.path.isfile(override_config_path): + return None + + with open(override_config_path, encoding="utf-8") as f: + override_hf_config = json.load(f) + + return get_quant_config( + override_hf_config, + override_quantized_path, + ) + + def _resolve_quant_config( *, hf_config: dict, @@ -317,21 +395,29 @@ def _resolve_quant_config( priority: model config.json -> safetensors metadata -> format-specific fallback """ quant_config = get_quant_config(hf_config, component_model_path) - if quant_config is None and server_args.transformer_weights_path: - for safetensors_file in safetensors_list: - quant_config = get_quant_config_from_safetensors_metadata(safetensors_file) - if quant_config is not None: - return quant_config + if quant_config is not None or not server_args.transformer_weights_path: + return quant_config - param_names_mapping_dict = ( - server_args.pipeline_config.dit_config.arch_config.param_names_mapping - ) - quant_config = build_nvfp4_config_from_safetensors_list( - safetensors_list, param_names_mapping_dict - ) + quant_config = _resolve_quant_config_from_transformer_override( + server_args.transformer_weights_path + ) + if quant_config is not None: + return quant_config + + for safetensors_file in safetensors_list: + quant_config = get_quant_config_from_safetensors_metadata(safetensors_file) if quant_config is not None: return quant_config + param_names_mapping_dict = ( + server_args.pipeline_config.dit_config.arch_config.param_names_mapping + ) + quant_config = build_nvfp4_config_from_safetensors_list( + safetensors_list, param_names_mapping_dict + ) + if quant_config is not None: + return quant_config + return quant_config diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index 7ebe0d5e5..35ba8902e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -172,9 +172,12 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): self.added_kv_proj_dim = added_kv_proj_dim self.added_proj_bias = added_proj_bias - # Fuse Q/K/V into a single linear when using NVFP4: the checkpoint stores them - # packed as one tensor, so a fused layer avoids splitting during weight loading. - self.use_fused_qkv = isinstance(quant_config, ModelOptFp4Config) + # Some FLUX.2 NVFP4 checkpoints store Q/K/V packed as a single tensor, while + # ModelOpt's standard diffusers export keeps the original to_q/to_k/to_v layout. + # Only enable the fused loader path for the packed checkpoint family. + self.use_fused_qkv = isinstance(quant_config, ModelOptFp4Config) and getattr( + quant_config, "checkpoint_uses_packed_qkv", False + ) self.use_fused_added_qkv = self.use_fused_qkv if self.use_fused_qkv: @@ -484,7 +487,10 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): param.data.copy_(loaded_weight) self.to_out.weight_loader = _loader - self.to_out.weight.weight_loader = _loader + if hasattr(self.to_out.weight, "_weight_loader"): + self.to_out.weight._weight_loader = _loader + else: + self.to_out.weight.weight_loader = _loader def forward( self, @@ -827,7 +833,11 @@ class Flux2PosEmbed(nn.Module): repeat_interleave_real=False, dtype=( torch.float64 - if current_platform.is_float64_supported() + if ( + current_platform.is_float64_supported() + if hasattr(current_platform, "is_float64_supported") + else True + ) else torch.float32 ), ) diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 553b48de8..0d772453a 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -1,9 +1,11 @@ import glob import json import os +import re from pathlib import Path from typing import Dict, List, Optional +import torch from safetensors import safe_open from sglang.multimodal_gen.runtime.layers.quantization import ( @@ -15,6 +17,34 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def _resolve_quant_method_name(quant_cfg: dict) -> str: + quant_method = quant_cfg.get("quant_method") + if quant_method != "modelopt": + return quant_method + + quant_algo = ( + quant_cfg.get("quant_algo") + or quant_cfg.get("quantization", {}).get("quant_algo") + or "" + ).upper() + if quant_algo == "MIXED_PRECISION": + raise ValueError( + "ModelOpt mixed precision is not supported by the current SGLang diffusion runtime." + ) + if "FP8" in quant_algo: + return "modelopt_fp8" + if "FP4" in quant_algo or "NVFP4" in quant_algo: + return "modelopt_fp4" + raise ValueError(f"Unsupported ModelOpt quant_algo for diffusion: {quant_algo}") + + +def _load_quant_cls(quant_cfg: dict): + quant_method = _resolve_quant_method_name(quant_cfg) + if not quant_method: + raise ValueError("Missing quant_method in quantization config.") + return get_quantization_config(quant_method) + + def find_quant_modelslim_config(model_config, component_model_path): quant_config_file = Path(component_model_path, "quant_model_description.json") quant_cfg = None @@ -41,74 +71,70 @@ def get_quant_config( packed_modules_mapping: Dict[str, List[str]] = {}, remap_prefix: Dict[str, str] | None = None, ) -> QuantizationConfig: - quant_cfg = find_quant_modelslim_config(model_config, component_model_path) if quant_cfg is not None: - quant_cls = get_quantization_config(quant_cfg["quant_method"]) + quant_cls = _load_quant_cls(quant_cfg) return quant_cls.from_config(quant_cfg) - else: - if "quantization_config" not in model_config: - return None - quant_cls = get_quantization_config( - model_config["quantization_config"]["quant_method"] + + if "quantization_config" not in model_config: + return None + + hf_quant_config = model_config["quantization_config"] + if hf_quant_config is not None and not isinstance(hf_quant_config, dict): + hf_quant_config = hf_quant_config.to_dict() + quant_cls = _load_quant_cls(hf_quant_config) + + # GGUF doesn't have config file + if hf_quant_config["quant_method"] == "gguf": + return quant_cls.from_config({}) + + # some vision model may keep quantization_config in their text_config + hf_text_config = getattr(model_config, "text_config", None) + if hf_quant_config is None and hf_text_config is not None: + hf_quant_config = getattr(hf_text_config, "quantization_config", None) + if hf_quant_config is None: + # compressed-tensors uses a compressions_config + hf_quant_config = getattr(model_config, "compression_config", None) + if hf_quant_config is not None: + hf_quant_config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(hf_quant_config) + + model_name_or_path = model_config["model_path"] + is_local = os.path.isdir(model_name_or_path) + hf_folder = model_name_or_path + + possible_config_filenames = quant_cls.get_config_filenames() + + # If the quantization config is not found, use the default config. + if not possible_config_filenames: + return quant_cls() + + config_files = glob.glob(os.path.join(hf_folder, "*.json")) + + quant_config_files = [ + f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) + ] + if len(quant_config_files) == 0: + raise ValueError( + f"Cannot find the config file for {model_config['quantization_config']['quant_method']}" + ) + if len(quant_config_files) > 1: + raise ValueError( + f"Found multiple config files for {model_config['quantization_config']['quant_method']}: " + f"{quant_config_files}" ) - # GGUF doesn't have config file - if model_config["quantization_config"]["quant_method"] == "gguf": - return quant_cls.from_config({}) - - # Read the quantization config from the HF model config, if available. - hf_quant_config = model_config["quantization_config"] - # some vision model may keep quantization_config in their text_config - hf_text_config = getattr(model_config, "text_config", None) - if hf_quant_config is None and hf_text_config is not None: - hf_quant_config = getattr(hf_text_config, "quantization_config", None) - if hf_quant_config is None: - # compressed-tensors uses a compressions_config - hf_quant_config = getattr(model_config, "compression_config", None) - if hf_quant_config is not None: - hf_quant_config["packed_modules_mapping"] = packed_modules_mapping - return quant_cls.from_config(hf_quant_config) - # In case of bitsandbytes/QLoRA, get quant config from the adapter model. - else: - model_name_or_path = model_config["model_path"] - is_local = os.path.isdir(model_name_or_path) - hf_folder = model_name_or_path - - possible_config_filenames = quant_cls.get_config_filenames() - - # If the quantization config is not found, use the default config. - if not possible_config_filenames: - return quant_cls() - - config_files = glob.glob(os.path.join(hf_folder, "*.json")) - - quant_config_files = [ - f - for f in config_files - if any(f.endswith(x) for x in possible_config_filenames) - ] - if len(quant_config_files) == 0: - raise ValueError( - f"Cannot find the config file for {model_config['quantization_config']['quant_method']}" - ) - if len(quant_config_files) > 1: - raise ValueError( - f"Found multiple config files for {model_config['quantization_config']['quant_method']}: " - f"{quant_config_files}" - ) - - quant_config_file = quant_config_files[0] - with open(quant_config_file) as f: - config = json.load(f) - if remap_prefix is not None: - exclude_modules = [ - replace_prefix(key, remap_prefix) - for key in config["quantization"]["exclude_modules"] - ] - config["quantization"]["exclude_modules"] = exclude_modules - config["packed_modules_mapping"] = packed_modules_mapping - return quant_cls.from_config(config) + quant_config_file = quant_config_files[0] + with open(quant_config_file) as f: + config = json.load(f) + if remap_prefix is not None and "quantization" in config: + exclude_modules = [ + replace_prefix(key, remap_prefix) + for key in config["quantization"]["exclude_modules"] + ] + config["quantization"]["exclude_modules"] = exclude_modules + config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(config) def handle_fp8_metadata_format(quant_config_dict): @@ -132,11 +158,23 @@ def get_quant_config_from_safetensors_metadata( return None quant_config_str = metadata.get("_quantization_metadata") - if not quant_config_str: - return None - try: - quant_config_dict = json.loads(quant_config_str) - except Exception as _e: + quant_config_dict = None + if quant_config_str: + try: + quant_config_dict = json.loads(quant_config_str) + except Exception: + quant_config_dict = None + + if quant_config_dict is None: + quant_config_str = metadata.get("quantization_config") + if not quant_config_str: + return None + try: + quant_config_dict = json.loads(quant_config_str) + except Exception: + return None + + if not quant_config_dict: return None # handle diffusers fp8 safetensors metadata format @@ -152,7 +190,7 @@ def get_quant_config_from_safetensors_metadata( return None try: - quant_cls = get_quantization_config(quant_method) + quant_cls = _load_quant_cls(quant_config_dict) config = quant_cls.from_config(quant_config_dict) logger.debug(f"Get quantization config from safetensors file: {file_path}") return config @@ -179,12 +217,14 @@ def _build_nvfp4_config_from_safetensors_files( safetensors. Building the config from only the first matching file can incorrectly exclude layers that are quantized in a later shard. """ - import torch - group_size = None quantized_bfl_modules: set[str] = set() non_quantized_bfl_modules: set[str] = set() files_with_nvfp4_metadata: list[str] = [] + checkpoint_uses_packed_qkv = False + packed_qkv_pattern = re.compile( + r"^(double_blocks\.\d+\.(img|txt)_attn\.qkv|single_blocks\.\d+\.linear1)\." + ) for file_path in file_paths: metadata = get_metadata_from_safetensors_file(file_path) @@ -216,6 +256,8 @@ def _build_nvfp4_config_from_safetensors_files( with safe_open(file_path, framework="pt", device="cpu") as f: all_keys = set(f.keys()) + if any(packed_qkv_pattern.match(k) for k in all_keys): + checkpoint_uses_packed_qkv = True if group_size is None: for layer_name in file_quantized_modules: @@ -263,13 +305,19 @@ def _build_nvfp4_config_from_safetensors_files( try: quant_cls = get_quantization_config("modelopt_fp4") result = quant_cls.from_config( - {"quant_algo": "NVFP4", "group_size": group_size, "ignore": exclude_modules} + { + "quant_algo": "NVFP4", + "group_size": group_size, + "ignore": exclude_modules, + "checkpoint_uses_packed_qkv": checkpoint_uses_packed_qkv, + } ) logger.info( - "Built NVFP4 quant config from %d safetensors: group_size=%d, %d excluded modules", + "Built NVFP4 quant config from %d safetensors: group_size=%d, %d excluded modules, packed_qkv=%s", len(files_with_nvfp4_metadata), group_size, len(exclude_modules), + checkpoint_uses_packed_qkv, ) return result except Exception as e: diff --git a/python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py b/python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py new file mode 100644 index 000000000..9a6843b48 --- /dev/null +++ b/python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py @@ -0,0 +1,498 @@ +"""Compare diffusion BF16 and quantized runs via trajectory-latent similarity. + +This tool runs two SGLang diffusion variants with the same prompt and seed, +captures intermediate denoising latents via `return_trajectory_latents`, and +reports cosine / error metrics for each timestep plus final frame metrics. + +The intended use is quant validation on reduced deterministic smoke settings: +- same prompt / seed / resolution / step count for both variants +- BF16 reference on the base model +- FP8 candidate via `--candidate-transformer-path` and/or component overrides + +Example: + + python -m sglang.multimodal_gen.tools.compare_diffusion_trajectory_similarity \ + --model-path /path/to/model \ + --prompt "A futuristic cyberpunk city at night" \ + --width 512 --height 512 --num-inference-steps 8 --seed 42 \ + --text-encoder-cpu-offload \ + --candidate-transformer-path /tmp/modelopt_flux2_fp8/sglang_transformer \ + --output-json /tmp/flux2_similarity.json +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any, Sequence + +import imageio.v3 as iio +import numpy as np +import torch +import torch.nn.functional as F + + +def parse_component_overrides(entries: Sequence[str] | None) -> dict[str, str]: + overrides: dict[str, str] = {} + for entry in entries or []: + if "=" not in entry: + raise ValueError( + f"Invalid component override '{entry}'. Expected format component=path." + ) + component, path = entry.split("=", 1) + component = component.strip().replace("-", "_") + path = path.strip() + if not component or not path: + raise ValueError( + f"Invalid component override '{entry}'. Expected format component=path." + ) + overrides[component] = path + return overrides + + +def _cosine_similarity(flat_a: torch.Tensor, flat_b: torch.Tensor) -> float: + norm_a = torch.linalg.vector_norm(flat_a).item() + norm_b = torch.linalg.vector_norm(flat_b).item() + if norm_a == 0.0 and norm_b == 0.0: + return 1.0 + if norm_a == 0.0 or norm_b == 0.0: + return 0.0 + return float(F.cosine_similarity(flat_a, flat_b, dim=0).item()) + + +def compute_tensor_metrics(lhs: Any, rhs: Any) -> dict[str, float]: + lhs_tensor = torch.as_tensor(lhs).detach().cpu().float() + rhs_tensor = torch.as_tensor(rhs).detach().cpu().float() + if lhs_tensor.shape != rhs_tensor.shape: + raise ValueError( + f"Metric shape mismatch: {tuple(lhs_tensor.shape)} vs {tuple(rhs_tensor.shape)}" + ) + + diff = lhs_tensor - rhs_tensor + mse = float(diff.square().mean().item()) + rmse = float(math.sqrt(mse)) + mae = float(diff.abs().mean().item()) + max_abs = float(diff.abs().max().item()) + l2 = float(torch.linalg.vector_norm(diff).item()) + cosine = _cosine_similarity(lhs_tensor.reshape(-1), rhs_tensor.reshape(-1)) + return { + "cosine_similarity": cosine, + "mae": mae, + "mse": mse, + "rmse": rmse, + "max_abs": max_abs, + "l2": l2, + } + + +def compute_uint8_frame_metrics(lhs: Any, rhs: Any) -> dict[str, float]: + metrics = compute_tensor_metrics(lhs, rhs) + mse = metrics["mse"] + metrics["psnr_db"] = ( + float("inf") if mse == 0.0 else 20 * math.log10(255.0) - 10 * math.log10(mse) + ) + return metrics + + +def _normalize_step_index(step_index: int, num_steps: int) -> int: + if num_steps <= 0: + raise ValueError("num_steps must be positive.") + if step_index < 0: + step_index += num_steps + if step_index < 0 or step_index >= num_steps: + raise IndexError( + f"Requested step index {step_index} is outside the valid range [0, {num_steps})." + ) + return step_index + + +def _maybe_scalar(timestep: torch.Tensor | None, index: int) -> float | None: + if timestep is None: + return None + value = timestep[index] + if isinstance(value, torch.Tensor): + value = value.detach().cpu() + if value.numel() == 1: + return float(value.item()) + return float(value) + + +def summarize_trajectory_metrics( + reference_latents: Any, + candidate_latents: Any, + *, + reference_timesteps: Any = None, + candidate_timesteps: Any = None, + step_index: int = -1, +) -> dict[str, Any]: + ref = torch.as_tensor(reference_latents).detach().cpu().float() + cand = torch.as_tensor(candidate_latents).detach().cpu().float() + if ref.shape != cand.shape: + raise ValueError( + f"Trajectory shape mismatch: {tuple(ref.shape)} vs {tuple(cand.shape)}" + ) + if ref.ndim < 2: + raise ValueError( + f"Expected trajectory latents with an explicit timestep dimension, got {tuple(ref.shape)}" + ) + + num_steps = ref.shape[1] + selected_step = _normalize_step_index(step_index, num_steps) + ref_t = ( + torch.as_tensor(reference_timesteps).detach().cpu() + if reference_timesteps is not None + else None + ) + cand_t = ( + torch.as_tensor(candidate_timesteps).detach().cpu() + if candidate_timesteps is not None + else None + ) + + per_step: list[dict[str, Any]] = [] + for idx in range(num_steps): + metrics = compute_tensor_metrics(ref[:, idx], cand[:, idx]) + metrics["step_index"] = idx + metrics["reference_timestep"] = _maybe_scalar(ref_t, idx) + metrics["candidate_timestep"] = _maybe_scalar(cand_t, idx) + per_step.append(metrics) + + return { + "trajectory_shape": list(ref.shape), + "num_steps": num_steps, + "selected_step_index": selected_step, + "selected_step_metrics": per_step[selected_step], + "per_step_metrics": per_step, + } + + +def summarize_output_frame_metrics( + reference_frames: Sequence[Any], + candidate_frames: Sequence[Any], +) -> dict[str, Any]: + if len(reference_frames) != len(candidate_frames): + raise ValueError( + f"Output frame count mismatch: {len(reference_frames)} vs {len(candidate_frames)}" + ) + if not reference_frames: + raise ValueError("No output frames available for comparison.") + + ref_stack = np.stack([np.asarray(frame) for frame in reference_frames], axis=0) + cand_stack = np.stack([np.asarray(frame) for frame in candidate_frames], axis=0) + + frame0_metrics = compute_uint8_frame_metrics(ref_stack[0], cand_stack[0]) + mid_index = len(reference_frames) // 2 + mid_metrics = compute_uint8_frame_metrics( + ref_stack[mid_index], cand_stack[mid_index] + ) + all_metrics = compute_uint8_frame_metrics(ref_stack, cand_stack) + + return { + "num_frames": len(reference_frames), + "frame0_metrics": frame0_metrics, + "mid_frame_index": mid_index, + "mid_frame_metrics": mid_metrics, + "all_frames_metrics": all_metrics, + } + + +def extract_result_frames(result: Any) -> list[np.ndarray]: + if result.frames is not None: + return [np.asarray(frame) for frame in result.frames] + + sample = result.samples + if sample is None: + if result.output_file_path: + output_path = Path(result.output_file_path) + if not output_path.exists(): + raise ValueError( + "GenerationResult did not contain frames or samples, and its " + f"output_file_path does not exist: {output_path}" + ) + if output_path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}: + return [np.asarray(iio.imread(output_path))] + return [np.asarray(frame) for frame in iio.imiter(output_path)] + raise ValueError( + "GenerationResult did not contain frames, samples, or a readable output_file_path." + ) + + if isinstance(sample, torch.Tensor): + tensor = sample.detach().cpu().float() + if tensor.ndim == 3: + tensor = tensor.unsqueeze(1) + if tensor.ndim != 4: + raise ValueError( + f"Unsupported tensor sample shape for frame extraction: {tuple(tensor.shape)}" + ) + tensor = (tensor * 255).clamp(0, 255).to(torch.uint8) + frames = tensor.permute(1, 2, 3, 0).contiguous().numpy() + return [frame for frame in frames] + + array = np.asarray(sample) + if array.ndim == 2: + array = array[..., None] + if array.ndim == 3: + if array.shape[-1] in (1, 3, 4): + array = array[None, ...] + else: + array = array[..., None] + if array.ndim != 4: + raise ValueError( + f"Unsupported numpy sample shape for frame extraction: {tuple(array.shape)}" + ) + if array.dtype != np.uint8: + array = (np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8) + return [frame for frame in array] + + +def build_server_kwargs(args: argparse.Namespace, *, variant: str) -> dict[str, Any]: + component_paths = parse_component_overrides( + getattr(args, f"{variant}_component_path") or [] + ) + transformer_path = getattr(args, f"{variant}_transformer_path") + + kwargs: dict[str, Any] = { + "model_path": args.model_path, + "model_id": args.model_id, + "backend": args.backend, + "num_gpus": args.num_gpus, + "dit_cpu_offload": args.dit_cpu_offload, + "dit_layerwise_offload": args.dit_layerwise_offload, + "text_encoder_cpu_offload": args.text_encoder_cpu_offload, + "vae_cpu_offload": args.vae_cpu_offload, + "pin_cpu_memory": args.pin_cpu_memory, + "enable_cfg_parallel": args.enable_cfg_parallel, + "ulysses_degree": args.ulysses_degree, + } + if args.sp_degree is not None: + kwargs["sp_degree"] = args.sp_degree + if transformer_path is not None: + kwargs["transformer_weights_path"] = transformer_path + if component_paths: + kwargs["component_paths"] = component_paths + return kwargs + + +def build_sampling_kwargs( + args: argparse.Namespace, *, output_dir: str | None = None +) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "prompt": args.prompt, + "width": args.width, + "height": args.height, + "num_inference_steps": args.num_inference_steps, + "guidance_scale": args.guidance_scale, + "seed": args.seed, + "return_frames": True, + "return_trajectory_latents": True, + "return_trajectory_decoded": args.return_trajectory_decoded, + "save_output": output_dir is not None, + } + if output_dir is not None: + kwargs["output_path"] = output_dir + if args.num_frames is not None: + kwargs["num_frames"] = args.num_frames + if args.guidance_scale_2 is not None: + kwargs["guidance_scale_2"] = args.guidance_scale_2 + return kwargs + + +def run_variant( + *, + server_kwargs: dict[str, Any], + sampling_kwargs: dict[str, Any], +): + from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import ( + DiffGenerator, + ) + + with DiffGenerator.from_pretrained(local_mode=True, **server_kwargs) as generator: + result = generator.generate(sampling_params_kwargs=sampling_kwargs) + + if isinstance(result, list): + if len(result) != 1: + raise ValueError( + f"Expected a single generation result, got {len(result)} results." + ) + result = result[0] + if result is None: + raise RuntimeError("Generation returned no result.") + return result + + +def _to_jsonable(result: dict[str, Any]) -> dict[str, Any]: + return json.loads(json.dumps(result, allow_nan=True)) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-path", required=True) + parser.add_argument( + "--model-id", + help=( + "Optional model ID override passed to DiffGenerator.from_pretrained. " + "Use this when --model-path points to a local directory whose name " + "does not match a registered native SGLang model." + ), + ) + parser.add_argument("--backend", default="sglang") + parser.add_argument("--prompt", required=True) + parser.add_argument("--output-json", required=True) + parser.add_argument("--width", type=int, required=True) + parser.add_argument("--height", type=int, required=True) + parser.add_argument("--num-frames", type=int) + parser.add_argument("--num-inference-steps", type=int, required=True) + parser.add_argument("--guidance-scale", type=float, required=True) + parser.add_argument("--guidance-scale-2", type=float) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--num-gpus", type=int, default=1) + parser.add_argument("--ulysses-degree", type=int, default=1) + parser.add_argument("--sp-degree", type=int) + parser.add_argument("--trajectory-step-index", type=int, default=-1) + parser.add_argument("--reference-transformer-path") + parser.add_argument("--candidate-transformer-path") + parser.add_argument( + "--reference-component-path", + action="append", + default=[], + help="Repeatable component override in the form component=path.", + ) + parser.add_argument( + "--candidate-component-path", + action="append", + default=[], + help="Repeatable component override in the form component=path.", + ) + parser.add_argument("--save-output-dir") + parser.add_argument( + "--return-trajectory-decoded", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--enable-cfg-parallel", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--text-encoder-cpu-offload", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--vae-cpu-offload", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--dit-cpu-offload", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--dit-layerwise-offload", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument( + "--pin-cpu-memory", + action=argparse.BooleanOptionalAction, + default=False, + ) + args = parser.parse_args() + + output_json = Path(args.output_json).expanduser().resolve() + output_json.parent.mkdir(parents=True, exist_ok=True) + + save_root: Path | None = None + if args.save_output_dir: + save_root = Path(args.save_output_dir).expanduser().resolve() + save_root.mkdir(parents=True, exist_ok=True) + + ref_server_kwargs = build_server_kwargs(args, variant="reference") + cand_server_kwargs = build_server_kwargs(args, variant="candidate") + + ref_sampling_kwargs = build_sampling_kwargs( + args, + output_dir=str(save_root / "reference") if save_root else None, + ) + cand_sampling_kwargs = build_sampling_kwargs( + args, + output_dir=str(save_root / "candidate") if save_root else None, + ) + + reference = run_variant( + server_kwargs=ref_server_kwargs, + sampling_kwargs=ref_sampling_kwargs, + ) + candidate = run_variant( + server_kwargs=cand_server_kwargs, + sampling_kwargs=cand_sampling_kwargs, + ) + + result = { + "model_path": args.model_path, + "prompt": args.prompt, + "seed": args.seed, + "server_kwargs": { + "reference": ref_server_kwargs, + "candidate": cand_server_kwargs, + }, + "sampling_kwargs": { + "width": args.width, + "height": args.height, + "num_frames": args.num_frames, + "num_inference_steps": args.num_inference_steps, + "guidance_scale": args.guidance_scale, + "guidance_scale_2": args.guidance_scale_2, + }, + "reference_generation": { + "generation_time_s": reference.generation_time, + "peak_memory_mb": reference.peak_memory_mb, + "output_file_path": reference.output_file_path, + }, + "candidate_generation": { + "generation_time_s": candidate.generation_time, + "peak_memory_mb": candidate.peak_memory_mb, + "output_file_path": candidate.output_file_path, + }, + "trajectory_metrics": summarize_trajectory_metrics( + reference.trajectory_latents, + candidate.trajectory_latents, + reference_timesteps=reference.trajectory_timesteps, + candidate_timesteps=candidate.trajectory_timesteps, + step_index=args.trajectory_step_index, + ), + "output_metrics": summarize_output_frame_metrics( + extract_result_frames(reference), + extract_result_frames(candidate), + ), + } + + output_json.write_text( + json.dumps(_to_jsonable(result), indent=2, sort_keys=True), encoding="utf-8" + ) + + selected = result["trajectory_metrics"]["selected_step_metrics"] + frame0 = result["output_metrics"]["frame0_metrics"] + print( + json.dumps( + { + "output_json": str(output_json), + "trajectory_selected_step": result["trajectory_metrics"][ + "selected_step_index" + ], + "trajectory_cosine": selected["cosine_similarity"], + "trajectory_mae": selected["mae"], + "frame0_psnr_db": frame0["psnr_db"], + "frame0_mae": frame0["mae"], + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py b/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py new file mode 100644 index 000000000..89a731fba --- /dev/null +++ b/python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py @@ -0,0 +1,494 @@ +"""Convert a ModelOpt diffusion FP8 export into an SGLang-loadable checkpoint. + +The core conversion path is model-agnostic: +- read the ModelOpt diffusers transformer export +- rebuild per-layer `weight_scale` / `input_scale` tensors from `backbone.pt` +- materialize SGLang-native `float8_e4m3fn` weights +- preserve ModelOpt `ignore` layers in their original dtype + +Some models still benefit from a small validated BF16 fallback set. Those +fallback profiles are intentionally isolated so the generic FP8 conversion path +remains reusable across future diffusion backbones. + +Example: + + python -m sglang.multimodal_gen.tools.convert_modelopt_fp8_checkpoint \ + --modelopt-hf-dir /tmp/modelopt_flux2_fp8/hf \ + --modelopt-backbone-ckpt /tmp/modelopt_flux2_fp8/ckpt/backbone.pt \ + --base-transformer-dir /path/to/FLUX.2-dev/transformer \ + --output-dir /tmp/modelopt_flux2_fp8/sglang_transformer +""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import re +import shutil +from collections import defaultdict +from pathlib import Path +from typing import Iterable, Mapping, Sequence + +import torch +from safetensors import safe_open +from safetensors.torch import load_file, save_file + +INDEX_FILENAMES = [ + "model.safetensors.index.json", + "diffusion_pytorch_model.safetensors.index.json", +] +FP8_E4M3_MAXBOUND = 448.0 +DEFAULT_FLUX2_KEEP_BF16_PATTERNS = [ + r"^time_guidance_embed\.(timestep_embedder|guidance_embedder)\.linear_[12]$", + r"^double_stream_modulation_(img|txt)\.linear$", + r"^single_stream_modulation\.linear$", + r"^x_embedder$", + r"^context_embedder$", + r"^norm_out\.linear$", +] +DEFAULT_FLUX1_KEEP_BF16_PATTERNS = [ + r"^transformer_blocks\.\d+\.norm1\.linear$", + r"^transformer_blocks\.\d+\.norm1_context\.linear$", + r"^transformer_blocks\.\d+\.ff\.net\.0\.proj$", + r"^transformer_blocks\.\d+\.ff\.net\.2$", + r"^transformer_blocks\.\d+\.ff_context\.net\.0\.proj$", + r"^transformer_blocks\.\d+\.ff_context\.net\.2$", + r"^single_transformer_blocks\.\d+\.norm\.linear$", +] + + +def _resolve_transformer_dir(path: str) -> str: + candidate = Path(path).expanduser().resolve() + if (candidate / "config.json").is_file(): + return str(candidate) + transformer_dir = candidate / "transformer" + if (transformer_dir / "config.json").is_file(): + return str(transformer_dir) + raise FileNotFoundError(f"Could not resolve a transformer directory from: {path}") + + +def _resolve_backbone_ckpt(path: str) -> str: + candidate = Path(path).expanduser().resolve() + if candidate.is_file(): + return str(candidate) + backbone_path = candidate / "backbone.pt" + if backbone_path.is_file(): + return str(backbone_path) + raise FileNotFoundError(f"Could not resolve backbone.pt from: {path}") + + +def _find_index_file(model_dir: str) -> str | None: + for filename in INDEX_FILENAMES: + candidate = os.path.join(model_dir, filename) + if os.path.isfile(candidate): + return filename + + matches = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors.index.json") + ) + return matches[0] if matches else None + + +def _load_weight_map(model_dir: str) -> tuple[dict[str, str], str | None]: + index_filename = _find_index_file(model_dir) + if index_filename is not None: + with open(os.path.join(model_dir, index_filename), encoding="utf-8") as f: + index_data = json.load(f) + return dict(index_data["weight_map"]), index_filename + + safetensors_files = sorted( + filename + for filename in os.listdir(model_dir) + if filename.endswith(".safetensors") + ) + if len(safetensors_files) != 1: + raise ValueError( + f"Expected an index file or a single safetensors shard in {model_dir}, " + f"found {len(safetensors_files)} shard(s)." + ) + + shard_name = safetensors_files[0] + with safe_open( + os.path.join(model_dir, shard_name), framework="pt", device="cpu" + ) as f: + weight_map = {key: shard_name for key in f.keys()} + index_filename = f"{Path(shard_name).stem}.safetensors.index.json" + return weight_map, index_filename + + +def _load_config(model_dir: str) -> dict: + config_path = os.path.join(model_dir, "config.json") + with open(config_path, encoding="utf-8") as f: + return json.load(f) + + +def get_default_keep_bf16_patterns( + *, model_type: str, class_name: str | None +) -> list[str]: + if model_type == "flux1": + return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS) + if model_type == "flux2": + return list(DEFAULT_FLUX2_KEEP_BF16_PATTERNS) + if model_type == "none": + return [] + if class_name == "FluxTransformer2DModel": + return list(DEFAULT_FLUX1_KEEP_BF16_PATTERNS) + if class_name == "Flux2Transformer2DModel": + return list(DEFAULT_FLUX2_KEEP_BF16_PATTERNS) + return [] + + +def should_keep_bf16( + weight_name: str, + keep_bf16_patterns: Sequence[str], +) -> bool: + if not keep_bf16_patterns: + return False + + module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name + return any(re.search(pattern, module_name) for pattern in keep_bf16_patterns) + + +def is_ignored_by_modelopt( + weight_name: str, + ignore_patterns: Sequence[str], +) -> bool: + if not ignore_patterns: + return False + + module_name = weight_name[:-7] if weight_name.endswith(".weight") else weight_name + for pattern in ignore_patterns: + regex_str = pattern.replace(".", r"\.").replace("*", r".*") + if re.fullmatch(regex_str, module_name): + return True + return False + + +def build_fp8_scale_map( + model_state_dict: Mapping[str, torch.Tensor], + *, + maxbound: float = FP8_E4M3_MAXBOUND, +) -> dict[str, dict[str, torch.Tensor]]: + scale_map: dict[str, dict[str, torch.Tensor]] = {} + for key, value in model_state_dict.items(): + if key.endswith(".weight_quantizer._amax"): + layer_name = key[: -len(".weight_quantizer._amax")] + scale_map.setdefault(f"{layer_name}.weight", {})["weight_scale"] = ( + value.detach().to(torch.float32).reshape(1).cpu() / maxbound + ) + elif key.endswith(".input_quantizer._amax"): + layer_name = key[: -len(".input_quantizer._amax")] + scale_map.setdefault(f"{layer_name}.weight", {})["input_scale"] = ( + value.detach().to(torch.float32).reshape(1).cpu() / maxbound + ) + + return { + weight_name: scale_tensors + for weight_name, scale_tensors in scale_map.items() + if {"weight_scale", "input_scale"} <= set(scale_tensors) + } + + +def quantize_fp8_weight( + weight: torch.Tensor, + weight_scale: torch.Tensor, +) -> torch.Tensor: + if weight.dtype == torch.float8_e4m3fn: + return weight.contiguous() + + scale = weight_scale.to(weight.device, dtype=torch.float32) + if scale.numel() != 1: + raise ValueError( + "Only per-tensor FP8 scales are supported for diffusion checkpoints, " + f"got shape {tuple(scale.shape)}." + ) + + quantized = (weight.to(torch.float32) / scale.reshape(1)).to(torch.float8_e4m3fn) + return quantized.cpu().contiguous() + + +def _copy_non_shard_files(source_dir: str, output_dir: str) -> None: + ignored = set(INDEX_FILENAMES) + for entry in os.listdir(source_dir): + if entry.endswith(".safetensors") or entry in ignored: + continue + source_path = os.path.join(source_dir, entry) + output_path = os.path.join(output_dir, entry) + if os.path.isdir(source_path): + shutil.copytree(source_path, output_path, dirs_exist_ok=True) + else: + shutil.copy2(source_path, output_path) + + +def _load_selected_tensors( + model_dir: str, + weight_map: Mapping[str, str], + tensor_names: Iterable[str], +) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + names_by_file: dict[str, list[str]] = defaultdict(list) + for name in tensor_names: + names_by_file[weight_map[name]].append(name) + + for filename, names in names_by_file.items(): + shard_path = os.path.join(model_dir, filename) + with safe_open(shard_path, framework="pt", device="cpu") as f: + for name in names: + tensors[name] = f.get_tensor(name).contiguous() + return tensors + + +def convert_modelopt_fp8_checkpoint( + *, + modelopt_hf_dir: str, + modelopt_backbone_ckpt: str, + output_dir: str, + base_transformer_dir: str | None = None, + model_type: str = "auto", + keep_bf16_patterns: Sequence[str] | None = None, + maxbound: float = FP8_E4M3_MAXBOUND, + overwrite: bool = False, +) -> dict[str, int]: + source_dir = _resolve_transformer_dir(modelopt_hf_dir) + backbone_ckpt_path = _resolve_backbone_ckpt(modelopt_backbone_ckpt) + base_dir = ( + _resolve_transformer_dir(base_transformer_dir) if base_transformer_dir else None + ) + + config = _load_config(source_dir) + quant_config = config.get("quantization_config") + if not isinstance(quant_config, dict): + raise ValueError( + "Expected a flat quantization_config dict in the ModelOpt export." + ) + if ( + quant_config.get("quant_method") != "modelopt" + or "FP8" not in str(quant_config.get("quant_algo", "")).upper() + ): + raise ValueError( + "This tool only supports ModelOpt diffusers FP8 exports " + "(quant_method=modelopt, quant_algo=FP8)." + ) + + class_name = config.get("_class_name") + ignore_patterns = list(quant_config.get("ignore", []) or []) + patterns = list( + get_default_keep_bf16_patterns(model_type=model_type, class_name=class_name) + ) + if keep_bf16_patterns: + patterns.extend(keep_bf16_patterns) + if patterns and base_dir is None: + raise ValueError( + "BF16 fallback patterns are enabled, but --base-transformer-dir was not provided." + ) + + output_path = Path(output_dir).expanduser().resolve() + if output_path.exists(): + if not overwrite: + raise FileExistsError( + f"Output directory already exists: {output_path}. " + "Use --overwrite to replace it." + ) + shutil.rmtree(output_path) + output_path.mkdir(parents=True, exist_ok=True) + + _copy_non_shard_files(source_dir, str(output_path)) + + source_weight_map, index_filename = _load_weight_map(source_dir) + base_weight_map: dict[str, str] = {} + if base_dir is not None: + base_weight_map, _ = _load_weight_map(base_dir) + + backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[ + "model_state_dict" + ] + fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound) + serialized_quant_config = json.dumps(quant_config, sort_keys=True) + + fallback_weight_names = sorted( + weight_name + for weight_name in source_weight_map + if weight_name.endswith(".weight") and should_keep_bf16(weight_name, patterns) + ) + fallback_tensors = ( + _load_selected_tensors(base_dir, base_weight_map, fallback_weight_names) + if fallback_weight_names + else {} + ) + fallback_scale_names = { + scale_name + for weight_name in fallback_weight_names + for scale_name in ( + weight_name[:-7] + ".weight_scale", + weight_name[:-7] + ".input_scale", + ) + } + + weights_by_file: dict[str, list[str]] = defaultdict(list) + for weight_name, filename in source_weight_map.items(): + weights_by_file[filename].append(weight_name) + + updated_weight_map: dict[str, str] = {} + total_size = 0 + added_scale_count = 0 + preserved_ignored_weight_count = 0 + + for filename, names in sorted(weights_by_file.items()): + shard_path = os.path.join(source_dir, filename) + shard_tensors = load_file(shard_path, device="cpu") + + with safe_open(shard_path, framework="pt", device="cpu") as f: + metadata = dict(f.metadata() or {}) + + metadata.setdefault("format", "pt") + metadata["quantization_config"] = serialized_quant_config + metadata["_quantization_metadata"] = serialized_quant_config + + for name in list(shard_tensors.keys()): + if "_quantizer." in name: + del shard_tensors[name] + continue + if name in fallback_scale_names: + del shard_tensors[name] + continue + if name.endswith(".weight") and is_ignored_by_modelopt( + name, ignore_patterns + ): + preserved_ignored_weight_count += 1 + continue + if name in fallback_tensors: + shard_tensors[name] = fallback_tensors[name] + if ( + name.endswith(".weight") + and name in fp8_scale_map + and name not in fallback_tensors + ): + scale_tensors = fp8_scale_map[name] + shard_tensors[name] = quantize_fp8_weight( + shard_tensors[name], scale_tensors["weight_scale"] + ) + shard_tensors[name[:-7] + ".weight_scale"] = scale_tensors[ + "weight_scale" + ] + shard_tensors[name[:-7] + ".input_scale"] = scale_tensors["input_scale"] + added_scale_count += 2 + + save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata) + + for name, tensor in shard_tensors.items(): + updated_weight_map[name] = filename + total_size += tensor.element_size() * tensor.numel() + + del shard_tensors + gc.collect() + + with open(output_path / index_filename, "w", encoding="utf-8") as f: + json.dump( + { + "metadata": {"total_size": total_size}, + "weight_map": updated_weight_map, + }, + f, + indent=2, + sort_keys=True, + ) + + return { + "quantized_weights": sum( + 1 + for name in source_weight_map + if name.endswith(".weight") + and name in fp8_scale_map + and not is_ignored_by_modelopt(name, ignore_patterns) + ), + "bf16_fallback_weights": len(fallback_weight_names), + "preserved_ignored_weights": preserved_ignored_weight_count, + "added_scale_tensors": added_scale_count, + "output_shards": len(weights_by_file), + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Inject FP8 scales from ModelOpt backbone.pt into a diffusers export so " + "SGLang diffusion can load it natively." + ) + ) + parser.add_argument( + "--modelopt-hf-dir", + required=True, + help="ModelOpt --hf-ckpt-dir output, or its transformer subdirectory.", + ) + parser.add_argument( + "--modelopt-backbone-ckpt", + required=True, + help="Path to backbone.pt, or the directory that contains it.", + ) + parser.add_argument( + "--output-dir", + required=True, + help="Directory to write the converted SGLang transformer checkpoint.", + ) + parser.add_argument( + "--base-transformer-dir", + help=( + "Original BF16 transformer directory (or parent model dir). Required when " + "BF16 fallback layers are enabled." + ), + ) + parser.add_argument( + "--model-type", + choices=["auto", "flux1", "flux2", "none"], + default="auto", + help=( + "Optional model-family BF16 fallback profile. 'none' uses the generic " + "conversion path. 'auto' enables the validated FLUX.1 / FLUX.2 " + "fallback set when the export config matches those transformer classes." + ), + ) + parser.add_argument( + "--keep-bf16-pattern", + action="append", + default=[], + help=( + "Regex matched against module names without the trailing .weight. " + "Matching weights are copied from --base-transformer-dir instead of " + "staying in FP8." + ), + ) + parser.add_argument( + "--maxbound", + type=float, + default=FP8_E4M3_MAXBOUND, + help="FP8 maxbound used to turn ModelOpt amax into a scale. E4M3 uses 448.", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Replace --output-dir if it already exists.", + ) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + stats = convert_modelopt_fp8_checkpoint( + modelopt_hf_dir=args.modelopt_hf_dir, + modelopt_backbone_ckpt=args.modelopt_backbone_ckpt, + output_dir=args.output_dir, + base_transformer_dir=args.base_transformer_dir, + model_type=args.model_type, + keep_bf16_patterns=args.keep_bf16_pattern, + maxbound=args.maxbound, + overwrite=args.overwrite, + ) + print(json.dumps(stats, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main()