[Fix] Fix FP8 Online Quantization (#26415)

Co-authored-by: yichiche@amd.com <jacky.cheng>
This commit is contained in:
Colin Z
2026-05-28 22:00:09 -07:00
committed by GitHub
co-authored by yichiche@amd.com
parent f16816f043
commit 226649e3b7
4 changed files with 18 additions and 7 deletions
@@ -75,7 +75,11 @@ logger = logging.getLogger(__name__)
class Fp8Config(QuantizationConfig): class Fp8Config(QuantizationConfig):
"""Config class for FP8.""" """Config class for FP8.
No-arg ``Fp8Config()`` selects online (post-load) weight quantization:
``is_checkpoint_fp8_serialized=False`` with ``activation_scheme="dynamic"``.
"""
def __init__( def __init__(
self, self,
@@ -44,7 +44,8 @@ class Mxfp4Config(QuantizationConfig):
""" """
MXFP4 quantization config for diffusion models. MXFP4 quantization config for diffusion models.
Supports online quantization from unquantized BF16/FP16 checkpoints. Supports online quantization from unquantized BF16/FP16 checkpoints;
no-arg ``Mxfp4Config()`` selects that online (post-load) path.
Note: MXFP4 requires ROCm and MI350+ (gfx95x). Note: MXFP4 requires ROCm and MI350+ (gfx95x).
""" """
@@ -503,12 +503,16 @@ def _resolve_quant_config(
) )
# modelslim requires a per-layer quant description file; load it from # modelslim requires a per-layer quant description file; load it from
# the component directory rather than returning an empty config. # the component directory rather than constructing an empty config.
if server_args.quantization == "modelslim": if server_args.quantization == "modelslim":
return get_quant_config(hf_config, component_model_path) return get_quant_config(hf_config, component_model_path)
# Online-quant convention: for `fp8` and `mxfp4`, a no-arg
# QuantizationConfig() selects the post-load path -- weights load
# in source dtype and are quantized in
# process_weights_after_loading.
quant_cls = get_quantization_config(server_args.quantization) quant_cls = get_quantization_config(server_args.quantization)
return quant_cls.from_config({}) return quant_cls()
quant_config = get_quant_config(hf_config, component_model_path) quant_config = get_quant_config(hf_config, component_model_path)
if quant_config is None and server_args.transformer_weights_path: if quant_config is None and server_args.transformer_weights_path:
@@ -1331,9 +1331,11 @@ class ServerArgs(DisaggArgsMixin):
help=( help=(
"Quantization method for the transformer. If omitted, the method is " "Quantization method for the transformer. If omitted, the method is "
"auto-detected from the checkpoint config or safetensors metadata when " "auto-detected from the checkpoint config or safetensors metadata when "
"possible. Applies to both pre-quantized checkpoints and online " "possible. Use this flag to override auto-detection. "
"quantization. Use this flag to override auto-detection. " "Online (post-load) quantization from a BF16/FP16 checkpoint "
"Options: 'fp8', 'mxfp8', 'mxfp4', 'mxfp4_npu', 'modelslim'. " "is supported for 'fp8' and 'mxfp4'. Other methods "
"('modelopt', 'modelopt_fp8', 'modelopt_fp4', 'mxfp8', "
"'mxfp4_npu', 'modelslim') require a pre-quantized checkpoint. "
"Note: 'mxfp4' targets ROCm + MI350+ (gfx95x); " "Note: 'mxfp4' targets ROCm + MI350+ (gfx95x); "
"'mxfp4_npu' / 'mxfp8' target Ascend NPU (A5 series for mxfp4_npu)." "'mxfp4_npu' / 'mxfp8' target Ascend NPU (A5 series for mxfp4_npu)."
), ),