779 lines
37 KiB
Python
779 lines
37 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch_npu
|
|
from vllm.config import get_current_vllm_config
|
|
from vllm.distributed import get_tensor_model_parallel_world_size
|
|
|
|
from vllm_ascend.ascend_config import get_ascend_config
|
|
from vllm_ascend.ascend_forward_context import _EXTRA_CTX
|
|
from vllm_ascend.distributed.parallel_state import get_mc2_group
|
|
from vllm_ascend.ops.fused_moe.experts_selector import select_experts
|
|
from vllm_ascend.ops.fused_moe.moe_runtime_args import build_fused_experts_input
|
|
from vllm_ascend.utils import COMPRESSED_TENSORS_METHOD, maybe_trans_nz
|
|
|
|
from .base import AscendLinearScheme, AscendMoEScheme, QuantType, get_moe_num_logical_experts
|
|
from .registry import register_scheme
|
|
|
|
|
|
@register_scheme("W4A8_DYNAMIC", "linear")
|
|
class AscendW4A8DynamicLinearMethod(AscendLinearScheme):
|
|
"""Linear method for Ascend W4A8_DYNAMIC.
|
|
|
|
This method supports only weights quantized by msModelSlim. It supports two
|
|
weight layouts, distinguished by ``quant_version`` which comes from
|
|
``quant_description["version"]`` in the vLLM quantization config. Version
|
|
``"1.0.0"`` is the newer layout: it reduces the checkpoint weight size and
|
|
precomputes the ``scale_bias`` offline, reducing weight loading time.
|
|
|
|
The names below use ``linear`` as the checkpoint prefix of a linear layer,
|
|
``input_size`` as the logical input dimension, ``output_size`` as the
|
|
logical output dimension, and ``group_size`` as the number of input
|
|
channels per weight quantization group.
|
|
|
|
For ``quant_version != "1.0.0"``, the original linear weights are:
|
|
|
|
- ``linear.weight``: ``torch.int8``, ``[output_size, input_size]``.
|
|
Each int8 element stores one 4-bit weight value.
|
|
- ``linear.weight_scale``: ``params_dtype``, ``[output_size, 1]``.
|
|
- ``linear.weight_offset``: ``params_dtype``, ``[output_size, 1]``.
|
|
- ``linear.weight_scale_second``: ``params_dtype``,
|
|
``[output_size, input_size // group_size]``.
|
|
- ``linear.weight_offset_second``: ``torch.int64``,
|
|
``[output_size, input_size // group_size]``.
|
|
|
|
For ``quant_version == "1.0.0"``, the original linear weights are:
|
|
|
|
- ``linear.weight``: ``torch.int8``, ``[output_size // 2, input_size]``.
|
|
Each int8 element stores two packed 4-bit weight values along the output
|
|
dimension.
|
|
- ``linear.weight_scale``: ``params_dtype``, ``[output_size, 1]``.
|
|
- ``linear.weight_offset``: ``params_dtype``, ``[output_size, 1]``.
|
|
- ``linear.weight_scale_second``: ``params_dtype``,
|
|
``[output_size, input_size // group_size]``.
|
|
- ``linear.weight_offset_second``: ``torch.int64``,
|
|
``[output_size, input_size // group_size]``.
|
|
- ``linear.scale_bias``: ``torch.float32``, ``[output_size, 1]`` for
|
|
column-parallel linear layers and ``[output_size, 16]`` for
|
|
row-parallel linear layers.
|
|
|
|
In :meth:`process_weights_after_loading`, ``linear.weight`` is transposed
|
|
from ``[output, input]`` to the operator-oriented ``[input, output]``
|
|
layout. Old-version weights are converted with
|
|
``torch_npu.npu_convert_weight_to_int4pack``; new-version weights are
|
|
already packed as int4 pairs in int8 storage and are reinterpreted as int32
|
|
by grouping four int8 values.
|
|
|
|
After processing, ``torch_npu.npu_weight_quant_batchmatmul`` is called with
|
|
``weight`` as ``torch.int32`` in the operator-required packed layout
|
|
with shape ``[input_size, output_size // 8]`` and
|
|
``antiquant_scale`` as ``weight_scale * weight_scale_second`` converted to
|
|
``x.dtype`` with shape ``[input_size // group_size, output_size]``.
|
|
"""
|
|
|
|
def __init__(self):
|
|
vllm_config = get_current_vllm_config()
|
|
self.group_size = vllm_config.quant_config.quant_description.get("group_size", 256)
|
|
quant_version = vllm_config.quant_config.quant_description.get("version", "0")
|
|
self.new_quant_version = quant_version == "1.0.0"
|
|
|
|
self.tp_size = get_tensor_model_parallel_world_size()
|
|
|
|
def get_weight(self, input_size: int, output_size: int, params_dtype: torch.dtype) -> dict[str, Any]:
|
|
"""Create weight parameters.
|
|
|
|
For new quantization version (double int4 pack into int8), the output dimension
|
|
is compressed by factor 2 (e.g., [2048, 3072] -> [1024, 3072]). The returned
|
|
dict includes "_packed_dim" and "_packed_factor" for vLLM's weight loader.
|
|
"""
|
|
params_dict = {}
|
|
|
|
if self.new_quant_version:
|
|
# double int4 pack into int8: output dimension is compressed
|
|
pack_factor = 2
|
|
actual_output_size = output_size // pack_factor
|
|
params_dict["weight"] = torch.empty(actual_output_size, input_size, dtype=torch.int8)
|
|
# Add packing information for vLLM's weight_loader
|
|
params_dict["_packed_dim"] = 0
|
|
params_dict["_packed_factor"] = pack_factor
|
|
else:
|
|
params_dict["weight"] = torch.empty(output_size, input_size, dtype=torch.int8)
|
|
|
|
return params_dict
|
|
|
|
def get_pergroup_param(
|
|
self, input_size: int, output_size: int, params_dtype: torch.dtype, layer_type: str | None = None
|
|
) -> dict[str, Any]:
|
|
"""Create per-group quantization parameters."""
|
|
params_dict = {}
|
|
params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype)
|
|
params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype)
|
|
params_dict["weight_scale_second"] = torch.empty(output_size, input_size // self.group_size, dtype=params_dtype)
|
|
params_dict["weight_offset_second"] = torch.empty(
|
|
output_size, input_size // self.group_size, dtype=params_dtype
|
|
)
|
|
|
|
# NOTE: In w4a8 quantization implementation,
|
|
# for down_proj and o_proj(layer_type == "row") scale_bias shape is [output_size, 16],
|
|
# others are [output_size, 1]
|
|
if self.new_quant_version:
|
|
scale_bias_dim = 16 if layer_type == "row" else 1
|
|
|
|
params_dict["scale_bias"] = torch.empty(output_size, scale_bias_dim, dtype=torch.float32)
|
|
return params_dict
|
|
|
|
@staticmethod
|
|
def process_scale_second(
|
|
weight: torch.Tensor, scale: torch.Tensor, per_group_scale: torch.Tensor, is_new_quant: bool = False
|
|
):
|
|
"""Process the scale for second-level quantization.
|
|
|
|
Args:
|
|
weight: weight tensor [k, n] (in new version, n is already compressed to n/2)
|
|
scale: first-level quantization scale [output_size]
|
|
per_group_scale: second-level per-group quantization scale [group_num, n_scale]
|
|
is_new_quant: whether it's the new quantization version (weight already compressed)
|
|
|
|
Returns:
|
|
(antiquant_scale, bias): dequantization scale and bias (bias=None for new version)
|
|
"""
|
|
k, n = weight.shape
|
|
group_num, n_scale = per_group_scale.shape
|
|
|
|
if is_new_quant:
|
|
# Restore logical dimension for compressed weight
|
|
n = n * 2
|
|
|
|
bias = None
|
|
if not is_new_quant:
|
|
weight_high = weight.to(torch.float32).reshape(group_num, -1, n) * per_group_scale.reshape(group_num, 1, n)
|
|
weight_high = weight_high.reshape(k, n)
|
|
bias = 8 * (weight_high.to(torch.float32) * scale).sum(dim=0)
|
|
# NOTE: scale_bias is not used currently
|
|
# because in msmodelslim w4a8 uses symmetric quantization
|
|
|
|
# TODO: support potential future asymmetric quantization
|
|
antiquant_scale = (scale * per_group_scale).reshape(group_num, n)
|
|
return antiquant_scale.npu(), bias
|
|
|
|
def apply(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
tp_rank: int | None = None,
|
|
) -> torch.Tensor:
|
|
# NOTE: activation `x` is not quantized
|
|
return torch_npu.npu_weight_quant_batchmatmul(
|
|
x,
|
|
layer.weight,
|
|
antiquant_scale=layer.weight_scale_second.to(x.dtype),
|
|
antiquant_group_size=self.group_size,
|
|
)
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module):
|
|
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
|
|
layer.weight.data = maybe_trans_nz(layer.weight.data)
|
|
layer.weight_scale.data = layer.weight_scale.data.flatten().to(torch.float32)
|
|
layer.weight_offset.data = layer.weight_offset.data.flatten()
|
|
layer.weight_scale_second.data, scale_bias = self.process_scale_second(
|
|
layer.weight.data,
|
|
layer.weight_scale.data,
|
|
layer.weight_scale_second.data.transpose(0, 1).contiguous(),
|
|
is_new_quant=self.new_quant_version,
|
|
)
|
|
|
|
if self.new_quant_version:
|
|
# Process the loaded data based on layer type
|
|
if hasattr(layer, "scale_bias"):
|
|
if layer.scale_bias.data.shape[1] == 1:
|
|
layer.scale_bias.data = layer.scale_bias.data.flatten()
|
|
else:
|
|
layer.scale_bias.data = layer.scale_bias.data.contiguous()
|
|
else:
|
|
if scale_bias is not None:
|
|
param = torch.nn.Parameter(scale_bias, requires_grad=False)
|
|
layer.register_parameter("weight_scale_bias", param)
|
|
|
|
# Convert to NPU-specific int4pack format
|
|
if self.new_quant_version:
|
|
# weights on disk are already in packed int4 format
|
|
# pack 4 int8(int4*2) to int32
|
|
assert layer.weight.data.shape[-1] % 4 == 0, (
|
|
f"the last dim of weight needs to be divided by 4 but got shape {layer.weight.data.shape}"
|
|
)
|
|
layer.weight.data = layer.weight.data.view(torch.int32).contiguous()
|
|
else:
|
|
# weights are not compressed
|
|
# need to be packed via npu_convert_weight_to_int4pack
|
|
layer.weight.data = torch_npu.npu_convert_weight_to_int4pack(layer.weight.data.to(torch.int32))
|
|
|
|
|
|
@register_scheme("W4A8_DYNAMIC", "moe")
|
|
class AscendW4A8DynamicFusedMoEMethod(AscendMoEScheme):
|
|
"""FusedMoE method for Ascend W4A8_DYNAMIC.
|
|
|
|
This method supports four MoE weight formats: three generated by
|
|
msModelSlim and one generated by LLM-Compressor. The LLM-Compressor path
|
|
is selected when ``ascend_quant_method`` in ``quant_description`` is
|
|
``COMPRESSED_TENSORS_METHOD``. Otherwise, the msModelSlim path is used.
|
|
msModelSlim layouts are first distinguished by
|
|
``quant_description["version"] == "1.0.0"``; for version ``"1.0.0"``,
|
|
``group_size == 0`` selects per-channel weight quantization and
|
|
``group_size > 0`` selects per-group weight quantization.
|
|
|
|
The names below use ``L`` for the layer index, ``E`` for the expert index,
|
|
``num_experts`` for the routed expert count, ``hidden_sizes`` for the
|
|
hidden dimension, ``moe_intermediate_size`` for the expert intermediate
|
|
dimension, ``group_size`` for per-group weight quantization, and
|
|
``tp_size`` for tensor parallel size.
|
|
|
|
Original MoE layer weights generated by msModelSlim with
|
|
``quant_version != "1.0.0"``:
|
|
|
|
- ``model.layers.L.mlp.experts.E.gate_proj.weight``:
|
|
``torch.int8``, ``[moe_intermediate_size, hidden_sizes]``.
|
|
- ``model.layers.L.mlp.experts.E.up_proj.weight``:
|
|
``torch.int8``, ``[moe_intermediate_size, hidden_sizes]``.
|
|
- ``model.layers.L.mlp.experts.E.down_proj.weight``:
|
|
``torch.int8``, ``[hidden_sizes, moe_intermediate_size]``.
|
|
- Each linear also has ``weight_scale`` and ``weight_offset``:
|
|
``torch.float32``, ``[out_features, 1]``.
|
|
- Each linear also has ``weight_scale_second`` and
|
|
``weight_offset_second``. The ``weight_scale_second`` dtype is
|
|
``torch.float32`` and the ``weight_offset_second`` dtype is
|
|
``torch.int64``; both use shape
|
|
``[out_features, in_features // group_size]``.
|
|
|
|
Original MoE layer weights generated by msModelSlim with
|
|
``quant_version == "1.0.0"`` and per-group quantization:
|
|
|
|
- Compared with the previous msModelSlim layout, ``weight`` stores two
|
|
packed 4-bit values in each int8 element along the output dimension.
|
|
Therefore ``gate_proj.weight`` and ``up_proj.weight`` are
|
|
``torch.int8`` with shape
|
|
``[moe_intermediate_size // 2, hidden_sizes]``, and
|
|
``down_proj.weight`` is ``torch.int8`` with shape
|
|
``[hidden_sizes // 2, moe_intermediate_size]``.
|
|
- Each linear additionally has ``scale_bias``: ``torch.float32``,
|
|
``[moe_intermediate_size, 1]`` for ``gate_proj`` and ``up_proj``, and
|
|
``[hidden_sizes, 16 // tp_size]`` for ``down_proj``.
|
|
|
|
Original MoE layer weights generated by msModelSlim with
|
|
``quant_version == "1.0.0"`` and per-channel quantization:
|
|
|
|
- ``weight`` has the same packed shape as the previous msModelSlim
|
|
``1.0.0`` per-group layout.
|
|
- ``weight_scale`` and ``weight_offset`` are per-channel tensors:
|
|
``torch.float32``, ``[out_features, 1]``. There are no
|
|
``weight_scale_second`` or ``weight_offset_second`` tensors.
|
|
- Each linear also has ``scale_bias``: ``torch.float32``,
|
|
``[moe_intermediate_size, 1]`` for ``gate_proj`` and ``up_proj``, and
|
|
``[hidden_sizes, 16 // tp_size]`` for ``down_proj``.
|
|
|
|
Original MoE layer weights generated by LLM-Compressor:
|
|
|
|
- ``model.layers.L.mlp.experts.E.gate_proj.weight``:
|
|
``torch.int8``, ``[moe_intermediate_size, hidden_sizes]``.
|
|
- ``model.layers.L.mlp.experts.E.up_proj.weight``:
|
|
``torch.int8``, ``[moe_intermediate_size, hidden_sizes]``.
|
|
- ``model.layers.L.mlp.experts.E.down_proj.weight``:
|
|
``torch.int8``, ``[hidden_sizes, moe_intermediate_size]``.
|
|
- Each linear also has ``weight_scale``: ``torch.bfloat16``,
|
|
``[out_features, in_features // group_size]`` for group quantization, or
|
|
``[out_features, 1]`` for channel quantization.
|
|
|
|
During loading, ``gate_proj`` and ``up_proj`` are fused into ``w13`` and
|
|
``down_proj`` is loaded as ``w2``. Before
|
|
:meth:`process_weights_after_loading`, their logical shapes are:
|
|
|
|
- msModelSlim old: ``w13_weight`` ``torch.int8``,
|
|
``[num_experts, 2 * moe_intermediate_size, hidden_sizes]``; and
|
|
``w2_weight`` ``torch.int8``,
|
|
``[num_experts, hidden_sizes, moe_intermediate_size]``.
|
|
- msModelSlim ``1.0.0`` per-group and per-channel: ``w13_weight`` ``torch.int8``,
|
|
``[num_experts, moe_intermediate_size, hidden_sizes]``; and
|
|
``w2_weight`` ``torch.int8``,
|
|
``[num_experts, hidden_sizes // 2, moe_intermediate_size]``.
|
|
- LLM-Compressor: ``w13_weight`` ``torch.int8``,
|
|
``[num_experts, 2 * moe_intermediate_size, hidden_sizes]``; and
|
|
``w2_weight`` ``torch.int8``,
|
|
``[num_experts, hidden_sizes, moe_intermediate_size]``.
|
|
|
|
After processing, ``apply`` passes these tensors to the fused MoE operator:
|
|
|
|
- Shared by all formats:
|
|
``w13_weight``: ``torch.int32``,
|
|
``[num_experts, hidden_sizes, moe_intermediate_size // 4]``.
|
|
``w2_weight``: ``torch.int32``,
|
|
``[num_experts, moe_intermediate_size, hidden_sizes // 8]``.
|
|
``w13_scale_bias``: ``torch.float32``, ``[num_experts, 2 * moe_intermediate_size]``.
|
|
``w2_scale_bias``: ``torch.float32``, ``[num_experts, hidden_sizes]``.
|
|
- per-group:
|
|
``w13_weight_scale``: ``torch.int64``,
|
|
``[num_experts, hidden_sizes // group_size,
|
|
2 * moe_intermediate_size]``.
|
|
``w2_weight_scale``: ``torch.int64``,
|
|
``[num_experts, moe_intermediate_size // group_size, hidden_sizes]``.
|
|
- per-channel:
|
|
``w13_weight_scale``: ``torch.int64``,
|
|
``[num_experts, 2 * moe_intermediate_size]``.
|
|
``w2_weight_scale``: ``torch.int64``,
|
|
``[num_experts, 1, hidden_sizes]``.
|
|
"""
|
|
|
|
# Declare the quantization type for this scheme
|
|
quant_type: QuantType = QuantType.W4A8
|
|
|
|
def __init__(self):
|
|
self.supports_eplb = True
|
|
|
|
vllm_config = get_current_vllm_config()
|
|
self.group_size = vllm_config.quant_config.quant_description.get("group_size", 256)
|
|
# NOTE: the weights are quantized from bf16 to int4 through a per-channel quantization process
|
|
self.is_per_channel_weight = self.group_size == 0
|
|
quant_version = vllm_config.quant_config.quant_description.get("version", "0")
|
|
# NOTE: new quantize weights: 2 int4 pack into int8
|
|
self.new_quant_version = quant_version == "1.0.0"
|
|
|
|
self.quant_method = vllm_config.quant_config.quant_description.get("ascend_quant_method", "")
|
|
if self.quant_method == COMPRESSED_TENSORS_METHOD:
|
|
self.weight_strategy = vllm_config.quant_config.quant_description.get("weight_strategy", "group")
|
|
|
|
self.tp_size = (
|
|
1 if vllm_config.parallel_config.enable_expert_parallel else get_tensor_model_parallel_world_size()
|
|
)
|
|
self.dynamic_eplb = get_ascend_config().eplb_config.dynamic_eplb
|
|
if self.new_quant_version and self.tp_size > 16:
|
|
raise ValueError("The current weight does not support moe part tp>16.")
|
|
|
|
try:
|
|
device_group = get_mc2_group().device_group
|
|
# TODO: Try local_rank = ep_group.rank_in_group
|
|
local_rank = torch.distributed.get_rank(group=device_group)
|
|
backend = device_group._get_backend(torch.device("npu"))
|
|
self.moe_all_to_all_group_name = backend.get_hccl_comm_name(local_rank)
|
|
except AttributeError:
|
|
self.moe_all_to_all_group_name = ""
|
|
|
|
def get_weight(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
if self.quant_method == COMPRESSED_TENSORS_METHOD:
|
|
return self.get_weight_compressed_tensors(
|
|
num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype
|
|
)
|
|
else:
|
|
return self.get_weight_modelslim(num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype)
|
|
|
|
def get_weight_compressed_tensors(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
param_dict = {}
|
|
E = num_experts
|
|
H = hidden_sizes
|
|
IN = intermediate_size_per_partition
|
|
|
|
param_dict["w13_weight"] = torch.empty(E, 2 * IN, H, dtype=torch.int8)
|
|
param_dict["w2_weight"] = torch.empty(E, H, IN, dtype=torch.int8)
|
|
return param_dict
|
|
|
|
def get_weight_modelslim(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
param_dict = {}
|
|
if self.new_quant_version:
|
|
w13_output_size = intermediate_size_per_partition
|
|
w2_output_size = hidden_sizes // 2
|
|
else:
|
|
w13_output_size = 2 * intermediate_size_per_partition
|
|
w2_output_size = hidden_sizes
|
|
|
|
param_dict["w13_weight"] = torch.empty(num_experts, w13_output_size, hidden_sizes, dtype=torch.int8)
|
|
param_dict["w2_weight"] = torch.empty(
|
|
num_experts, w2_output_size, intermediate_size_per_partition, dtype=torch.int8
|
|
)
|
|
return param_dict
|
|
|
|
def get_dynamic_quant_param(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
if self.quant_method == COMPRESSED_TENSORS_METHOD:
|
|
return self.get_dynamic_quant_param_compressed_tensors(
|
|
num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype
|
|
)
|
|
else:
|
|
return self.get_dynamic_quant_param_modelslim(
|
|
num_experts, intermediate_size_per_partition, hidden_sizes, params_dtype
|
|
)
|
|
|
|
def get_dynamic_quant_param_compressed_tensors(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
param_dict = {}
|
|
|
|
E = num_experts
|
|
H = hidden_sizes
|
|
IN = intermediate_size_per_partition
|
|
g = self.group_size
|
|
|
|
# Per-row scale columns
|
|
def _n_scale_cols(in_features: int) -> int:
|
|
return 1 if g <= 0 else (in_features // g)
|
|
|
|
param_dict["w13_weight_scale"] = torch.empty(E, 2 * IN, _n_scale_cols(H), dtype=torch.bfloat16)
|
|
|
|
param_dict["w2_weight_scale"] = torch.empty(E, H, _n_scale_cols(IN), dtype=torch.bfloat16)
|
|
|
|
return param_dict
|
|
|
|
def get_dynamic_quant_param_modelslim(
|
|
self, num_experts: int, intermediate_size_per_partition: int, hidden_sizes: int, params_dtype: torch.dtype
|
|
) -> dict[str, Any]:
|
|
param_dict = {}
|
|
param_dict["w13_weight_scale"] = torch.empty(
|
|
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
|
)
|
|
|
|
param_dict["w13_weight_offset"] = torch.empty(
|
|
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
|
)
|
|
|
|
param_dict["w2_weight_scale"] = torch.empty(num_experts, hidden_sizes, 1, dtype=torch.float32)
|
|
param_dict["w2_weight_offset"] = torch.empty(num_experts, hidden_sizes, 1, dtype=torch.float32)
|
|
if not self.is_per_channel_weight:
|
|
param_dict["w13_weight_scale_second"] = torch.empty(
|
|
num_experts, 2 * intermediate_size_per_partition, hidden_sizes // self.group_size, dtype=torch.float32
|
|
)
|
|
param_dict["w13_weight_offset_second"] = torch.empty(
|
|
num_experts, 2 * intermediate_size_per_partition, hidden_sizes // self.group_size, dtype=torch.float32
|
|
)
|
|
|
|
param_dict["w2_weight_scale_second"] = torch.empty(
|
|
num_experts, hidden_sizes, intermediate_size_per_partition // self.group_size, dtype=torch.float32
|
|
)
|
|
param_dict["w2_weight_offset_second"] = torch.empty(
|
|
num_experts, hidden_sizes, intermediate_size_per_partition // self.group_size, dtype=torch.float32
|
|
)
|
|
|
|
if self.new_quant_version:
|
|
param_dict["w13_scale_bias"] = torch.empty(
|
|
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
|
|
)
|
|
param_dict["w2_scale_bias"] = torch.empty(
|
|
num_experts, hidden_sizes, 16 // self.tp_size, dtype=torch.float32
|
|
)
|
|
|
|
return param_dict
|
|
|
|
def apply(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
router_logits: torch.Tensor,
|
|
top_k: int,
|
|
renormalize: bool,
|
|
use_grouped_topk: bool = False,
|
|
num_experts: int = -1,
|
|
expert_map: torch.Tensor | None = None,
|
|
topk_group: int | None = None,
|
|
num_expert_group: int | None = None,
|
|
custom_routing_function: Callable | None = None,
|
|
scoring_func: str = "softmax",
|
|
routed_scaling_factor: float = 1.0,
|
|
e_score_correction_bias: torch.Tensor | None = None,
|
|
is_prefill: bool = True,
|
|
enable_force_load_balance: bool = False,
|
|
log2phy: torch.Tensor | None = None,
|
|
global_redundant_expert_num: int = 0,
|
|
pertoken_scale: torch.Tensor | None = None,
|
|
activation: str = "silu",
|
|
apply_router_weight_on_input: bool = False,
|
|
mc2_mask: torch.Tensor | None = None,
|
|
tid2eid: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
num_shared_experts = getattr(layer, "n_shared_experts", 0)
|
|
if num_shared_experts is None:
|
|
num_shared_experts = 0
|
|
num_logical_experts = get_moe_num_logical_experts(
|
|
layer,
|
|
num_experts,
|
|
global_redundant_expert_num=global_redundant_expert_num,
|
|
num_shared_experts=num_shared_experts,
|
|
)
|
|
assert router_logits.shape[1] == num_logical_experts, (
|
|
"Number of global experts mismatch (excluding redundancy): "
|
|
f"router_logits.shape[1]={router_logits.shape[1]}, num_logical_experts={num_logical_experts}"
|
|
)
|
|
|
|
# NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern
|
|
topk_weights, topk_ids = select_experts(
|
|
hidden_states=x,
|
|
router_logits=router_logits,
|
|
top_k=top_k,
|
|
use_grouped_topk=use_grouped_topk,
|
|
renormalize=renormalize,
|
|
topk_group=topk_group,
|
|
num_expert_group=num_expert_group,
|
|
custom_routing_function=custom_routing_function,
|
|
scoring_func=scoring_func,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
e_score_correction_bias=e_score_correction_bias,
|
|
num_experts=num_logical_experts,
|
|
tid2eid=tid2eid,
|
|
)
|
|
|
|
# this is a naive implementation for experts load balance so as
|
|
# to avoid accumulating too much tokens on a single rank.
|
|
# currently it is only activated when doing profile runs.
|
|
if enable_force_load_balance:
|
|
random_matrix = torch.rand(topk_ids.size(0), num_logical_experts, device=topk_ids.device)
|
|
topk_ids = torch.argsort(random_matrix, dim=1)[:, : topk_ids.size(1)].to(topk_ids.dtype)
|
|
|
|
topk_weights = topk_weights.to(x.dtype)
|
|
|
|
if self.dynamic_eplb:
|
|
w1 = [i.view(torch.int32) for i in layer.w13_weight_list]
|
|
w1_scale = layer.w13_weight_scale_list
|
|
w2 = [i.view(torch.int32) for i in layer.w2_weight_list]
|
|
w2_scale = layer.w2_weight_scale_list
|
|
w1_scale_bias = layer.w13_scale_bias_list
|
|
w2_scale_bias = layer.w2_scale_bias_list
|
|
else:
|
|
w1 = [layer.w13_weight]
|
|
w1_scale = [layer.w13_weight_scale]
|
|
w2 = [layer.w2_weight]
|
|
w2_scale = [layer.w2_weight_scale]
|
|
w1_scale_bias = [layer.w13_scale_bias.detach()] if hasattr(layer, "w13_scale_bias") else None
|
|
w2_scale_bias = [layer.w2_scale_bias.detach()] if hasattr(layer, "w2_scale_bias") else None
|
|
|
|
moe_comm_method = _EXTRA_CTX.moe_comm_method
|
|
return moe_comm_method.fused_experts(
|
|
fused_experts_input=build_fused_experts_input(
|
|
hidden_states=x,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
w1=w1,
|
|
w2=w2,
|
|
quant_type=self.quant_type,
|
|
dynamic_eplb=self.dynamic_eplb,
|
|
expert_map=expert_map,
|
|
global_redundant_expert_num=global_redundant_expert_num,
|
|
mc2_mask=mc2_mask,
|
|
apply_router_weight_on_input=apply_router_weight_on_input,
|
|
log2phy=log2phy,
|
|
pertoken_scale=pertoken_scale,
|
|
activation=activation,
|
|
w1_scale=w1_scale,
|
|
w2_scale=w2_scale,
|
|
w1_scale_bias=w1_scale_bias,
|
|
w2_scale_bias=w2_scale_bias,
|
|
is_per_channel_weight=self.is_per_channel_weight,
|
|
swiglu_limit=layer.swiglu_limit,
|
|
)
|
|
)
|
|
|
|
def process_scale(self, weight: torch.Tensor, scale, per_group_scale):
|
|
scale = scale.transpose(1, 2).contiguous()
|
|
if self.is_per_channel_weight:
|
|
scale_np = scale.cpu().numpy()
|
|
scale_np.dtype = np.uint32
|
|
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
|
|
return scale_uint64_tensor, None
|
|
per_group_scale = per_group_scale.transpose(1, 2).contiguous()
|
|
group_num, k, n = weight.shape
|
|
# the weight of the new version is reduced by half by pack n, so it needs to be restored
|
|
if self.new_quant_version:
|
|
n = n * 2
|
|
per_group_scale = per_group_scale.reshape(group_num, -1, n)
|
|
group_num, quantgroup_num, n = per_group_scale.shape
|
|
bias = None
|
|
if not self.new_quant_version:
|
|
weight_high = weight.to(torch.float32).reshape(
|
|
[group_num, quantgroup_num, -1, n]
|
|
) * per_group_scale.reshape([group_num, quantgroup_num, 1, n])
|
|
weight_high = weight_high.reshape([group_num, k, n])
|
|
bias = 8 * (weight_high.to(torch.float32) * scale).sum(axis=1)
|
|
scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32)
|
|
scale_fp32_np = scale_fp32.cpu().numpy()
|
|
scale_fp32_np.dtype = np.uint32
|
|
sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32)
|
|
|
|
sscale_uint64[..., ::2] = scale_fp32_np
|
|
|
|
sscale_uint64_buffer = np.frombuffer(sscale_uint64.tobytes(), dtype=np.int64).copy()
|
|
sscale_uint64_tensor = torch.from_numpy(sscale_uint64_buffer).reshape(group_num, quantgroup_num, n)
|
|
sscale_uint64_tensor = sscale_uint64_tensor.npu()
|
|
return sscale_uint64_tensor, bias
|
|
|
|
def update_bias(self, layer, w13_bias, w2_bias):
|
|
if self.new_quant_version:
|
|
layer.w13_scale_bias.data = layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
|
|
layer.w2_scale_bias.data = layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
|
|
else:
|
|
w13_scale_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
|
|
layer.register_parameter("w13_scale_bias", w13_scale_bias)
|
|
w2_scale_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
|
|
layer.register_parameter("w2_scale_bias", w2_scale_bias)
|
|
|
|
def pack_to_int32(self, weight: torch.Tensor):
|
|
if self.new_quant_version or self.quant_method == COMPRESSED_TENSORS_METHOD:
|
|
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
|
|
assert weight.shape[-1] % 4 == 0, (
|
|
f"the last dim of weight needs to be divided by 4 but got shape {weight.shape}"
|
|
)
|
|
return weight.view(torch.int32).contiguous()
|
|
else:
|
|
return torch_npu.npu_quantize(
|
|
weight.to(torch.float32), torch.tensor([1.0]).npu(), None, torch.quint4x2, -1, False
|
|
)
|
|
|
|
def pack_int4_to_int8(self, weight: torch.Tensor) -> torch.Tensor:
|
|
shape = weight.shape
|
|
weight = weight.reshape(-1, 2)
|
|
weight0 = weight[:, :1]
|
|
weight1 = weight[:, 1:]
|
|
weight1_4 = torch.bitwise_left_shift(weight1, 4)
|
|
weight2_4 = weight0 & 0b00001111
|
|
weight_add = torch.bitwise_or(weight1_4, weight2_4)
|
|
# The clone() call is used to break the view chain
|
|
return weight_add.reshape(shape[:-1] + (shape[-1] // 2,)).clone()
|
|
|
|
@staticmethod
|
|
def maybe_squeeze_per_channel_weight_scale(scale: torch.Tensor) -> torch.Tensor:
|
|
if scale.dim() > 1 and scale.shape[1] == 1:
|
|
return scale.squeeze(1)
|
|
return scale
|
|
|
|
def process_weights_after_loading(self, layer):
|
|
if self.quant_method == COMPRESSED_TENSORS_METHOD:
|
|
self.process_weights_after_loading_compressed_tensors(layer)
|
|
else:
|
|
self.process_weights_after_loading_modelslim(layer)
|
|
|
|
def process_weights_after_loading_compressed_tensors(self, layer):
|
|
layer.w13_weight.data = layer.w13_weight.data.transpose(1, 2).contiguous()
|
|
layer.w2_weight.data = layer.w2_weight.data.transpose(1, 2).contiguous()
|
|
|
|
def process_scale_compressed_tensors(scale: torch.Tensor, squeeze: bool = True):
|
|
scale = scale.transpose(1, 2).to(torch.float32).contiguous()
|
|
scale_np = scale.cpu().numpy()
|
|
scale_np.dtype = np.uint32
|
|
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
|
|
if self.is_per_channel_weight and squeeze:
|
|
return self.maybe_squeeze_per_channel_weight_scale(scale_uint64_tensor)
|
|
return scale_uint64_tensor
|
|
|
|
def update_bias_compressed_tensors(weight: torch.Tensor, scale: torch.Tensor, strategy: str):
|
|
group_num, k, n = weight.shape
|
|
scale = scale.transpose(1, 2).contiguous()
|
|
scale = scale.reshape(group_num, -1, n)
|
|
group_num, quantgroup_num, n = scale.shape
|
|
|
|
bias = None
|
|
if strategy == "group":
|
|
tmp = weight.to(torch.float32).reshape([group_num, quantgroup_num, -1, n]) * scale.reshape(
|
|
[group_num, quantgroup_num, 1, n]
|
|
)
|
|
tmp = tmp.reshape([group_num, k, n])
|
|
bias = 8 * tmp.sum(axis=1)
|
|
elif strategy == "channel":
|
|
bias = 8 * (weight.to(torch.float32) * scale).sum(axis=1)
|
|
else:
|
|
raise ValueError(f"Unsupported weight strategy: {strategy}")
|
|
return bias
|
|
|
|
w13_bias = update_bias_compressed_tensors(
|
|
layer.w13_weight.data, layer.w13_weight_scale.data, self.weight_strategy
|
|
)
|
|
w2_bias = update_bias_compressed_tensors(layer.w2_weight.data, layer.w2_weight_scale.data, self.weight_strategy)
|
|
|
|
layer.w13_weight_scale.data = process_scale_compressed_tensors(layer.w13_weight_scale.data)
|
|
# To use torch_npu.npu_grouped_matmul, keep w2_weigh_scale unsqueezed
|
|
layer.w2_weight_scale.data = process_scale_compressed_tensors(layer.w2_weight_scale.data, False)
|
|
|
|
w13_scale_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
|
|
layer.register_parameter("w13_scale_bias", w13_scale_bias)
|
|
w2_scale_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
|
|
layer.register_parameter("w2_scale_bias", w2_scale_bias)
|
|
|
|
# Packs 2 int4 into 1 int8 on-the-fly to mirror the modelslim new_quant_version path
|
|
layer.w13_weight.data = self.pack_int4_to_int8(layer.w13_weight.data)
|
|
layer.w2_weight.data = self.pack_int4_to_int8(layer.w2_weight.data)
|
|
layer.w13_weight.data = maybe_trans_nz(layer.w13_weight.data)
|
|
layer.w2_weight.data = maybe_trans_nz(layer.w2_weight.data)
|
|
layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data)
|
|
layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data)
|
|
|
|
def process_weights_after_loading_modelslim(self, layer):
|
|
layer.w13_weight.data = layer.w13_weight.data.transpose(1, 2).contiguous()
|
|
layer.w2_weight.data = layer.w2_weight.data.transpose(1, 2).contiguous()
|
|
|
|
w13_weight_scale_second = (
|
|
layer.w13_weight_scale_second.data if hasattr(layer, "w13_weight_scale_second") else None
|
|
)
|
|
w2_weight_scale_second = layer.w2_weight_scale_second.data if hasattr(layer, "w2_weight_scale_second") else None
|
|
layer.w13_weight_scale.data, w13_bias = self.process_scale(
|
|
layer.w13_weight, layer.w13_weight_scale.data, w13_weight_scale_second
|
|
)
|
|
layer.w2_weight_scale.data, w2_bias = self.process_scale(
|
|
layer.w2_weight, layer.w2_weight_scale.data, w2_weight_scale_second
|
|
)
|
|
if hasattr(layer, "w13_weight_scale_second"):
|
|
# scale_second is no longer used, release this part of the memory
|
|
del layer.w13_weight_scale_second
|
|
del layer.w2_weight_scale_second
|
|
del layer.w13_weight_offset_second
|
|
del layer.w2_weight_offset_second
|
|
|
|
self.update_bias(layer, w13_bias, w2_bias)
|
|
|
|
if self.is_per_channel_weight:
|
|
layer.w13_weight_scale.data = self.maybe_squeeze_per_channel_weight_scale(layer.w13_weight_scale.data)
|
|
layer.w13_weight.data = maybe_trans_nz(layer.w13_weight.data)
|
|
layer.w2_weight.data = maybe_trans_nz(layer.w2_weight.data)
|
|
|
|
if self.dynamic_eplb:
|
|
layer.w13_weight_list = [weight.clone() for weight in layer.w13_weight.data.unbind(dim=0)]
|
|
layer.w2_weight_list = [weight.clone() for weight in layer.w2_weight.data.unbind(dim=0)]
|
|
layer.w13_weight_scale_list = [weight.clone() for weight in layer.w13_weight_scale.data.unbind(dim=0)]
|
|
layer.w2_weight_scale_list = [weight.clone() for weight in layer.w2_weight_scale.data.unbind(dim=0)]
|
|
layer.w13_scale_bias_list = (
|
|
[weight.clone() for weight in layer.w13_scale_bias.data.unbind(dim=0)]
|
|
if hasattr(layer, "w13_scale_bias")
|
|
else None
|
|
)
|
|
layer.w2_scale_bias_list = (
|
|
[weight.clone() for weight in layer.w2_scale_bias.data.unbind(dim=0)]
|
|
if hasattr(layer, "w2_scale_bias")
|
|
else None
|
|
)
|
|
del layer.w13_weight
|
|
del layer.w2_weight
|
|
del layer.w13_weight_scale
|
|
del layer.w2_weight_scale
|
|
del layer.w13_scale_bias
|
|
del layer.w2_scale_bias
|
|
else:
|
|
layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data)
|
|
layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data)
|