Revert "[AMD][Quantization] Online MXFP4 quantization 2/N - FP8 to MXFP4 requantization on AMD GPUs" (#28213)
This commit is contained in:
Lianmin Zheng
@@ -137,9 +137,6 @@ def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
|
||||
return config.index_n_heads
|
||||
|
||||
|
||||
REQUANTIZATION_METHODS = ["quark_mxfp4"]
|
||||
|
||||
|
||||
def get_num_indexer_layers(config) -> int:
|
||||
"""Layer count for the global indexer-topk capturer's host buffer.
|
||||
|
||||
@@ -1068,7 +1065,7 @@ class ModelConfig:
|
||||
log_str = f"quant={quant_method}"
|
||||
|
||||
# Append interesting fields if they exist
|
||||
for field in ["bits", "quant_algo", "fmt", "requantization_method"]:
|
||||
for field in ["bits", "quant_algo", "fmt"]:
|
||||
if field in quant_cfg:
|
||||
log_str += f", {field}={quant_cfg[field]}"
|
||||
|
||||
@@ -1263,10 +1260,6 @@ class ModelConfig:
|
||||
f"Using draft model's detected quantization: {quant_method}"
|
||||
)
|
||||
self.quantization = quant_method
|
||||
elif self.quantization in REQUANTIZATION_METHODS:
|
||||
logger.info_once(
|
||||
f"Requantizing from quant_method='{quant_method}' to the requested online quantization='{self.quantization}'. Beware that requantization may incur a loss in accuracy, the requantized model should be re-validated/re-evaluated. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Quantization method specified in the model config "
|
||||
|
||||
@@ -77,7 +77,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
|
||||
"IPEXAWQLinearMethod",
|
||||
"PetitNvFp4LinearMethod",
|
||||
"QuarkInt4Fp8LinearMethod",
|
||||
"QuarkLinearMethod",
|
||||
]
|
||||
|
||||
_is_cpu = is_cpu()
|
||||
|
||||
@@ -1,44 +0,0 @@
|
||||
"""
|
||||
Utilities to manage the dequantization of weights.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
block_quant_dequant,
|
||||
inverse_transform_scale_ue8m0,
|
||||
)
|
||||
from sglang.srt.utils import set_weight_attrs
|
||||
|
||||
|
||||
def copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None:
|
||||
"""Copies any attrs present in `old` but not in `new` to `new`"""
|
||||
new_attrs = set(dir(new))
|
||||
attrs_to_set = {}
|
||||
for attr in dir(old):
|
||||
if attr not in new_attrs:
|
||||
attrs_to_set[attr] = getattr(old, attr)
|
||||
set_weight_attrs(new, attrs_to_set)
|
||||
|
||||
|
||||
def dequantize_fp8(
|
||||
w_q: torch.Tensor,
|
||||
w_s: torch.Tensor,
|
||||
block_size: list[int],
|
||||
format_ue8m0: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Dequantizes `w_q` to bfloat16.
|
||||
"""
|
||||
if format_ue8m0:
|
||||
# TODO this is only needed for Blackwell
|
||||
w_s = inverse_transform_scale_ue8m0(w_s, mn=w_q.shape[-2])
|
||||
|
||||
w_dequant = block_quant_dequant(
|
||||
w_q,
|
||||
w_s,
|
||||
block_size=block_size,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
return w_dequant
|
||||
@@ -357,9 +357,8 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
self.use_aiter_fp8_per_token = envs.SGLANG_USE_AITER_FP8_PER_TOKEN.get()
|
||||
self.use_per_token_if_dynamic = False
|
||||
|
||||
@staticmethod
|
||||
def validate_block_quant_shapes(
|
||||
quant_config,
|
||||
self,
|
||||
input_size: int,
|
||||
input_size_per_partition: int,
|
||||
output_size: int,
|
||||
@@ -369,8 +368,8 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
):
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
block_n, block_k = (
|
||||
quant_config.weight_block_size[0],
|
||||
quant_config.weight_block_size[1],
|
||||
self.quant_config.weight_block_size[0],
|
||||
self.quant_config.weight_block_size[1],
|
||||
)
|
||||
|
||||
if skip_block_quant_check:
|
||||
@@ -398,114 +397,6 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
f"weight quantization block_n = {block_n}."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_fp8_weight_(
|
||||
layer: torch.nn.Module,
|
||||
block_quant: bool,
|
||||
quant_config,
|
||||
use_mxfp8: bool,
|
||||
output_size_per_partition: int,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader,
|
||||
is_checkpoint_fp8_serialized: bool,
|
||||
skip_block_quant_check: bool = False,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
"""
|
||||
Registers weights into `layer`. This static method can be reused by other quantization methods that require loading FP8 checkpoints first (e.g. requantization to other formats as MXFP4).
|
||||
"""
|
||||
# Copy the layer attributes
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
|
||||
if block_quant:
|
||||
block_n, block_k = quant_config.weight_block_size
|
||||
Fp8LinearMethod.validate_block_quant_shapes(
|
||||
quant_config,
|
||||
input_size,
|
||||
input_size_per_partition,
|
||||
output_size,
|
||||
output_size_per_partition,
|
||||
output_partition_sizes,
|
||||
skip_block_quant_check,
|
||||
)
|
||||
else:
|
||||
block_n, block_k = None, None
|
||||
|
||||
# Create the weight
|
||||
weight_dtype = (
|
||||
torch.float8_e4m3fn if is_checkpoint_fp8_serialized else params_dtype
|
||||
)
|
||||
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,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
if is_checkpoint_fp8_serialized:
|
||||
if block_quant:
|
||||
if hasattr(quant_config, "activation_scheme"):
|
||||
assert quant_config.activation_scheme == "dynamic"
|
||||
elif hasattr(quant_config, "linear_activation_scheme"):
|
||||
assert quant_config.linear_activation_scheme == "dynamic"
|
||||
if use_mxfp8 and not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"MXFP8 requires fp8-serialized checkpoint for linear layers."
|
||||
)
|
||||
|
||||
scale_dtype = torch.uint8 if use_mxfp8 else torch.float32
|
||||
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.empty
|
||||
scale = BlockQuantScaleParameter(
|
||||
data=scale_init(
|
||||
(output_size_per_partition + block_n - 1) // block_n,
|
||||
(input_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scale.format_ue8m0 = use_mxfp8
|
||||
if scale_dtype != torch.uint8:
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale_inv", scale)
|
||||
else:
|
||||
scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale", scale)
|
||||
|
||||
# INPUT ACTIVATION SCALE
|
||||
if (
|
||||
hasattr(quant_config, "activation_scheme")
|
||||
and quant_config.activation_scheme == "static"
|
||||
) or (
|
||||
hasattr(quant_config, "linear_activation_scheme")
|
||||
and quant_config.linear_activation_scheme == "static"
|
||||
):
|
||||
scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("input_scale", scale)
|
||||
else:
|
||||
layer.register_parameter("input_scale", None)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -525,21 +416,85 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
layer.orig_dtype = params_dtype
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
|
||||
Fp8LinearMethod.create_fp8_weight_(
|
||||
layer,
|
||||
block_quant=self.block_quant,
|
||||
quant_config=self.quant_config,
|
||||
use_mxfp8=self.use_mxfp8,
|
||||
output_size_per_partition=output_size_per_partition,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
weight_loader=weight_loader,
|
||||
skip_block_quant_check=skip_block_quant_check,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
is_checkpoint_fp8_serialized=self.is_checkpoint_fp8_serialized,
|
||||
params_dtype=params_dtype,
|
||||
if self.block_quant:
|
||||
block_n, block_k = self.quant_config.weight_block_size
|
||||
self.validate_block_quant_shapes(
|
||||
input_size,
|
||||
input_size_per_partition,
|
||||
output_size,
|
||||
output_size_per_partition,
|
||||
output_partition_sizes,
|
||||
skip_block_quant_check,
|
||||
)
|
||||
|
||||
# Create the weight
|
||||
weight_dtype = (
|
||||
torch.float8_e4m3fn if self.is_checkpoint_fp8_serialized else params_dtype
|
||||
)
|
||||
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,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# If checkpoint is serialized fp8, load them.
|
||||
# Otherwise, wait until process_weights_after_loading.
|
||||
if self.is_checkpoint_fp8_serialized:
|
||||
# WEIGHT SCALE
|
||||
if self.block_quant:
|
||||
if hasattr(self.quant_config, "activation_scheme"):
|
||||
assert self.quant_config.activation_scheme == "dynamic"
|
||||
elif hasattr(self.quant_config, "linear_activation_scheme"):
|
||||
assert self.quant_config.linear_activation_scheme == "dynamic"
|
||||
if self.use_mxfp8 and not self.is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"MXFP8 requires fp8-serialized checkpoint for linear layers."
|
||||
)
|
||||
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
|
||||
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.empty
|
||||
scale = BlockQuantScaleParameter(
|
||||
data=scale_init(
|
||||
(output_size_per_partition + block_n - 1) // block_n,
|
||||
(input_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scale.format_ue8m0 = self.use_mxfp8
|
||||
if scale_dtype != torch.uint8:
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale_inv", scale)
|
||||
else:
|
||||
scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale", scale)
|
||||
|
||||
# INPUT ACTIVATION SCALE
|
||||
if (
|
||||
hasattr(self.quant_config, "activation_scheme")
|
||||
and self.quant_config.activation_scheme == "static"
|
||||
) or (
|
||||
hasattr(self.quant_config, "linear_activation_scheme")
|
||||
and self.quant_config.linear_activation_scheme == "static"
|
||||
):
|
||||
scale = PerTensorScaleParameter(
|
||||
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("input_scale", scale)
|
||||
else:
|
||||
layer.register_parameter("input_scale", None)
|
||||
|
||||
def process_weights_after_loading_block_quant(self, layer: Module) -> None:
|
||||
# If ROCm, normalize the weights and scales to e4m3fnuz
|
||||
@@ -921,29 +876,21 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def create_fp8_moe_weight_(
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
block_quant: bool,
|
||||
quant_config,
|
||||
use_mxfp8: bool,
|
||||
is_checkpoint_fp8_serialized: bool,
|
||||
is_fp4_expert: bool,
|
||||
params_dtype: torch.dtype,
|
||||
extra_weight_attrs: dict,
|
||||
with_bias: bool,
|
||||
with_bias: bool = False,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
"""
|
||||
Registers weights into `layer`. This static method can be reused by other quantization methods that require loading FP8 checkpoints first (e.g. requantization to other formats as MXFP4).
|
||||
"""
|
||||
self.with_bias = with_bias
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
if is_checkpoint_fp8_serialized:
|
||||
if self.quant_config.is_checkpoint_fp8_serialized:
|
||||
params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes(
|
||||
@@ -953,10 +900,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
is_packed=False,
|
||||
)
|
||||
|
||||
if block_quant:
|
||||
if self.block_quant:
|
||||
block_n, block_k = (
|
||||
quant_config.weight_block_size[0],
|
||||
quant_config.weight_block_size[1],
|
||||
self.quant_config.weight_block_size[0],
|
||||
self.quant_config.weight_block_size[1],
|
||||
)
|
||||
|
||||
padding_size = get_moe_padding_size(_use_aiter)
|
||||
@@ -979,7 +926,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
)
|
||||
|
||||
# WEIGHTS
|
||||
if is_fp4_expert:
|
||||
if self.is_fp4_expert:
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
@@ -1049,7 +996,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# BIAS (optional, e.g. GPT-OSS)
|
||||
if with_bias:
|
||||
if self.with_bias:
|
||||
w13_up_dim = (
|
||||
2 * intermediate_size_per_partition
|
||||
if layer.moe_runner_config.is_gated
|
||||
@@ -1070,7 +1017,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
set_weight_attrs(w2_weight_bias, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
if is_fp4_expert:
|
||||
if self.is_fp4_expert:
|
||||
fp4_block_k = 32
|
||||
fp4_scale_dtype = torch.float8_e8m0fnu if _use_aiter else torch.float32
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
@@ -1093,8 +1040,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
)
|
||||
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
|
||||
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
|
||||
elif block_quant:
|
||||
scale_dtype = torch.uint8 if use_mxfp8 else torch.float32
|
||||
elif self.block_quant:
|
||||
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
|
||||
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.ones
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
scale_init(
|
||||
@@ -1115,12 +1062,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
requires_grad=False,
|
||||
)
|
||||
# w13_weight and w2_weight are always requanted together
|
||||
w13_weight_scale.format_ue8m0 = use_mxfp8
|
||||
w2_weight_scale.format_ue8m0 = use_mxfp8
|
||||
w13_weight_scale.format_ue8m0 = self.use_mxfp8
|
||||
w2_weight_scale.format_ue8m0 = self.use_mxfp8
|
||||
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
|
||||
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
|
||||
|
||||
assert quant_config.activation_scheme == "dynamic"
|
||||
assert self.quant_config.activation_scheme == "dynamic"
|
||||
if get_moe_runner_backend().is_cutlass():
|
||||
self._ensure_cutlass_buffers_initialized(layer)
|
||||
|
||||
else:
|
||||
# Allocate 2 scales for w1 and w3 respectively.
|
||||
@@ -1155,14 +1103,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
# to ensure the weight scales are loaded in properly
|
||||
extra_weight_attrs.update(
|
||||
{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
|
||||
if block_quant
|
||||
if self.block_quant
|
||||
else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
|
||||
)
|
||||
|
||||
# If loading fp8 checkpoint, pass the weight loaders.
|
||||
# If loading an fp16 checkpoint, do not (we will quantize in
|
||||
# process_weights_after_loading()
|
||||
if quant_config.is_checkpoint_fp8_serialized:
|
||||
if self.quant_config.is_checkpoint_fp8_serialized:
|
||||
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
|
||||
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
|
||||
|
||||
@@ -1174,8 +1121,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
set_weight_attrs(w2_weight_scale1, extra_weight_attrs)
|
||||
|
||||
# INPUT_SCALES
|
||||
if quant_config.activation_scheme == "static":
|
||||
if not quant_config.is_checkpoint_fp8_serialized:
|
||||
if self.quant_config.activation_scheme == "static":
|
||||
if not self.quant_config.is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"Found static activation scheme for checkpoint that "
|
||||
"was not serialized fp8."
|
||||
@@ -1197,34 +1144,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Module,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size_per_partition: int,
|
||||
params_dtype: torch.dtype,
|
||||
with_bias: bool = False,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
Fp8MoEMethod.create_fp8_moe_weight_(
|
||||
layer=layer,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
block_quant=self.block_quant,
|
||||
quant_config=self.quant_config,
|
||||
use_mxfp8=self.use_mxfp8,
|
||||
is_checkpoint_fp8_serialized=self.quant_config.is_checkpoint_fp8_serialized,
|
||||
is_fp4_expert=self.is_fp4_expert,
|
||||
params_dtype=params_dtype,
|
||||
with_bias=with_bias,
|
||||
extra_weight_attrs=extra_weight_attrs,
|
||||
)
|
||||
|
||||
if self.block_quant and get_moe_runner_backend().is_cutlass():
|
||||
self._ensure_cutlass_buffers_initialized(layer)
|
||||
|
||||
def process_weights_after_loading_block_quant(self, layer: Module) -> None:
|
||||
# AMD FP4 experts: use aiter's native MXFP4 MoE path
|
||||
if _use_aiter and self.is_fp4_expert:
|
||||
|
||||
@@ -1216,8 +1216,6 @@ def block_quant_dequant(
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
*_, n, k = x_q_block.shape
|
||||
|
||||
# NOTE: This is very memory inefficient, results in *16384 memory requirement for scales
|
||||
# with block_size = [128, 128].
|
||||
# ... n_scale k_scale -> ... (n_scale block_n) (k_scale block_k)
|
||||
x_scale_repeat = x_s.repeat_interleave(block_n, dim=-2).repeat_interleave(
|
||||
block_k, dim=-1
|
||||
|
||||
@@ -1,23 +0,0 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/pull/31914
|
||||
import torch
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
|
||||
class CopyNumelCounter(TorchDispatchMode):
|
||||
"""
|
||||
Tracks total number of elements modified with `copy_`. Useful for keeping
|
||||
track of weight loading where underlying weights can be arbitrarily
|
||||
transformed (such as with `narrow`) before calling copy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.copied_numel = 0
|
||||
|
||||
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
out = func(*args, **kwargs)
|
||||
if func == torch.ops.aten.copy_.default:
|
||||
self.copied_numel += args[0].numel()
|
||||
return out
|
||||
@@ -14,7 +14,6 @@ from sglang.srt.layers.quantization.base_config import ( # noqa: E501
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
|
||||
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from sglang.srt.layers.quantization.quark.schemes import (
|
||||
QuarkLinearScheme,
|
||||
@@ -55,14 +54,13 @@ class QuarkConfig(QuantizationConfig):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: dict[str, Any] | None = None,
|
||||
quant_config: Optional[dict[str, Any]] = None,
|
||||
hf_config: "PretrainedConfig | None" = None,
|
||||
kv_cache_group: Optional[list[str]] = None,
|
||||
kv_cache_config: Optional[dict[str, Any]] = None,
|
||||
pack_method: str = "reorder",
|
||||
is_prequantized: bool = False,
|
||||
online_scheme: Optional[str] = None,
|
||||
dequantization_config: Optional[QuantizationConfig] = None,
|
||||
):
|
||||
super().__init__()
|
||||
if kv_cache_group is None:
|
||||
@@ -86,22 +84,17 @@ class QuarkConfig(QuantizationConfig):
|
||||
self.pack_method = pack_method
|
||||
self.exclude_layers = cast(list[str], self.quant_config.get("exclude", []))
|
||||
self.is_prequantized = is_prequantized
|
||||
self.dequantization_config = dequantization_config
|
||||
|
||||
self.packed_modules_mapping = self.quant_config["packed_modules_mapping"]
|
||||
self._quantized_layers = set()
|
||||
|
||||
if isinstance(self.dequantization_config, Fp8Config):
|
||||
self.weight_block_size = self.dequantization_config.weight_block_size
|
||||
|
||||
@property
|
||||
def quantized_layers(self) -> tuple[dict[str, int], int]:
|
||||
# Count layers per type (last two parts after ".")
|
||||
type_counts: dict[str, int] = {}
|
||||
for name in self._quantized_layers:
|
||||
parts = name.split(".")
|
||||
layer_type = ".".join(parts[-2:]) if len(parts) >= 2 else parts[-1]
|
||||
type_counts[layer_type] = type_counts.get(layer_type, 0) + 1
|
||||
return dict(sorted(type_counts.items())), len(self._quantized_layers)
|
||||
def quantized_layers(self) -> tuple[list[str], int]:
|
||||
# Extract unique layer types (last part after ".")
|
||||
layer_types = sorted(
|
||||
set(name.split(".")[-1] for name in self._quantized_layers)
|
||||
)
|
||||
return layer_types, len(self._quantized_layers)
|
||||
|
||||
def get_linear_method(self) -> "QuarkLinearMethod":
|
||||
return QuarkLinearMethod(self)
|
||||
@@ -130,18 +123,14 @@ class QuarkConfig(QuantizationConfig):
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional["QuantizeMethodBase"]:
|
||||
# Check if the layer is skipped for quantization.
|
||||
|
||||
if should_ignore_layer(
|
||||
prefix,
|
||||
ignore=self.exclude_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
if isinstance(layer, LinearBase):
|
||||
if self.dequantization_config is not None:
|
||||
# In case of online requantization, "exclude" means keeping the original precision.
|
||||
# NOTE: Only FP8 supported for now.
|
||||
return Fp8LinearMethod(quant_config=self.dequantization_config)
|
||||
else:
|
||||
return UnquantizedLinearMethod()
|
||||
return UnquantizedLinearMethod()
|
||||
elif isinstance(layer, RadixAttention):
|
||||
return QuarkKVCacheMethod(self)
|
||||
return None
|
||||
@@ -167,32 +156,6 @@ class QuarkConfig(QuantizationConfig):
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":
|
||||
if config["quant_method"] != "quark":
|
||||
assert "requantization_method" in config
|
||||
|
||||
if (
|
||||
config["quant_method"] == "fp8"
|
||||
and config["requantization_method"] == "quark_mxfp4"
|
||||
and config["activation_scheme"] == "dynamic"
|
||||
):
|
||||
hf_config = config["hf_config"]
|
||||
quant_config = QuarkConfig._create_online_mxfp4_config(
|
||||
model_type=hf_config.model_type
|
||||
)
|
||||
dequantization_config = Fp8Config.from_config(config)
|
||||
quark_config = cls(
|
||||
quant_config=quant_config,
|
||||
hf_config=hf_config,
|
||||
is_prequantized=False,
|
||||
dequantization_config=dequantization_config,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Requantization into {config['requantization_method']} is not supported, from the original quant_method={config['quant_method']} and activation_scheme={config['activation_scheme']}. "
|
||||
)
|
||||
|
||||
return quark_config
|
||||
|
||||
export_config = config.get("export")
|
||||
if export_config is None:
|
||||
raise ValueError(
|
||||
@@ -277,17 +240,7 @@ class QuarkConfig(QuantizationConfig):
|
||||
"re:.*embed_tokens",
|
||||
]
|
||||
|
||||
# Exclusion for accuracy adapted from
|
||||
# https://huggingface.co/amd/DeepSeek-V3.2-mxfp4/blob/main/config.json
|
||||
if model_type in ["deepseek_v3", "deepseek_v32"]:
|
||||
exclude.extend(
|
||||
[
|
||||
"re:.*model.layers.61.*",
|
||||
"re:.*self_attn.*",
|
||||
"re:.*mlp.gate$",
|
||||
]
|
||||
)
|
||||
elif model_type == "qwen3_5_moe":
|
||||
if model_type == "qwen3_5_moe":
|
||||
# Exclusion for accuracy adapted from
|
||||
# https://huggingface.co/amd/Qwen3.5-397B-A17B-MXFP4/blob/main/config.json
|
||||
exclude.extend(
|
||||
@@ -492,7 +445,6 @@ class QuarkConfig(QuantizationConfig):
|
||||
weight_config,
|
||||
input_config,
|
||||
is_checkpoint_mxfp4_serialized=self.is_prequantized,
|
||||
dequantization_config=self.dequantization_config,
|
||||
)
|
||||
if self._is_fp8_w8a8(weight_config, input_config):
|
||||
is_fp8_w8a8_supported = self._check_scheme_supported(
|
||||
@@ -543,7 +495,6 @@ class QuarkConfig(QuantizationConfig):
|
||||
weight_config,
|
||||
input_config,
|
||||
is_checkpoint_mxfp4_serialized=self.is_prequantized,
|
||||
dequantization_config=self.dequantization_config,
|
||||
)
|
||||
elif self._is_fp8_w8a8(weight_config, input_config):
|
||||
return QuarkW8A8FP8MoE(weight_config, input_config)
|
||||
@@ -590,7 +541,6 @@ class QuarkLinearMethod(LinearMethodBase):
|
||||
|
||||
def __init__(self, quantization_config: QuarkConfig):
|
||||
self.quantization_config = quantization_config
|
||||
self.quant_config = quantization_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scheme.process_weights_after_loading(layer)
|
||||
|
||||
@@ -1,19 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
|
||||
from sglang.srt.layers.quantization import QuantizationConfig
|
||||
from sglang.srt.layers.quantization.dequantization import (
|
||||
copy_missing_attrs,
|
||||
dequantize_fp8,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
|
||||
from sglang.srt.layers.quantization.online_quantization import CopyNumelCounter
|
||||
from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme
|
||||
from sglang.srt.utils import is_hip
|
||||
from sglang.srt.utils.common import direct_register_custom_op, mxfp_supported
|
||||
@@ -170,14 +162,12 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
|
||||
weight_quant_spec: dict[str, Any],
|
||||
input_quant_spec: dict[str, Any],
|
||||
is_checkpoint_mxfp4_serialized: bool = True,
|
||||
dequantization_config: QuantizationConfig | None = None,
|
||||
):
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.qscheme = "per_group"
|
||||
self.weight_quant_spec = weight_quant_spec
|
||||
self.input_quant_spec = input_quant_spec
|
||||
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
|
||||
self.dequantization_config = dequantization_config
|
||||
|
||||
if not self.is_checkpoint_mxfp4_serialized:
|
||||
if not mxfp_supported():
|
||||
@@ -186,7 +176,7 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
|
||||
"FP4 hardware support (gfx95x, e.g. MI355x)."
|
||||
)
|
||||
logger.info_once(
|
||||
"Using online MXFP4 quantization in dense linear layers from a higher precision checkpoint. Beware that this optimization may degrade prediction quality - please validate your model accuracy. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
|
||||
"Using online MXFP4 quantization from a higher precision checkpoint. Beware that this optimization may degrade prediction quality - please validate your model accuracy. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -214,96 +204,38 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
|
||||
|
||||
layer.logical_widths = output_partition_sizes
|
||||
|
||||
# If dequantization_config is provided, we need to create FP8 weights first
|
||||
# for dequantization from FP8 checkpoint to MXFP4
|
||||
if self.dequantization_config is not None:
|
||||
if not isinstance(self.dequantization_config, Fp8Config):
|
||||
raise NotImplementedError(
|
||||
f"Requantization in QuarkW4A4MXFP4 from {self.dequantization_config.__class__.__name__} is not supported, only Fp8Config is supported."
|
||||
)
|
||||
# Create FP8 weights for re-quantization from FP8 checkpoint
|
||||
# Extract necessary parameters from dequantization_config
|
||||
self.weight_block_size = self.dequantization_config.weight_block_size
|
||||
original_weight_loader = weight_loader
|
||||
if not self.is_checkpoint_mxfp4_serialized:
|
||||
weight_loader = self.get_online_mxfp4_weight_loader(layer, weight_loader)
|
||||
|
||||
if self.dequantization_config.use_mxfp8:
|
||||
raise NotImplementedError(
|
||||
"use_mxfp8=True is not supported in Quark MXFP4 requantization."
|
||||
)
|
||||
# WEIGHT
|
||||
# Both serialized and online quantization use packed uint8 format
|
||||
weight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=2,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
block_quant = self.weight_block_size is not None
|
||||
|
||||
if not block_quant:
|
||||
raise NotImplementedError(
|
||||
"Only block_quant=True is supported in Quark MXFP4 requantization, got block_quant=False."
|
||||
)
|
||||
|
||||
layer._fp8_weight_loaded_numel = 0
|
||||
layer._load_device = torch.get_default_device()
|
||||
layer._fp8_weight_loading_lock = threading.Lock()
|
||||
layer._fp8_weight_materialized = False
|
||||
|
||||
# Wrap the weight loader to handle FP8->MXFP4 conversion
|
||||
fp8_to_mxfp4_weight_loader = self.get_online_fp8_to_mxfp4_weight_loader(
|
||||
layer, weight_loader
|
||||
)
|
||||
|
||||
# Create FP8 MoE weight parameters on meta device to avoid device memory overhead during weight loading, as the resulting model uses MXFP4 using less device memory.
|
||||
# The weight loader handles progressive FP8 weight materialization on device.
|
||||
with torch.device("meta"):
|
||||
Fp8LinearMethod.create_fp8_weight_(
|
||||
layer=layer,
|
||||
block_quant=block_quant,
|
||||
quant_config=self.dequantization_config,
|
||||
use_mxfp8=False,
|
||||
output_size_per_partition=output_size_per_partition,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
weight_loader=fp8_to_mxfp4_weight_loader,
|
||||
is_checkpoint_fp8_serialized=True,
|
||||
params_dtype=params_dtype,
|
||||
skip_block_quant_check=False,
|
||||
input_size=kwargs.get("input_size", input_size_per_partition),
|
||||
output_size=kwargs.get("output_size", output_size_per_partition),
|
||||
)
|
||||
|
||||
# NOTE: ideally, weight_loader should be refactored to be aware of `param_name`.
|
||||
layer.weight._param_name = "weight"
|
||||
layer.weight_scale_inv._param_name = "weight_scale_inv"
|
||||
else:
|
||||
original_weight_loader = weight_loader
|
||||
if not self.is_checkpoint_mxfp4_serialized:
|
||||
weight_loader = self.get_online_mxfp4_weight_loader(
|
||||
layer, weight_loader
|
||||
)
|
||||
|
||||
# WEIGHT
|
||||
# Both serialized and online quantization use packed uint8 format
|
||||
weight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // 2,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_dim=1,
|
||||
packed_factor=2,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // OCP_MX_BLOCK_SIZE,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=original_weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
# WEIGHT SCALE
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // OCP_MX_BLOCK_SIZE,
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=original_weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def get_online_mxfp4_weight_loader(
|
||||
self,
|
||||
@@ -338,119 +270,6 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
|
||||
|
||||
return online_mxfp4_weight_loader
|
||||
|
||||
def get_online_fp8_to_mxfp4_weight_loader(
|
||||
self,
|
||||
layer,
|
||||
original_weight_loader: Callable,
|
||||
) -> Callable:
|
||||
"""
|
||||
Wrap the original weight loader to perform FP8 to MXFP4 requantization.
|
||||
|
||||
This loader handles:
|
||||
1. Loading FP8 weights and weight_scale_inv parameters
|
||||
2. Waiting for all shards (e.g., q_proj, k_proj, v_proj) to be loaded
|
||||
3. Dequantizing FP8 -> BF16
|
||||
4. Requantizing BF16 -> MXFP4
|
||||
"""
|
||||
|
||||
def online_fp8_to_mxfp4_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
shard_id: int | str | None = None,
|
||||
):
|
||||
param_name = getattr(param, "_param_name", None)
|
||||
|
||||
is_weight_or_weight_scale = "weight" in param_name
|
||||
is_weight = param_name == "weight"
|
||||
is_weight_scale_inv = param_name == "weight_scale_inv"
|
||||
|
||||
# Sanity multi-threaded load check.
|
||||
assert torch.cuda.current_device() == layer._load_device.index
|
||||
|
||||
with layer._fp8_weight_loading_lock:
|
||||
# Materialize FP8 parameters on first load on device (there may be several shards for a single layer parameter, e.g. q_proj, k_proj, v_proj).
|
||||
|
||||
if is_weight_or_weight_scale and not layer._fp8_weight_materialized:
|
||||
|
||||
# Sanity check.
|
||||
assert layer.weight.device.type == "meta"
|
||||
|
||||
materialized_tensor = layer.weight.__class__(
|
||||
data=torch.empty_like(
|
||||
layer.weight.data, device=layer._load_device
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=layer.weight._weight_loader,
|
||||
)
|
||||
copy_missing_attrs(layer.weight, materialized_tensor)
|
||||
layer.weight = materialized_tensor
|
||||
|
||||
# Sanity check.
|
||||
assert layer.weight_scale_inv.device.type == "meta" # Sanity check.
|
||||
|
||||
materialized_tensor = layer.weight_scale_inv.__class__(
|
||||
data=torch.empty_like(
|
||||
layer.weight_scale_inv.data, device=layer._load_device
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=layer.weight_scale_inv._weight_loader,
|
||||
)
|
||||
copy_missing_attrs(layer.weight_scale_inv, materialized_tensor)
|
||||
layer.weight_scale_inv = materialized_tensor
|
||||
|
||||
# Mark as materialized to prevent other threads from doing it again.
|
||||
layer._fp8_weight_materialized = True
|
||||
|
||||
if is_weight:
|
||||
param = layer.weight
|
||||
elif is_weight_scale_inv:
|
||||
param = layer.weight_scale_inv
|
||||
|
||||
kwargs = {}
|
||||
if shard_id is not None:
|
||||
kwargs["loaded_shard_id"] = shard_id
|
||||
|
||||
# Track how much data we are actually loading (`narrow` used in weight loader)
|
||||
copy_numel_counter = CopyNumelCounter()
|
||||
with copy_numel_counter:
|
||||
original_weight_loader(param, loaded_weight, **kwargs)
|
||||
|
||||
with layer._fp8_weight_loading_lock:
|
||||
if is_weight_or_weight_scale:
|
||||
layer._fp8_weight_loaded_numel += copy_numel_counter.copied_numel
|
||||
|
||||
target_numel = layer.weight.numel() + layer.weight_scale_inv.numel()
|
||||
|
||||
# Perform requantization outside the lock (but only if we're the chosen thread)
|
||||
if layer._fp8_weight_loaded_numel == target_numel and hasattr(
|
||||
layer, "weight_scale_inv"
|
||||
):
|
||||
assert layer.weight.device.type != "meta"
|
||||
|
||||
# FP8 -> BF16 dequantization.
|
||||
weight_bf16 = dequantize_fp8(
|
||||
layer.weight,
|
||||
layer.weight_scale_inv,
|
||||
block_size=self.weight_block_size,
|
||||
)
|
||||
|
||||
# BF16 -> MXFP4 requantization.
|
||||
weight_mxfp4, weight_mxfp4_scale = dynamic_mxfp4_quant(weight_bf16)
|
||||
|
||||
layer.weight = torch.nn.Parameter(weight_mxfp4, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
weight_mxfp4_scale, requires_grad=False
|
||||
)
|
||||
|
||||
# Clean up FP8 parameters and tracking attributes
|
||||
del layer.weight_scale_inv
|
||||
del layer._load_device
|
||||
del weight_bf16
|
||||
|
||||
return online_fp8_to_mxfp4_weight_loader
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
|
||||
@@ -3,20 +3,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.utils import get_moe_weight_sizes
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
from sglang.srt.layers.quantization.dequantization import (
|
||||
copy_missing_attrs,
|
||||
dequantize_fp8,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
|
||||
from sglang.srt.layers.quantization.online_quantization import CopyNumelCounter
|
||||
from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme
|
||||
from sglang.srt.utils import (
|
||||
get_bool_env_var,
|
||||
@@ -59,12 +51,10 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
||||
weight_config: dict[str, Any],
|
||||
input_config: dict[str, Any],
|
||||
is_checkpoint_mxfp4_serialized: bool = True,
|
||||
dequantization_config: QuantizationConfig | None = None,
|
||||
):
|
||||
self.weight_quant = weight_config
|
||||
self.input_quant = input_config
|
||||
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
|
||||
self.dequantization_config = dequantization_config
|
||||
|
||||
weight_qscheme = self.weight_quant.get("qscheme")
|
||||
input_qscheme = self.input_quant.get("qscheme")
|
||||
@@ -106,60 +96,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
||||
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
original_weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
with_bias = extra_weight_attrs.get("with_bias", False)
|
||||
|
||||
# Handle FP8 to MXFP4 requantization
|
||||
if self.dequantization_config is not None:
|
||||
if not isinstance(self.dequantization_config, Fp8Config):
|
||||
raise NotImplementedError(
|
||||
f"Requantization in QuarkW4A4MXFp4MoEMethod from {self.dequantization_config.__class__.__name__} is not supported, only Fp8Config is supported."
|
||||
)
|
||||
|
||||
if self.dequantization_config.use_mxfp8:
|
||||
raise NotImplementedError(
|
||||
"use_mxfp8=True is not supported in Quark MXFP4 requantization."
|
||||
)
|
||||
|
||||
block_quant = self.dequantization_config.weight_block_size is not None
|
||||
|
||||
if not block_quant:
|
||||
raise NotImplementedError(
|
||||
"Only block_quant=True is supported in Quark MXFP4 requantization, got block_quant=False."
|
||||
)
|
||||
|
||||
# `_fp8_loaded_numel` is used to trigger FP8 -> MXFP4 requantization once all weights are loaded.
|
||||
# `_fp8_materialized` is used to ensure only one thread materializes weights from meta device.
|
||||
layer._fp8_loaded_numel = 0
|
||||
layer._fp8_materialized = False
|
||||
layer._load_device = torch.get_default_device()
|
||||
layer._fp8_loading_lock = threading.Lock()
|
||||
|
||||
# Custom weight loader handling FP8->MXFP4 conversion.
|
||||
fp8_to_mxfp4_weight_loader = self.get_online_fp8_to_mxfp4_weight_loader(
|
||||
layer, original_weight_loader
|
||||
)
|
||||
|
||||
extra_weight_attrs["weight_loader"] = fp8_to_mxfp4_weight_loader
|
||||
# Create FP8 MoE weight parameters on meta device to avoid device memory overhead during weight loading, as the resulting model uses MXFP4 using less device memory.
|
||||
# The weight loader handles progressive FP8 weight materialization on device.
|
||||
with torch.device("meta"):
|
||||
Fp8MoEMethod.create_fp8_moe_weight_(
|
||||
layer=layer,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
block_quant=block_quant,
|
||||
quant_config=self.dequantization_config,
|
||||
use_mxfp8=False,
|
||||
is_checkpoint_fp8_serialized=True,
|
||||
is_fp4_expert=False,
|
||||
params_dtype=params_dtype,
|
||||
extra_weight_attrs=extra_weight_attrs,
|
||||
with_bias=with_bias,
|
||||
)
|
||||
return
|
||||
|
||||
w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes(
|
||||
intermediate_size_per_partition,
|
||||
is_aiter_moe=_use_aiter,
|
||||
@@ -176,61 +112,40 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
||||
},
|
||||
)
|
||||
|
||||
if self.is_checkpoint_mxfp4_serialized:
|
||||
weight_loader = original_weight_loader
|
||||
weight_device = torch.get_default_device()
|
||||
weight_dtype = torch.uint8
|
||||
else:
|
||||
# Online quantization: use original dtype and meta device
|
||||
weight_loader = self.get_online_weight_loader(layer, original_weight_loader)
|
||||
weight_device = torch.device("meta")
|
||||
weight_dtype = params_dtype
|
||||
|
||||
params_dtype = torch.uint8
|
||||
|
||||
layer._load_device = torch.get_default_device()
|
||||
layer._w13_loaded_numel = 0
|
||||
layer._w2_loaded_numel = 0
|
||||
original_weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
|
||||
if self.is_checkpoint_mxfp4_serialized:
|
||||
weight_loader = original_weight_loader
|
||||
else:
|
||||
weight_loader = self.get_online_weight_loader(layer, original_weight_loader)
|
||||
|
||||
extra_weight_attrs["weight_loader"] = weight_loader
|
||||
|
||||
# WEIGHTS
|
||||
w13_shape = (
|
||||
num_experts,
|
||||
2 * intermediate_size_per_partition,
|
||||
hidden_size // 2 if self.is_checkpoint_mxfp4_serialized else hidden_size,
|
||||
)
|
||||
# WEIGHTS — always uint8 (packed mxfp4), always on device
|
||||
w13_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
w13_shape,
|
||||
dtype=weight_dtype,
|
||||
device=weight_device,
|
||||
num_experts,
|
||||
w13_up_dim,
|
||||
hidden_size // 2,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_weight", w13_weight)
|
||||
|
||||
set_weight_attrs(w13_weight, extra_weight_attrs)
|
||||
|
||||
w2_shape = (
|
||||
num_experts,
|
||||
hidden_size,
|
||||
(
|
||||
intermediate_size_per_partition // 2
|
||||
if self.is_checkpoint_mxfp4_serialized
|
||||
else intermediate_size_per_partition
|
||||
),
|
||||
)
|
||||
w2_weight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
w2_shape,
|
||||
dtype=weight_dtype,
|
||||
device=weight_device,
|
||||
num_experts,
|
||||
hidden_size,
|
||||
w2_down_dim,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_weight", w2_weight)
|
||||
|
||||
set_weight_attrs(w2_weight, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_SCALES
|
||||
@@ -265,7 +180,7 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
||||
|
||||
def get_online_weight_loader(self, layer, original_weight_loader):
|
||||
"""
|
||||
Wrap the original weight loader to perform online MXFP4 quantization for MoE layers.
|
||||
Wrap the original weight loader to perform online MXFP4 quantization.
|
||||
"""
|
||||
|
||||
def online_mxfp4_moe_weight_loader(
|
||||
@@ -280,290 +195,29 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
|
||||
"Online MXFP4 quantization for MoE is only supported on AMD GPUs."
|
||||
)
|
||||
|
||||
# Determine which weight parameter we're loading (w13 or w2)
|
||||
is_w13 = "w13" in weight_name
|
||||
is_w2 = "w2" in weight_name
|
||||
# Materialize on device the loaded weight.
|
||||
loaded_weight = loaded_weight.to(param.device)
|
||||
|
||||
# Initialize weight on device if first load
|
||||
if is_w13 and layer._w13_loaded_numel == 0:
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
torch.empty_like(param.data, device=layer._load_device),
|
||||
requires_grad=False,
|
||||
)
|
||||
param = layer.w13_weight
|
||||
elif is_w2 and layer._w2_loaded_numel == 0:
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
torch.empty_like(param.data, device=layer._load_device),
|
||||
requires_grad=False,
|
||||
)
|
||||
param = layer.w2_weight
|
||||
# Quantize the high-precision shard loaded_weight to MXFP4.
|
||||
qweight, weight_scale = dynamic_mxfp4_quant(loaded_weight)
|
||||
|
||||
# Move to device for faster quantization
|
||||
loaded_weight = loaded_weight.to(layer._load_device)
|
||||
original_weight_loader(param, qweight, weight_name, shard_id, expert_id)
|
||||
|
||||
if is_w13:
|
||||
param = layer.w13_weight
|
||||
elif is_w2:
|
||||
param = layer.w2_weight
|
||||
|
||||
# In case TP>1, the weight loader logic uses narrow so we cannot directly rely on `param.shape` or `loaded_weight.shape`.
|
||||
copy_numel_counter = CopyNumelCounter()
|
||||
with copy_numel_counter:
|
||||
original_weight_loader(
|
||||
param, loaded_weight, weight_name, shard_id, expert_id
|
||||
)
|
||||
|
||||
if is_w13:
|
||||
layer._w13_loaded_numel += copy_numel_counter.copied_numel
|
||||
target_loaded_numel = layer.w13_weight.numel()
|
||||
current_loaded = layer._w13_loaded_numel
|
||||
elif is_w2:
|
||||
layer._w2_loaded_numel += copy_numel_counter.copied_numel
|
||||
target_loaded_numel = layer.w2_weight.numel()
|
||||
current_loaded = layer._w2_loaded_numel
|
||||
if "w13" in weight_name:
|
||||
scale_param = layer.w13_weight_scale
|
||||
scale_weight_name = "w13_weight_scale"
|
||||
else:
|
||||
raise ValueError("Expected w13 or w2.")
|
||||
# w2.
|
||||
scale_param = layer.w2_weight_scale
|
||||
scale_weight_name = "w2_weight_scale"
|
||||
|
||||
assert (
|
||||
current_loaded <= target_loaded_numel
|
||||
), f"target_loaded_numel={target_loaded_numel}, current_loaded={current_loaded}"
|
||||
|
||||
# Delay online quantization until all tensor shards (e.g. w1 and w3) are loaded, to avoid having to re-quantize later on.
|
||||
if is_w13 and layer._w13_loaded_numel == target_loaded_numel:
|
||||
self._quantize_w13_online(layer, dynamic_mxfp4_quant)
|
||||
elif is_w2 and layer._w2_loaded_numel == target_loaded_numel:
|
||||
self._quantize_w2_online(layer, dynamic_mxfp4_quant)
|
||||
scale_param.weight_loader(
|
||||
scale_param, weight_scale, scale_weight_name, shard_id, expert_id
|
||||
)
|
||||
|
||||
return online_mxfp4_moe_weight_loader
|
||||
|
||||
def get_online_fp8_to_mxfp4_weight_loader(self, layer, original_weight_loader):
|
||||
"""
|
||||
Wrap the original weight loader to perform FP8 to MXFP4 requantization for MoE layers.
|
||||
|
||||
This loader handles:
|
||||
1. Loading FP8 weights (w13_weight, w2_weight) and weight_scale_inv parameters
|
||||
2. Waiting for all experts to be loaded
|
||||
3. Dequantizing FP8 -> BF16 using weight_scale_inv
|
||||
4. Requantizing BF16 -> MXFP4
|
||||
"""
|
||||
|
||||
def online_fp8_to_mxfp4_moe_weight_loader(
|
||||
param: torch.nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
weight_name: str,
|
||||
shard_id: str,
|
||||
expert_id: int,
|
||||
):
|
||||
is_w13_weight = "w13_weight" in weight_name and "scale" not in weight_name
|
||||
is_w2_weight = "w2_weight" in weight_name and "scale" not in weight_name
|
||||
is_w13_scale = "w13_weight_scale_inv" in weight_name
|
||||
is_w2_scale = "w2_weight_scale_inv" in weight_name
|
||||
|
||||
# Sanity multi-threaded load check.
|
||||
assert torch.cuda.current_device() == layer._load_device.index
|
||||
|
||||
# Materialize FP8 parameters on first load from meta device. Adds a small but manageable overhead compared to materializing one by one - but weights are loaded in order layer by layer so it is fine.
|
||||
with layer._fp8_loading_lock:
|
||||
|
||||
if not layer._fp8_materialized:
|
||||
# w13_weight
|
||||
assert layer.w13_weight.device.type == "meta"
|
||||
materialized = torch.nn.Parameter(
|
||||
torch.empty_like(
|
||||
layer.w13_weight.data, device=layer._load_device
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
copy_missing_attrs(layer.w13_weight, materialized)
|
||||
layer.w13_weight = materialized
|
||||
|
||||
# w13_weight_scale_inv
|
||||
materialized = torch.nn.Parameter(
|
||||
torch.empty_like(
|
||||
layer.w13_weight_scale_inv.data, device=layer._load_device
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
copy_missing_attrs(layer.w13_weight_scale_inv, materialized)
|
||||
layer.w13_weight_scale_inv = materialized
|
||||
|
||||
# w2_weight
|
||||
assert layer.w2_weight.device.type == "meta"
|
||||
materialized = torch.nn.Parameter(
|
||||
torch.empty_like(
|
||||
layer.w2_weight.data, device=layer._load_device
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
copy_missing_attrs(layer.w2_weight, materialized)
|
||||
layer.w2_weight = materialized
|
||||
|
||||
# w2_weight_scale_inv
|
||||
assert layer.w2_weight_scale_inv.device.type == "meta"
|
||||
materialized = torch.nn.Parameter(
|
||||
torch.empty_like(
|
||||
layer.w2_weight_scale_inv.data, device=layer._load_device
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
copy_missing_attrs(layer.w2_weight_scale_inv, materialized)
|
||||
layer.w2_weight_scale_inv = materialized
|
||||
|
||||
# Mark as materialized to prevent other threads from doing it again.
|
||||
layer._fp8_materialized = True
|
||||
|
||||
if is_w13_weight:
|
||||
param = layer.w13_weight
|
||||
elif is_w2_weight:
|
||||
param = layer.w2_weight
|
||||
elif is_w13_scale:
|
||||
param = layer.w13_weight_scale_inv
|
||||
elif is_w2_scale:
|
||||
param = layer.w2_weight_scale_inv
|
||||
|
||||
# Track how much data we are actually loading (`narrow` used in weight loader)
|
||||
copy_numel_counter = CopyNumelCounter()
|
||||
with copy_numel_counter:
|
||||
original_weight_loader(
|
||||
param, loaded_weight, weight_name, shard_id, expert_id
|
||||
)
|
||||
|
||||
with layer._fp8_loading_lock:
|
||||
layer._fp8_loaded_numel += copy_numel_counter.copied_numel
|
||||
|
||||
total_target_numel = (
|
||||
layer.w13_weight.numel()
|
||||
+ layer.w2_weight.numel()
|
||||
+ layer.w13_weight_scale_inv.numel()
|
||||
+ layer.w2_weight_scale_inv.numel()
|
||||
)
|
||||
|
||||
# Sanity check
|
||||
assert layer._fp8_loaded_numel <= total_target_numel
|
||||
|
||||
# Perform dequantization and requantization only when all data is loaded AND no other threads are still loading.
|
||||
if layer._fp8_loaded_numel == total_target_numel:
|
||||
if dynamic_mxfp4_quant is None:
|
||||
raise NotImplementedError(
|
||||
"MXFP4 quantization for MoE is only supported on AMD GPUs."
|
||||
)
|
||||
|
||||
assert layer.w13_weight.device.type == "cuda"
|
||||
assert layer.w13_weight_scale_inv.device.type == "cuda"
|
||||
assert layer.w13_weight.dtype != torch.uint8
|
||||
|
||||
# Dequantize and requantize w13
|
||||
w13_bf16 = dequantize_fp8(
|
||||
layer.w13_weight,
|
||||
layer.w13_weight_scale_inv,
|
||||
block_size=self.dequantization_config.weight_block_size,
|
||||
)
|
||||
|
||||
qw13_weight_list = []
|
||||
w13_weight_scale_list = []
|
||||
for expert_idx in range(w13_bf16.shape[0]):
|
||||
# NOTE: dynamic_mxfp4_quant does not accept 3D inputs.
|
||||
qweight, weight_scale = dynamic_mxfp4_quant(
|
||||
w13_bf16[expert_idx]
|
||||
)
|
||||
qw13_weight_list.append(qweight)
|
||||
w13_weight_scale_list.append(weight_scale)
|
||||
|
||||
qw13_weight = torch.stack(qw13_weight_list)
|
||||
w13_weight_scale = torch.stack(w13_weight_scale_list)
|
||||
|
||||
# Dequantize and requantize w2
|
||||
w2_bf16 = dequantize_fp8(
|
||||
layer.w2_weight,
|
||||
layer.w2_weight_scale_inv,
|
||||
block_size=self.dequantization_config.weight_block_size,
|
||||
)
|
||||
|
||||
qw2_weight_list = []
|
||||
w2_weight_scale_list = []
|
||||
for expert_idx in range(w2_bf16.shape[0]):
|
||||
# NOTE: dynamic_mxfp4_quant does not accept 3D inputs.
|
||||
qweight, weight_scale = dynamic_mxfp4_quant(w2_bf16[expert_idx])
|
||||
qw2_weight_list.append(qweight)
|
||||
w2_weight_scale_list.append(weight_scale)
|
||||
|
||||
qw2_weight = torch.stack(qw2_weight_list)
|
||||
w2_weight_scale = torch.stack(w2_weight_scale_list)
|
||||
|
||||
# Replace FP8 parameters with MXFP4 parameters
|
||||
layer.w13_weight = torch.nn.Parameter(
|
||||
qw13_weight, requires_grad=False
|
||||
)
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
w13_weight_scale, requires_grad=False
|
||||
)
|
||||
layer.w2_weight = torch.nn.Parameter(
|
||||
qw2_weight, requires_grad=False
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
w2_weight_scale, requires_grad=False
|
||||
)
|
||||
|
||||
# Clean up FP8 parameters and tracking attributes
|
||||
del layer.w13_weight_scale_inv
|
||||
del layer.w2_weight_scale_inv
|
||||
del layer._fp8_materialized
|
||||
del layer._load_device
|
||||
del layer._fp8_loading_lock
|
||||
|
||||
return online_fp8_to_mxfp4_moe_weight_loader
|
||||
|
||||
def _quantize_w13_online(self, layer, dynamic_mxfp4_quant):
|
||||
qw13_weight = torch.empty(
|
||||
layer.w13_weight.shape[0],
|
||||
layer.w13_weight.shape[1],
|
||||
layer.w13_weight.shape[2] // 2,
|
||||
dtype=torch.uint8,
|
||||
device=layer._load_device,
|
||||
)
|
||||
|
||||
for expert in range(layer.w13_weight.shape[0]):
|
||||
qweight, weight_scale = dynamic_mxfp4_quant(layer.w13_weight.data[expert])
|
||||
assert qw13_weight[expert].shape == qweight.shape
|
||||
assert qw13_weight[expert].dtype == qweight.dtype
|
||||
qw13_weight[expert] = qweight
|
||||
|
||||
assert layer.w13_weight_scale[expert].shape == weight_scale.shape
|
||||
assert layer.w13_weight_scale[expert].dtype == weight_scale.dtype
|
||||
layer.w13_weight_scale[expert] = weight_scale
|
||||
|
||||
layer.w13_weight = torch.nn.Parameter(qw13_weight, requires_grad=False)
|
||||
|
||||
def _quantize_w2_online(self, layer, dynamic_mxfp4_quant):
|
||||
qw2_weight = torch.empty(
|
||||
layer.w2_weight.shape[0],
|
||||
layer.w2_weight.shape[1],
|
||||
layer.w2_weight.shape[2] // 2,
|
||||
dtype=torch.uint8,
|
||||
device=layer._load_device,
|
||||
)
|
||||
|
||||
for expert in range(layer.w2_weight.shape[0]):
|
||||
qweight, weight_scale = dynamic_mxfp4_quant(layer.w2_weight.data[expert])
|
||||
qw2_weight[expert] = qweight
|
||||
layer.w2_weight_scale[expert] = weight_scale
|
||||
|
||||
layer.w2_weight = torch.nn.Parameter(qw2_weight, requires_grad=False)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
if (
|
||||
not self.is_checkpoint_mxfp4_serialized
|
||||
or self.dequantization_config is not None
|
||||
):
|
||||
# Quantization already happened during weight loading.
|
||||
# This covers both:
|
||||
# - Online quantization from BF16/FP16 -> MXFP4
|
||||
# - Requantization from FP8 -> MXFP4
|
||||
assert layer.w13_weight.dtype == torch.uint8
|
||||
assert layer.w2_weight.dtype == torch.uint8
|
||||
assert layer.w13_weight_scale.dtype == torch.uint8
|
||||
assert layer.w2_weight_scale.dtype == torch.uint8
|
||||
|
||||
# Pre-shuffle weight scales
|
||||
s0, s1, _ = layer.w13_weight_scale.shape
|
||||
w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1)
|
||||
|
||||
@@ -210,9 +210,5 @@ def quark_post_load_weights(self_attn: nn.Module, w: torch.Tensor, quant_format:
|
||||
w_vc, w_s_vc = b_dynamic_mxfp4_quant(w_vc)
|
||||
w_s_kc = w_s_kc.transpose(1, 2).contiguous().transpose(1, 2)
|
||||
w_s_vc = w_s_vc.contiguous().transpose(1, 2)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unexpected w.dtype: {w.dtype} (should be bfloat16 or uint8)"
|
||||
)
|
||||
|
||||
return w_kc, w_s_kc, w_vc, w_s_vc
|
||||
|
||||
@@ -1508,16 +1508,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
getattr(self.model, "quant_config", None), "quantized_layers", None
|
||||
)
|
||||
if (
|
||||
hasattr(self.model, "quant_config")
|
||||
and hasattr(self.model.quant_config, "quantized_layers")
|
||||
and self.server_args.quantization is not None
|
||||
self.server_args.quantization is not None
|
||||
and isinstance(quantized_layers, tuple)
|
||||
and len(quantized_layers) == 2
|
||||
):
|
||||
type_counts, quantized_layers_count = (
|
||||
self.model.quant_config.quantized_layers
|
||||
)
|
||||
type_summary = ", ".join(f"{t}: {c}" for t, c in type_counts.items())
|
||||
layer_types, quantized_layers_count = quantized_layers
|
||||
logger.info(
|
||||
f"Online {self.server_args.quantization} quantization: quantized {quantized_layers_count} layers in total ({type_summary})."
|
||||
f"Online {self.server_args.quantization} quantization: quantized {quantized_layers_count} layers of types: {layer_types}"
|
||||
)
|
||||
|
||||
if self.server_args.debug_tensor_dump_output_folder is not None:
|
||||
|
||||
@@ -321,20 +321,6 @@ def maybe_executor_submit(
|
||||
if func_kwargs is None:
|
||||
func_kwargs = {}
|
||||
if use_async:
|
||||
# CRITICAL: Capture current CUDA device and restore it in worker thread.
|
||||
# torch.cuda.current_device() is thread-local and is NOT correctly passed to threads.
|
||||
# This may result in errors in case the `func` relies on torch current device to be already correctly specified.
|
||||
# See details in https://github.com/pytorch/pytorch/issues/56588.
|
||||
current_device = (
|
||||
torch.cuda.current_device() if torch.cuda.is_available() else None
|
||||
)
|
||||
|
||||
def device_aware_wrapper(*args, **kwargs):
|
||||
# Set CUDA device in worker thread to match parent thread
|
||||
if current_device is not None:
|
||||
torch.cuda.set_device(current_device)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
futures.append(executor.submit(device_aware_wrapper, *func_args, **func_kwargs))
|
||||
futures.append(executor.submit(func, *func_args, **func_kwargs))
|
||||
else:
|
||||
func(*func_args, **func_kwargs)
|
||||
|
||||
@@ -40,7 +40,7 @@ from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from sglang.srt.configs.load_config import LoadConfig
|
||||
from sglang.srt.configs.model_config import REQUANTIZATION_METHODS, ModelConfig
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
@@ -261,12 +261,6 @@ def get_quant_config(
|
||||
if not isinstance(hf_quant_config, dict):
|
||||
hf_quant_config = hf_quant_config.to_dict()
|
||||
hf_quant_config["packed_modules_mapping"] = packed_modules_mapping
|
||||
hf_quant_config["hf_config"] = model_config.hf_config
|
||||
|
||||
# This is only used by quantization methods that support requantization (e.g. from fp8 to mxfp4).
|
||||
if model_config.quantization in REQUANTIZATION_METHODS:
|
||||
hf_quant_config["requantization_method"] = model_config.quantization
|
||||
|
||||
return quant_cls.from_config(hf_quant_config)
|
||||
|
||||
# In case of bitsandbytes/QLoRA, get quant config from the adapter model.
|
||||
|
||||
Reference in New Issue
Block a user