wehub-resource-sync
830 lines
32 KiB
Python
830 lines
32 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import json
|
|
import math
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import regex as re
|
|
import torch
|
|
|
|
import vllm.utils.humming as _hm
|
|
from vllm import envs
|
|
from vllm.model_executor.layers.fused_moe import (
|
|
FusedMoEConfig,
|
|
FusedMoEMethodBase,
|
|
FusedMoEQuantConfig,
|
|
RoutedExperts,
|
|
SharedExperts,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
|
|
UnquantizedFusedMoEMethod,
|
|
)
|
|
from vllm.model_executor.layers.linear import (
|
|
LinearBase,
|
|
LinearMethodBase,
|
|
UnquantizedLinearMethod,
|
|
)
|
|
from vllm.model_executor.layers.quantization import QuantizationMethods
|
|
from vllm.model_executor.layers.quantization.base_config import (
|
|
QuantizationConfig,
|
|
QuantizeMethodBase,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.humming_utils import (
|
|
convert_to_humming_moe_kernel_format,
|
|
get_humming_moe_quant_config,
|
|
input_schema_to_quant_key,
|
|
make_humming_moe_kernel,
|
|
select_humming_moe_experts,
|
|
weight_schema_to_quant_key,
|
|
)
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.parameter import (
|
|
BasevLLMParameter,
|
|
BlockQuantScaleParameter,
|
|
ChannelQuantScaleParameter,
|
|
GroupQuantScaleParameter,
|
|
ModelWeightParameter,
|
|
PackedvLLMParameter,
|
|
PerTensorScaleParameter,
|
|
RowvLLMParameter,
|
|
)
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.model_executor.models.utils import WeightsMapper
|
|
from vllm.utils.humming import (
|
|
BaseInputSchema,
|
|
BaseWeightSchema,
|
|
HummingInputSchema,
|
|
HummingWeightSchema,
|
|
)
|
|
|
|
|
|
def prepare_padded_shape(shape, x):
|
|
padded_shape = math.ceil(shape / x) * x
|
|
return padded_shape, padded_shape - shape
|
|
|
|
|
|
def prepare_param(tensor, name, extra_attrs):
|
|
extra_attrs = extra_attrs.copy()
|
|
scale_type = extra_attrs.pop("scale_type", None)
|
|
param_cls_name_map = {
|
|
"block": BlockQuantScaleParameter,
|
|
"tensor": PerTensorScaleParameter,
|
|
"group": GroupQuantScaleParameter,
|
|
"channel": ChannelQuantScaleParameter,
|
|
"input_scale": PerTensorScaleParameter,
|
|
}
|
|
|
|
param_cls: type[BasevLLMParameter]
|
|
if "packed_dim" in extra_attrs:
|
|
param_cls = PackedvLLMParameter
|
|
elif scale_type in param_cls_name_map:
|
|
param_cls = param_cls_name_map[scale_type]
|
|
elif "output_dim" in extra_attrs and "input_dim" in extra_attrs:
|
|
param_cls = ModelWeightParameter
|
|
elif "input_dim" in extra_attrs:
|
|
param_cls = RowvLLMParameter
|
|
elif "output_dim" in extra_attrs:
|
|
param_cls = ChannelQuantScaleParameter
|
|
else:
|
|
param_cls = BasevLLMParameter
|
|
|
|
kwargs_keys = [
|
|
"input_dim",
|
|
"output_dim",
|
|
"packed_dim",
|
|
"packed_factor",
|
|
"weight_loader",
|
|
]
|
|
cls_kwargs = {}
|
|
for key in extra_attrs.copy():
|
|
if key in kwargs_keys:
|
|
cls_kwargs[key] = extra_attrs.pop(key)
|
|
|
|
param = param_cls(data=tensor, **cls_kwargs)
|
|
set_weight_attrs(param, extra_attrs)
|
|
|
|
param.param_name = name
|
|
param.ignore_warning = True
|
|
if scale_type in ["tensor", "input_scale"]:
|
|
param.needs_scalar_to_array = True
|
|
|
|
return param
|
|
|
|
|
|
def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, Any]):
|
|
param = torch.nn.Parameter(tensor, requires_grad=False)
|
|
if "scale_type" in extra_attrs:
|
|
extra_attrs["quant_method"] = extra_attrs["scale_type"]
|
|
|
|
if "input_dim" in extra_attrs and "output_dim" in extra_attrs:
|
|
input_dim = extra_attrs["input_dim"]
|
|
output_dim = extra_attrs["output_dim"]
|
|
extra_attrs["is_transposed"] = input_dim < output_dim
|
|
|
|
set_weight_attrs(param, extra_attrs)
|
|
param.param_name = name
|
|
return param
|
|
|
|
|
|
def may_pad_loaded_weight(param, loaded_weight):
|
|
pad_shape = getattr(param, "pad_shape", None)
|
|
if pad_shape is None:
|
|
return loaded_weight
|
|
value = 1 if loaded_weight.dtype == torch.float8_e8m0fnu else 0
|
|
padding = []
|
|
for x in pad_shape[::-1][: loaded_weight.ndim]:
|
|
padding += [0, x]
|
|
loaded_weight = torch.nn.functional.pad(
|
|
input=loaded_weight,
|
|
pad=padding,
|
|
value=value,
|
|
)
|
|
return loaded_weight
|
|
|
|
|
|
def compressed_tensors_get_config(config: dict[str, Any], key: str):
|
|
assert key in ["weights", "input_activations"]
|
|
target_group_config = None
|
|
for group_config in config["config_groups"].values():
|
|
if "Linear" in group_config["targets"]:
|
|
if "weights" not in group_config:
|
|
return None
|
|
if key not in group_config or group_config[key] is None:
|
|
return None
|
|
target_group_config = group_config[key].copy()
|
|
break
|
|
|
|
if target_group_config is None:
|
|
return None
|
|
target_group_config["quant_method"] = config["quant_method"]
|
|
if config["quant_method"] == "compressed-tensors":
|
|
target_group_config["format"] = config["format"]
|
|
elif config["quant_method"] == "modelopt":
|
|
target_group_config["quant_algo"] = config["quant_algo"]
|
|
return target_group_config
|
|
|
|
|
|
class HummingConfig(QuantizationConfig):
|
|
packed_modules_mapping: dict[str, list[str]] = {}
|
|
|
|
def __init__(self, full_config: dict[str, Any] | None = None):
|
|
self.full_config: dict[str, Any] = full_config or {}
|
|
|
|
@classmethod
|
|
def get_name(cls) -> QuantizationMethods:
|
|
return "humming"
|
|
|
|
@classmethod
|
|
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
|
return [torch.bfloat16, torch.half]
|
|
|
|
@classmethod
|
|
def get_min_capability(cls) -> int:
|
|
return 75
|
|
|
|
@classmethod
|
|
def get_config_filenames(cls) -> list[str]:
|
|
return []
|
|
|
|
@classmethod
|
|
def from_config(cls, config: dict[str, Any]) -> "HummingConfig":
|
|
return cls(full_config=config)
|
|
|
|
@classmethod
|
|
def override_quantization_method(
|
|
cls, hf_quant_cfg, user_quant, hf_config=None
|
|
) -> QuantizationMethods | None:
|
|
if user_quant == "humming" and hf_config is not None:
|
|
model_type = hf_config.model_type
|
|
quant_method = hf_quant_cfg.get("quant_method", None)
|
|
if model_type == "gpt_oss" and quant_method == "mxfp4":
|
|
msg = (
|
|
"For gpt-oss model, use '--moe-backend humming' "
|
|
"instead of '--quantization humming'."
|
|
)
|
|
raise ValueError(msg)
|
|
return "humming" if user_quant == "humming" else None
|
|
|
|
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
|
|
self.hf_to_vllm_mapper = hf_to_vllm_mapper
|
|
|
|
def is_layer_skipped(self, config: dict[str, Any], prefix: str):
|
|
keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
|
|
ignored_layers = self.get_from_keys_or(config, keys, []) or []
|
|
if hasattr(self, "hf_to_vllm_mapper"):
|
|
ignored_layers = self.hf_to_vllm_mapper.apply_list(ignored_layers)
|
|
|
|
if any(module_name in prefix for module_name in ignored_layers):
|
|
return True
|
|
if "lm_head" in prefix:
|
|
return True
|
|
|
|
for regex in config.get("dynamic", {}):
|
|
if regex[:1] != "-":
|
|
continue
|
|
if re.match(regex[2:], prefix):
|
|
return True
|
|
|
|
return False
|
|
|
|
def get_layer_weight_schema(self, config: dict[str, Any], prefix: str):
|
|
if self.is_layer_skipped(config, prefix):
|
|
return None
|
|
|
|
if config["quant_method"] in ["compressed-tensors", "modelopt"]:
|
|
group_config = compressed_tensors_get_config(config, "weights")
|
|
if group_config is None:
|
|
return None
|
|
config = group_config
|
|
|
|
layer_config = config
|
|
layer_dynamic = config.get("dynamic", {})
|
|
if not isinstance(layer_dynamic, dict):
|
|
layer_dynamic = {}
|
|
for regex, override_config in layer_dynamic.items():
|
|
if regex[:1] != "+":
|
|
continue
|
|
if re.match(regex[2:], prefix):
|
|
layer_config = config.copy()
|
|
layer_config.update(override_config)
|
|
break
|
|
|
|
if "quant_method" in layer_config:
|
|
return _hm.BaseWeightSchema.from_config(layer_config)
|
|
return None
|
|
|
|
def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
|
|
if self.is_layer_skipped(config, prefix):
|
|
return None
|
|
if config["quant_method"] in ["compressed-tensors", "modelopt"]:
|
|
group_config = compressed_tensors_get_config(config, "input_activations")
|
|
if group_config is None:
|
|
return None
|
|
config = group_config
|
|
|
|
if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP:
|
|
return _hm.BaseInputSchema.from_config(config)
|
|
return None
|
|
|
|
def get_quant_config_for_layer(
|
|
self, prefix: str, layer_type: str
|
|
) -> "HummingLayerQuantizationConfig | None":
|
|
weight_schema: BaseWeightSchema | None = None
|
|
force_weight_schema: HummingWeightSchema | None = None
|
|
|
|
if self.full_config:
|
|
weight_schema = self.get_layer_weight_schema(self.full_config, prefix)
|
|
|
|
is_online_quant = False
|
|
online_quant_config = envs.VLLM_HUMMING_ONLINE_QUANT_CONFIG or {}
|
|
if not self.full_config or online_quant_config.get("force_requant", False):
|
|
online_quant_config["quant_method"] = "humming"
|
|
schema = self.get_layer_weight_schema(online_quant_config, prefix)
|
|
if not self.full_config:
|
|
weight_schema = schema
|
|
is_online_quant = True
|
|
else:
|
|
force_weight_schema = schema
|
|
|
|
if weight_schema is not None:
|
|
input_schema = None
|
|
force_input_schema = None
|
|
|
|
if self.full_config:
|
|
input_schema = self.get_layer_input_schema(self.full_config, prefix)
|
|
|
|
if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
|
|
quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy()
|
|
quant_config["quant_method"] = "humming"
|
|
force_input_schema = self.get_layer_input_schema(quant_config, prefix)
|
|
if input_schema is None:
|
|
input_schema = force_input_schema
|
|
|
|
if force_weight_schema is not None and force_input_schema is None:
|
|
force_input_schema = _hm.HummingInputSchema()
|
|
|
|
return HummingLayerQuantizationConfig(
|
|
weight_schema=weight_schema,
|
|
input_schema=input_schema,
|
|
force_weight_schema=force_weight_schema,
|
|
force_input_schema=force_input_schema,
|
|
is_online_quant=is_online_quant,
|
|
)
|
|
return None
|
|
|
|
def get_quant_method(
|
|
self, layer: torch.nn.Module, prefix: str
|
|
) -> "QuantizeMethodBase | None":
|
|
layer_type = "other"
|
|
if isinstance(layer, RoutedExperts):
|
|
layer_type = "moe"
|
|
elif isinstance(layer, LinearBase):
|
|
layer_type = "linear"
|
|
|
|
quant_config = self.get_quant_config_for_layer(prefix, layer_type)
|
|
if quant_config is None:
|
|
if isinstance(layer, RoutedExperts):
|
|
return UnquantizedFusedMoEMethod(layer.moe_config)
|
|
elif isinstance(layer, LinearBase):
|
|
return UnquantizedLinearMethod()
|
|
elif isinstance(layer, LinearBase):
|
|
return HummingLinearMethod(quant_config)
|
|
elif isinstance(layer, RoutedExperts):
|
|
return HummingMoEMethod(quant_config, layer.moe_config)
|
|
return None
|
|
|
|
|
|
class HummingLayerQuantizationConfig(HummingConfig):
|
|
def __init__(
|
|
self,
|
|
weight_schema: "BaseWeightSchema",
|
|
input_schema: "BaseInputSchema | None" = None,
|
|
force_weight_schema: "HummingWeightSchema | None" = None,
|
|
force_input_schema: "HummingInputSchema | None" = None,
|
|
is_online_quant: bool = False,
|
|
):
|
|
self.weight_schema = weight_schema
|
|
if input_schema is None:
|
|
input_schema = _hm.HummingInputSchema()
|
|
self.input_schema = input_schema
|
|
self.force_weight_schema = force_weight_schema
|
|
self.force_input_schema = force_input_schema
|
|
self.is_online_quant = is_online_quant
|
|
|
|
@classmethod
|
|
def from_config(cls, config):
|
|
weight_schema = _hm.BaseWeightSchema.from_config(config)
|
|
return cls(weight_schema)
|
|
|
|
def get_quant_method(
|
|
self, layer: torch.nn.Module, prefix: str
|
|
) -> QuantizeMethodBase | None:
|
|
raise NotImplementedError
|
|
|
|
|
|
class HummingLinearMethod(LinearMethodBase):
|
|
def __init__(self, quant_config: HummingLayerQuantizationConfig):
|
|
self.quant_config = quant_config
|
|
self.weight_schema = quant_config.weight_schema
|
|
self.input_schema = quant_config.input_schema
|
|
self.force_weight_schema = quant_config.force_weight_schema
|
|
self.force_input_schema = quant_config.force_input_schema
|
|
self.is_online_quant = self.quant_config.is_online_quant
|
|
|
|
def prepare_weight_loader(self, layer: torch.nn.Module, weight_loader: Callable):
|
|
def new_weight_loader(
|
|
param: torch.nn.Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
shard_id: str | int | None = None,
|
|
):
|
|
name = param.param_name
|
|
float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
|
is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
|
|
if is_unquantized and self.is_online_quant:
|
|
# online quant (fp16/bf16 -> quant_type)
|
|
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
|
f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
|
|
has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
|
|
tensor_list = _hm.quantize_weight(
|
|
weight=loaded_weight,
|
|
dtype=self.weight_schema.b_dtype,
|
|
scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
|
|
group_size=self.weight_schema.weight_scale_group_size,
|
|
has_zero_point=self.weight_schema.has_zero_point,
|
|
has_global_scale=has_global_scale,
|
|
is_fp_zero_point=self.weight_schema.is_fp_zero_point,
|
|
pack=True,
|
|
)
|
|
|
|
key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
|
|
for key, tensor in zip(key_list, tensor_list):
|
|
if tensor is None or tensor.nelement() == 0:
|
|
continue
|
|
param = getattr(layer, key)
|
|
param.weight_loader(param, tensor, shard_id)
|
|
|
|
return None
|
|
elif is_unquantized and not self.is_online_quant:
|
|
# fallback to unquantized linear
|
|
# some model skip some layer when quantizing model, but
|
|
# don't mark the layer as unquantized.
|
|
if not layer.is_fallback:
|
|
layer.is_fallback = True
|
|
for name, _ in list(layer.named_parameters()):
|
|
if name != "bias":
|
|
delattr(layer, name)
|
|
delattr(layer, "locks")
|
|
self.__class__ = UnquantizedLinearMethod # type: ignore
|
|
tensor = torch.empty(
|
|
(
|
|
layer.output_partition_sizes_sum,
|
|
layer.input_size_per_partition,
|
|
),
|
|
dtype=layer.param_dtype,
|
|
device=param.device,
|
|
)
|
|
extra_weight_attrs = layer.extra_weight_attrs.copy()
|
|
orig_weight_loader = extra_weight_attrs.pop("weight_loader")
|
|
layer.weight = ModelWeightParameter(
|
|
data=tensor,
|
|
input_dim=1,
|
|
output_dim=0,
|
|
weight_loader=orig_weight_loader,
|
|
)
|
|
layer.weight.tp_size = layer.tp_size
|
|
layer.weight.tp_rank = layer.tp_rank
|
|
set_weight_attrs(layer.weight, extra_weight_attrs)
|
|
|
|
param = layer.weight
|
|
if shard_id is not None:
|
|
return layer.weight.weight_loader(param, loaded_weight, shard_id)
|
|
return layer.weight.weight_loader(param, loaded_weight)
|
|
|
|
# weight processing logic for specific quantization schema
|
|
loaded_weight = self.weight_schema.process_loaded_weight(
|
|
tensor=loaded_weight,
|
|
name=name,
|
|
)
|
|
if shard_id is not None:
|
|
return weight_loader(param, loaded_weight, shard_id)
|
|
return weight_loader(param, loaded_weight)
|
|
|
|
return new_weight_loader
|
|
|
|
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,
|
|
):
|
|
layer.is_fallback = False
|
|
layer.param_dtype = params_dtype
|
|
layer.input_size = input_size
|
|
layer.output_size = output_size
|
|
layer.input_size_per_partition = input_size_per_partition
|
|
layer.output_partition_sizes_sum = sum(output_partition_sizes)
|
|
layer.output_partition_sizes = output_partition_sizes
|
|
layer.extra_weight_attrs = extra_weight_attrs.copy()
|
|
|
|
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
|
|
new_weight_loader = self.prepare_weight_loader(layer, weight_loader)
|
|
extra_weight_attrs["weight_loader"] = new_weight_loader
|
|
|
|
for key in ["weight_block_size", "block_structure"]:
|
|
block_size = getattr(self.weight_schema, key, None)
|
|
if block_size is not None:
|
|
layer.weight_block_size = block_size
|
|
|
|
weight_tensor_attrs = self.weight_schema.get_tensors_attrs(
|
|
shape_n=layer.output_partition_sizes_sum,
|
|
shape_k=layer.input_size_per_partition,
|
|
param_dtype=params_dtype,
|
|
stack_size=len(layer.output_partition_sizes),
|
|
)
|
|
|
|
input_tensor_attrs = self.input_schema.get_tensors_attrs(
|
|
shape_k=layer.input_size_per_partition,
|
|
param_dtype=params_dtype,
|
|
stack_size=len(layer.output_partition_sizes),
|
|
)
|
|
|
|
tensors_attrs = weight_tensor_attrs | input_tensor_attrs
|
|
|
|
for name, attrs in tensors_attrs.items():
|
|
tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
|
|
extra_attrs = attrs.get("extra_attrs", {}).copy()
|
|
extra_attrs.update(extra_weight_attrs)
|
|
param = prepare_param(tensor, name, extra_attrs)
|
|
setattr(layer, name, param)
|
|
|
|
locks = torch.zeros(1024, dtype=torch.int32)
|
|
layer.register_buffer("locks", locks)
|
|
|
|
if self.force_input_schema is not None:
|
|
self.input_schema = self.force_input_schema
|
|
|
|
if not hasattr(layer, "weight"):
|
|
param = prepare_param(torch.tensor(0), "weight", extra_weight_attrs)
|
|
layer.weight = param
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
if layer.is_fallback:
|
|
return None
|
|
|
|
# convert from checkpoint format to humming format
|
|
if not isinstance(self.weight_schema, _hm.HummingWeightSchema):
|
|
self.weight_schema, tensors = self.weight_schema.convert_humming(
|
|
tensors=layer.state_dict(),
|
|
shape_n_stacks=layer.output_partition_sizes,
|
|
shape_k_stacks=[layer.input_size_per_partition],
|
|
param_dtype=layer.param_dtype,
|
|
)
|
|
|
|
self.input_schema, _ = self.input_schema.convert_humming(
|
|
tensors=layer.state_dict(),
|
|
shape_n_stacks=layer.output_partition_sizes,
|
|
shape_k_stacks=[layer.input_size_per_partition],
|
|
param_dtype=layer.param_dtype,
|
|
)
|
|
|
|
for name, _ in list(layer.named_parameters()):
|
|
delattr(layer, name)
|
|
|
|
for name, tensor in tensors.items():
|
|
param = torch.nn.Parameter(tensor, requires_grad=False)
|
|
setattr(layer, name, param)
|
|
|
|
del tensors
|
|
|
|
# force requant (origin quant setting -> fp16/bf16 -> new_quant setting)
|
|
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
|
force_requant = self.force_weight_schema is not None
|
|
if force_requant and self.weight_schema != self.force_weight_schema:
|
|
tensors = self.weight_schema.requant_tensors(
|
|
tensors=layer.state_dict(),
|
|
target_weight_schema=self.force_weight_schema,
|
|
param_dtype=layer.param_dtype,
|
|
)
|
|
|
|
self.weight_schema = self.force_weight_schema
|
|
|
|
for name, _ in list(layer.named_parameters()):
|
|
if name != "bias":
|
|
delattr(layer, name)
|
|
|
|
for name, tensor in tensors.items():
|
|
param = torch.nn.Parameter(tensor, requires_grad=False)
|
|
setattr(layer, name, param)
|
|
|
|
del tensors
|
|
|
|
# prepare layer config from humming kernel
|
|
_hm.HummingMethod.prepare_layer_meta(
|
|
layer=layer,
|
|
shape_n=layer.output_partition_sizes_sum,
|
|
shape_k=layer.input_size_per_partition,
|
|
weight_schema=self.weight_schema,
|
|
input_schema=self.input_schema,
|
|
pad_n_to_multiple=256,
|
|
pad_k_to_multiple=128,
|
|
has_bias=layer.has_bias,
|
|
torch_dtype=layer.param_dtype,
|
|
)
|
|
|
|
# preprocess weight for inference
|
|
_hm.HummingMethod.transform_humming_layer(layer)
|
|
|
|
# compute_config: kernel configs that do not directly affect weights
|
|
# but significantly impact kernel behavior or computation precision.
|
|
# see https://github.com/inclusionAI/humming/blob/main/docs/config.md
|
|
compute_config = {
|
|
"use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
|
|
"use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
|
|
"gemm_type": "dense",
|
|
}
|
|
self.compute_config = json.dumps(compute_config)
|
|
|
|
def apply(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
flatten_inputs = x.view(-1, x.size(-1))
|
|
output = _hm.HummingMethod.forward_layer(
|
|
layer=layer,
|
|
inputs=flatten_inputs,
|
|
compute_config=self.compute_config,
|
|
)
|
|
output = output.view(*x.shape[:-1], output.size(-1))
|
|
return output
|
|
|
|
|
|
class HummingMoEMethod(FusedMoEMethodBase):
|
|
def __init__(
|
|
self, quant_config: HummingLayerQuantizationConfig, moe: "FusedMoEConfig"
|
|
) -> None:
|
|
super().__init__(moe)
|
|
self.quant_config = quant_config
|
|
self.weight_schema = quant_config.weight_schema
|
|
self.input_schema = quant_config.input_schema
|
|
self.force_weight_schema = quant_config.force_weight_schema
|
|
self.force_input_schema = quant_config.force_input_schema
|
|
|
|
# Derive QuantKeys from humming schemas.
|
|
# Prefer force schemas (the final format after requant) over base.
|
|
weight_key = weight_schema_to_quant_key(
|
|
self.force_weight_schema or self.weight_schema
|
|
)
|
|
activation_key = input_schema_to_quant_key(
|
|
self.force_input_schema or self.input_schema
|
|
)
|
|
|
|
# Select Humming MoE experts
|
|
self.experts_cls = select_humming_moe_experts(
|
|
config=self.moe,
|
|
weight_key=weight_key,
|
|
activation_key=activation_key,
|
|
)
|
|
|
|
def prepare_weight_loader(self, layer, weight_loader):
|
|
def new_weight_loader(
|
|
param: torch.nn.Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
weight_name: str,
|
|
shard_id: str,
|
|
expert_id: int | None = None,
|
|
return_success: bool = False,
|
|
):
|
|
name = param.param_name
|
|
float_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
|
is_unquantized = name == "weight" and loaded_weight.dtype in float_dtypes
|
|
# online quant (fp16/bf16 -> quant_type)
|
|
if is_unquantized:
|
|
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
|
|
f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
|
|
has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
|
|
tensor_list = _hm.quantize_weight(
|
|
weight=loaded_weight,
|
|
dtype=self.weight_schema.b_dtype,
|
|
scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
|
|
group_size=self.weight_schema.weight_scale_group_size,
|
|
has_zero_point=self.weight_schema.has_zero_point,
|
|
has_global_scale=has_global_scale,
|
|
is_fp_zero_point=self.weight_schema.is_fp_zero_point,
|
|
pack=True,
|
|
)
|
|
|
|
key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
|
|
success = True
|
|
for key, tensor in zip(key_list, tensor_list):
|
|
if tensor is None or tensor.nelement() == 0:
|
|
continue
|
|
sublayer_name = "w2" if shard_id == "w2" else "w13"
|
|
|
|
param = getattr(layer, sublayer_name + "_" + key)
|
|
part_success = param.weight_loader(
|
|
param=param,
|
|
loaded_weight=tensor.cpu(),
|
|
weight_name=shard_id + "_" + key,
|
|
shard_id=shard_id,
|
|
expert_id=expert_id,
|
|
return_success=return_success,
|
|
)
|
|
success = success and part_success
|
|
|
|
return success if return_success else None
|
|
|
|
# weight processing logic for specific quantization schema
|
|
loaded_weight = self.weight_schema.process_loaded_weight(
|
|
tensor=loaded_weight,
|
|
name=name,
|
|
)
|
|
return weight_loader(
|
|
param,
|
|
loaded_weight,
|
|
weight_name,
|
|
shard_id=shard_id,
|
|
expert_id=expert_id,
|
|
return_success=return_success,
|
|
)
|
|
|
|
return new_weight_loader
|
|
|
|
def create_weights(
|
|
self,
|
|
layer: RoutedExperts,
|
|
num_experts: int,
|
|
hidden_size: int,
|
|
intermediate_size_per_partition: int,
|
|
params_dtype: torch.dtype,
|
|
**extra_weight_attrs,
|
|
):
|
|
layer.num_experts = num_experts
|
|
layer.param_dtype = params_dtype
|
|
layer.intermediate_size = intermediate_size_per_partition
|
|
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
|
|
weight_loader = self.prepare_weight_loader(layer, weight_loader)
|
|
extra_weight_attrs["weight_loader"] = weight_loader
|
|
|
|
# sublayer: a layer contains multiple sets of weights for quantized GEMM
|
|
# (e.g., weight, weight_scale, etc.).
|
|
# The weight names of sublayer start with the prefix "{sublayer_name}_"
|
|
layer.sublayer_configs = {
|
|
"w13": {
|
|
"shape_n": intermediate_size_per_partition * 2,
|
|
"shape_k": hidden_size,
|
|
"tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
|
|
shape_n=intermediate_size_per_partition * 2,
|
|
shape_k=hidden_size,
|
|
num_experts=num_experts,
|
|
param_dtype=params_dtype,
|
|
has_bias=self.moe.has_bias,
|
|
),
|
|
},
|
|
"w2": {
|
|
"shape_n": hidden_size,
|
|
"shape_k": intermediate_size_per_partition,
|
|
"tensors_attrs": self.weight_schema.get_padded_tensors_attrs(
|
|
shape_n=hidden_size,
|
|
shape_k=intermediate_size_per_partition,
|
|
num_experts=num_experts,
|
|
param_dtype=params_dtype,
|
|
has_bias=self.moe.has_bias,
|
|
),
|
|
},
|
|
}
|
|
|
|
for sublayer_name, configs in layer.sublayer_configs.items():
|
|
for name, attrs in configs["tensors_attrs"].items():
|
|
tensor = torch.empty(attrs["shape"], dtype=attrs["dtype"])
|
|
param = torch.nn.Parameter(tensor, requires_grad=False)
|
|
extra_attrs = attrs.get("extra_attrs", {}).copy()
|
|
extra_attrs.update(extra_weight_attrs)
|
|
param = prepare_moe_param(tensor, name, extra_attrs)
|
|
setattr(layer, f"{sublayer_name}_{name}", param)
|
|
|
|
if self.force_input_schema is not None:
|
|
self.input_schema = self.force_input_schema
|
|
|
|
locks = torch.zeros(1024, dtype=torch.int32)
|
|
layer.register_buffer("locks", locks)
|
|
|
|
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
|
|
return get_humming_moe_quant_config(layer)
|
|
|
|
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
|
|
if getattr(self, "processed", False):
|
|
return
|
|
self.processed = True
|
|
|
|
# Convert weights to Humming kernel format
|
|
convert_to_humming_moe_kernel_format(
|
|
layer=layer,
|
|
sublayer_configs=layer.sublayer_configs,
|
|
weight_schema=self.weight_schema,
|
|
input_schema=self.input_schema,
|
|
force_weight_schema=self.force_weight_schema,
|
|
)
|
|
|
|
# Build the MoE kernel
|
|
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
|
|
assert self.moe_quant_config is not None
|
|
assert self.experts_cls is not None
|
|
self.moe_kernel = make_humming_moe_kernel(
|
|
self.moe_quant_config,
|
|
self.moe,
|
|
self.experts_cls,
|
|
layer=layer,
|
|
routing_tables=layer._expert_routing_tables(),
|
|
)
|
|
|
|
def apply(
|
|
self,
|
|
layer: RoutedExperts,
|
|
x: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
shared_experts: SharedExperts | None,
|
|
shared_experts_input: torch.Tensor | None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Apply Humming-quantized MoE computation using the standard kernel flow.
|
|
|
|
This method uses FusedMoEKernel.apply() which orchestrates:
|
|
1. Preparation (quantization if needed - skipped for Humming via
|
|
expects_unquantized_inputs=True to prevent double quantization)
|
|
2. Expert computation (via experts.apply())
|
|
3. Finalization (weight application & reduction - no-op for Humming
|
|
since it's already done internally)
|
|
|
|
Humming handles all quantization, weight application, and reduction
|
|
internally in the experts.apply() method via HummingMethod calls.
|
|
|
|
Note: Although w1/w2 weights are passed to the kernel for interface
|
|
consistency, Humming's experts.apply() reads weights directly from
|
|
the layer object via HummingMethod.forward_layer() and ignores the
|
|
w1/w2 parameters.
|
|
"""
|
|
assert self.moe_kernel is not None
|
|
return self.moe_kernel.apply(
|
|
hidden_states=x,
|
|
w1=layer.w13_weight,
|
|
w2=layer.w2_weight,
|
|
topk_ids=topk_ids,
|
|
topk_weights=topk_weights,
|
|
activation=layer.activation,
|
|
global_num_experts=layer.global_num_experts,
|
|
expert_map=layer.expert_map,
|
|
apply_router_weight_on_input=False,
|
|
shared_experts=shared_experts,
|
|
shared_experts_input=shared_experts_input,
|
|
)
|