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

[Diffusion] modelopt diffusion fp8 support for flux1/flux2 and wan2.2 (#22365)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-04-10 20:56:57 +08:00
committed by GitHub
parent 8ba9646044
commit 1ff51555f2
10 changed files with 1891 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py
+142
View File
@@ -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"]))
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-modelopt-quant/SKILL.md
+308
View File
@@ -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 <model-name> \
--override-model-path <hf-repo-or-local-model> \
--model-dtype <Half|BFloat16> \
--format <fp8|nvfp4> \
--batch-size 1 \
--calib-size <calib-size> \
--n-steps <calib-steps> \
--quantize-mha \
--prompts-file <prompt-file> \
--quantized-torch-ckpt-save-path <out>/ckpt \
--hf-ckpt-dir <out>/hf
```
For multi-transformer models:
- quantize each backbone deliberately
- keep each output directory separate
- save both `backbone.pt` and the matching `hf/<component>` 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 <out>/hf \
--modelopt-backbone-ckpt <out>/ckpt/backbone.pt \
--base-transformer-dir <base-model-transformer-dir> \
--output-dir <out>/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 <base-model> \
--transformer-path <quantized-transformer> \
--prompt "<prompt>" \
--seed <seed> \
--save-output
```
Multi-transformer example:
```bash
sglang generate \
--model-path <base-model> \
--transformer-path <quantized-transformer> \
--transformer-2-path <another-transformer-or-bf16-override> \
--prompt "<prompt>" \
--seed <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 `--<component>-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 <base-model> \
--model-id <optional-native-model-id> \
--prompt "<prompt>" \
--width <w> \
--height <h> \
--num-inference-steps <steps> \
--guidance-scale <cfg> \
--seed <seed> \
--candidate-transformer-path <quantized-transformer> \
--output-json <report.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 |
python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py
+3 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+175 -25
View File
@@ -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."""
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+98 -12
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/flux_2.py
+15 -5
View File
@@ -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
),
)
python/sglang/multimodal_gen/runtime/utils/quantization_utils.py
+121 -73
View File
@@ -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:
python/sglang/multimodal_gen/tools/compare_diffusion_trajectory_similarity.py
+498
View File
@@ -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()
python/sglang/multimodal_gen/tools/convert_modelopt_fp8_checkpoint.py
+494
View File
@@ -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()