465 lines
24 KiB
Python
465 lines
24 KiB
Python
import itertools
|
||
from abc import abstractmethod
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch.nn.parameter import Parameter, UninitializedParameter
|
||
|
||
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
|
||
get_tensor_model_parallel_world_size,
|
||
split_tensor_along_last_dim,
|
||
tensor_model_parallel_all_gather,
|
||
tensor_model_parallel_all_reduce)
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.quantization.base_config import (
|
||
QuantizationConfig, QuantizeMethodBase)
|
||
# yapf: disable
|
||
from vllm.model_executor.parameter import (BasevLLMParameter,
|
||
BlockQuantScaleParameter,
|
||
PackedColumnParameter,
|
||
PackedvLLMParameter,
|
||
PerTensorScaleParameter,
|
||
RowvLLMParameter)
|
||
# yapf: enable
|
||
from vllm.model_executor.utils import set_weight_attrs
|
||
|
||
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
|
||
ReplicatedLinear,
|
||
WEIGHT_LOADER_V2_SUPPORTED,
|
||
LinearBase,
|
||
RowParallelLinear)
|
||
|
||
def ReplicatedLinear__init__(self,
|
||
input_size: int,
|
||
output_size: int,
|
||
bias: bool = True,
|
||
skip_bias_add: bool = False,
|
||
params_dtype: Optional[torch.dtype] = None,
|
||
quant_config: Optional[QuantizationConfig] = None,
|
||
prefix: str = ""):
|
||
super(ReplicatedLinear,self).__init__(input_size,
|
||
output_size,
|
||
skip_bias_add,
|
||
params_dtype,
|
||
quant_config,
|
||
prefix=prefix)
|
||
|
||
# All the linear layer supports quant method.
|
||
assert self.quant_method is not None
|
||
|
||
self.scale_k = 1 # quant_block_k 128 需要除以 scale_k， 如设置为2 即 quant_block_k 是 64
|
||
self.scale_k_slice = 1
|
||
self.scale_n = 1
|
||
self.scale_n_slice = 1
|
||
if quant_config is not None and hasattr(quant_config, "weight_block_size") and quant_config.weight_block_size is not None:
|
||
gcd_value = quant_config.weight_block_size[1]
|
||
import math
|
||
if input_size % quant_config.weight_block_size[1]:
|
||
gcd_value = math.gcd(input_size % quant_config.weight_block_size[1], quant_config.weight_block_size[1])
|
||
self.scale_k =self.scale_k * quant_config.weight_block_size[1] // gcd_value
|
||
self.scale_k_slice = input_size // gcd_value
|
||
if output_size % quant_config.weight_block_size[0]:
|
||
gcd_value = math.gcd(output_size % quant_config.weight_block_size[0], quant_config.weight_block_size[0])
|
||
self.scale_n = self.scale_n * quant_config.weight_block_size[0] // gcd_value
|
||
self.scale_n_slice = output_size // gcd_value
|
||
|
||
self.quant_method.create_weights(self,
|
||
self.input_size, [self.output_size],
|
||
self.input_size,
|
||
self.output_size,
|
||
self.params_dtype,
|
||
scale_k = self.scale_k,
|
||
scale_n = self.scale_n,
|
||
weight_loader=self.weight_loader)
|
||
|
||
if bias:
|
||
self.bias = Parameter(
|
||
torch.empty(self.output_size, dtype=self.params_dtype))
|
||
set_weight_attrs(self.bias, {
|
||
"output_dim": 0,
|
||
"weight_loader": self.weight_loader,
|
||
})
|
||
else:
|
||
self.register_parameter("bias", None)
|
||
|
||
def ReplicatedLinear_weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
|
||
# If the weight on disk does not have a shape, give it one
|
||
# (such scales for AutoFp8).
|
||
if len(loaded_weight.shape) == 0:
|
||
loaded_weight = loaded_weight.reshape(1)
|
||
|
||
if len(loaded_weight.shape) == 0:
|
||
assert loaded_weight.numel() == 1
|
||
loaded_weight = loaded_weight.reshape(1)
|
||
if self.quant_method.__class__.__name__ in ['Fp8LinearMethod', 'Fp8MoEMethod'] and torch.finfo(loaded_weight.dtype).bits > 8:
|
||
if self.scale_k > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_k, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,2,0).reshape([loaded_weight.shape[1], -1])[:, :self.scale_k_slice]
|
||
#[1,n,k] -> [scale_k,n,k] -> [n,k,scale_k] -> [n, k*scale_k]
|
||
|
||
if self.scale_n > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_n, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,0,2).reshape([-1, loaded_weight.shape[2]])[:self.scale_n_slice]
|
||
|
||
assert param.size() == loaded_weight.size(), f'{param.size()}, {loaded_weight.size()}'
|
||
param.data.copy_(loaded_weight)
|
||
|
||
def refine_block(block_size:list[int],
|
||
weight_size:list[int],
|
||
dim:int=0,
|
||
pingpong_size:int = 2.5*1024*1024, #bytes
|
||
core_number:int = 4,
|
||
data_type:int = 2, #bfloat16
|
||
max_iter_number:int = 2):
|
||
'''
|
||
对于不均匀分core， 需要每个core <= 2.5M 才能保证可以pingpong,
|
||
core间相差数量为 block_size[dim] * weight_size[1-dim]
|
||
缩小block_size可以减小core间差距，使得更平均一些，直到大core数据量不超
|
||
如果均匀分core已经超了或者没有超，就没必要调整
|
||
'''
|
||
if dim < 0:
|
||
dim = 2 + dim
|
||
|
||
pingpong_size = pingpong_size / data_type # number of data
|
||
|
||
block_size_refine = block_size[dim]
|
||
all_block_number = weight_size[dim] // block_size_refine
|
||
|
||
if all_block_number % core_number == 0:
|
||
#均分,这种情况不管有没有超，都无需调整
|
||
return block_size_refine
|
||
|
||
block_number_tiny = all_block_number // core_number
|
||
block_number_big = all_block_number // core_number + 1
|
||
if block_number_tiny * block_size_refine * weight_size[1-dim] >= pingpong_size or \
|
||
block_number_big * block_size_refine * weight_size[1-dim] <= pingpong_size :
|
||
# 小的已经超了，无法再调整了
|
||
# 大的没有超，无需调整
|
||
return block_size_refine
|
||
|
||
all_block_number_tmp = all_block_number
|
||
block_size_refine_tmp = block_size_refine
|
||
for iter_index in range(max_iter_number):
|
||
all_block_number_tmp = all_block_number_tmp * 2
|
||
block_size_refine_tmp = block_size_refine_tmp // 2
|
||
if all_block_number_tmp % core_number == 0:
|
||
block_number_tiny = all_block_number // core_number
|
||
if block_number_tiny * block_size_refine_tmp * weight_size[1-dim] <= pingpong_size:
|
||
return block_size_refine_tmp
|
||
else:
|
||
#均分还是超了，无需调整
|
||
return block_size_refine
|
||
else:
|
||
block_number_big = all_block_number_tmp // core_number + 1
|
||
if block_number_big * block_size_refine_tmp * weight_size[1-dim] <= pingpong_size:
|
||
return block_size_refine_tmp
|
||
|
||
return block_size_refine
|
||
|
||
def ColumnParallelLinear__init__(self,
|
||
input_size: int,
|
||
output_size: int,
|
||
bias: bool = True,
|
||
gather_output: bool = False,
|
||
skip_bias_add: bool = False,
|
||
params_dtype: Optional[torch.dtype] = None,
|
||
quant_config: Optional[QuantizationConfig] = None,
|
||
output_sizes: Optional[List[int]] = None,
|
||
prefix: str = "",
|
||
*,
|
||
return_bias: bool = True,
|
||
disable_tp: bool = False,):
|
||
# Divide the weight matrix along the last dimension.
|
||
self.tp_rank = (get_tensor_model_parallel_rank()
|
||
if not disable_tp else 0)
|
||
self.tp_size = (get_tensor_model_parallel_world_size()
|
||
if not disable_tp else 1)
|
||
self.input_size_per_partition = input_size
|
||
self.output_size_per_partition = divide(output_size, self.tp_size)
|
||
self.output_partition_sizes = [self.output_size_per_partition]
|
||
# If QKV or MergedColumn, use output size of each partition.
|
||
if hasattr(self, "output_sizes"):
|
||
self.output_partition_sizes = [
|
||
divide(output_size, self.tp_size)
|
||
for output_size in self.output_sizes
|
||
]
|
||
super(ColumnParallelLinear,self).__init__(input_size,
|
||
output_size,
|
||
skip_bias_add,
|
||
params_dtype,
|
||
quant_config,
|
||
prefix,
|
||
return_bias=return_bias,
|
||
disable_tp=disable_tp)
|
||
|
||
self.gather_output = gather_output
|
||
|
||
if output_sizes is None:
|
||
output_sizes = [output_size]
|
||
|
||
self.scale_n = 1
|
||
if quant_config is not None and hasattr(quant_config, "weight_block_size") and quant_config.weight_block_size is not None:
|
||
gcd_value = quant_config.weight_block_size[0]
|
||
|
||
import math
|
||
if hasattr(self, "output_sizes"):
|
||
# 对于Merge类型的ColumnParallelLinear来说，需要根据每个Part Linear的shape，去计算最小公约数
|
||
output_size_no_merge = self.output_partition_sizes
|
||
block_values = [o % quant_config.weight_block_size[0] for o in output_size_no_merge]
|
||
is_gcd_recompute = sum(block_values)
|
||
|
||
if is_gcd_recompute:
|
||
import math
|
||
block_values.append(quant_config.weight_block_size[0])
|
||
gcd_value = math.gcd(*block_values)
|
||
# Notice:
|
||
# 这儿对于非对齐的Part-Weight， 可能需要验证一下流程
|
||
# 对于DeepSeek来说，仅存在于MLP&MOE中的MergeColumnLinear，都是Shape一致的PartWeight
|
||
# 对于QWen3来说，会存在QKVColumnLinear，是Shape不一致的PartWeight，但是由于QWen3当下的切分方案，对于gcd_value无感，无需重计算所以暂时不会进来
|
||
if hasattr(self, "output_sizes") and len(output_size_no_merge) == 2 and output_size_no_merge[0] == output_size_no_merge[1]:
|
||
#only refine mlp w13
|
||
gcd_value = refine_block([gcd_value, quant_config.weight_block_size[1]], [output_size_no_merge[0], input_size])
|
||
self.scale_n =self.scale_n * quant_config.weight_block_size[0] // gcd_value
|
||
else:
|
||
# 对于非Merge的ColumnParallelLinear来说, 仅仅根据当下shape去计算最小公约数
|
||
output_size_no_merge = self.output_size_per_partition
|
||
is_gcd_recompute = output_size_no_merge % quant_config.weight_block_size[0]
|
||
if is_gcd_recompute:
|
||
gcd_value = math.gcd(output_size_no_merge % quant_config.weight_block_size[0], quant_config.weight_block_size[0])
|
||
self.scale_n =self.scale_n * quant_config.weight_block_size[0] // gcd_value
|
||
|
||
|
||
|
||
self.quant_method.create_weights(
|
||
layer=self,
|
||
input_size_per_partition=self.input_size,
|
||
output_partition_sizes=self.output_partition_sizes,
|
||
input_size=self.input_size,
|
||
output_size=self.output_size,
|
||
params_dtype=self.params_dtype,
|
||
scale_n = self.scale_n,
|
||
weight_loader=(
|
||
self.weight_loader_v2 if self.quant_method.__class__.__name__
|
||
in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
|
||
if bias:
|
||
self.bias = Parameter(
|
||
torch.empty(self.output_size_per_partition,
|
||
dtype=params_dtype))
|
||
set_weight_attrs(self.bias, {
|
||
"output_dim": 0,
|
||
"weight_loader": self.weight_loader,
|
||
})
|
||
else:
|
||
self.register_parameter("bias", None)
|
||
|
||
self.update_param_tp_status()
|
||
|
||
def ColumnParallelLinear_weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor):
|
||
# Special case for loading scales off disk, which often do not
|
||
# have a shape (such as in the case of AutoFP8).
|
||
if len(loaded_weight.shape) == 0:
|
||
assert loaded_weight.numel() == 1
|
||
loaded_weight = loaded_weight.reshape(1)
|
||
if self.quant_method.__class__.__name__ in ['Fp8LinearMethod', 'Fp8MoEMethod'] and torch.finfo(loaded_weight.dtype).bits > 8:
|
||
if self.scale_n > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_n, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,0,2).reshape([-1, loaded_weight.shape[-1]])
|
||
#[1,n,k] -> [scale_n,n,k] -> [n,scale_n,n,k] -> [n*scale_n, k]
|
||
param.load_column_parallel_weight(loaded_weight=loaded_weight)
|
||
|
||
|
||
class MergedColumnParallelLinear(ColumnParallelLinear):
|
||
def weight_loader_v2(self,
|
||
param: BasevLLMParameter,
|
||
loaded_weight: torch.Tensor,
|
||
loaded_shard_id: Optional[int] = None):
|
||
|
||
if self.quant_method.__class__.__name__ in ['Fp8LinearMethod', 'Fp8MoEMethod'] and torch.finfo(loaded_weight.dtype).bits > 8:
|
||
if self.scale_n > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_n, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,0,2).reshape([-1, loaded_weight.shape[-1]])
|
||
#[1,n,k] -> [scale_n,n,k] -> [n,scale_n,n,k] -> [n*scale_n, k]
|
||
|
||
if self.quant_method.__class__.__name__ in ['GPTQLinearMethod']:
|
||
if self.quant_method.scale_k > 1 and len(loaded_weight.shape) == 2 and loaded_weight.dtype in [torch.float16, torch.bfloat16, torch.float32]:
|
||
loaded_weight = loaded_weight.unsqueeze(1) #[k,1,n]
|
||
loaded_weight = loaded_weight.expand(loaded_weight.shape[0], self.quant_method.scale_k, loaded_weight.shape[2]).reshape([-1, loaded_weight.shape[2]])
|
||
#[k,1,n] -> [k,scale_k,n]] -> [k*scale_k, n]
|
||
|
||
if loaded_shard_id is None:
|
||
if isinstance(param, PerTensorScaleParameter):
|
||
param.load_merged_column_weight(loaded_weight=loaded_weight,
|
||
shard_id=0)
|
||
return
|
||
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
|
||
param.load_merged_column_weight(loaded_weight=loaded_weight)
|
||
return
|
||
# TODO: @dsikka - move to parameter.py
|
||
self._load_fused_module_from_checkpoint(param, loaded_weight)
|
||
return
|
||
|
||
assert loaded_shard_id < len(self.output_sizes)
|
||
|
||
tp_size = get_tensor_model_parallel_world_size()
|
||
|
||
if isinstance(param, BlockQuantScaleParameter):
|
||
from vllm.model_executor.layers.quantization.fp8 import (
|
||
Fp8LinearMethod, Fp8MoEMethod)
|
||
assert self.quant_method is not None
|
||
assert isinstance(self.quant_method,
|
||
(Fp8LinearMethod, Fp8MoEMethod))
|
||
weight_block_size = self.quant_method.quant_config.weight_block_size
|
||
assert weight_block_size is not None
|
||
block_n, _ = weight_block_size[0] // self.scale_n, weight_block_size[1]
|
||
shard_offset = (
|
||
(sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
|
||
block_n) // tp_size
|
||
shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
|
||
block_n // tp_size)
|
||
else:
|
||
shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
|
||
shard_size = self.output_sizes[loaded_shard_id] // tp_size
|
||
|
||
param.load_merged_column_weight(loaded_weight=loaded_weight,
|
||
shard_id=loaded_shard_id,
|
||
shard_offset=shard_offset,
|
||
shard_size=shard_size)
|
||
|
||
def RowParallelLinear__init__(
|
||
self,
|
||
input_size: int,
|
||
output_size: int,
|
||
bias: bool = True,
|
||
input_is_parallel: bool = True,
|
||
skip_bias_add: bool = False,
|
||
params_dtype: Optional[torch.dtype] = None,
|
||
reduce_results: bool = True,
|
||
quant_config: Optional[QuantizationConfig] = None,
|
||
prefix: str = "",
|
||
*,
|
||
return_bias: bool = True,
|
||
disable_tp: bool = False,
|
||
):
|
||
|
||
# Divide the weight matrix along the first dimension.
|
||
self.tp_rank = (get_tensor_model_parallel_rank()
|
||
if not disable_tp else 0)
|
||
self.tp_size = (get_tensor_model_parallel_world_size()
|
||
if not disable_tp else 1)
|
||
self.input_size_per_partition = divide(input_size, self.tp_size)
|
||
self.output_size_per_partition = output_size
|
||
self.output_partition_sizes = [output_size]
|
||
super(RowParallelLinear, self).__init__(input_size,
|
||
output_size,
|
||
skip_bias_add,
|
||
params_dtype,
|
||
quant_config,
|
||
prefix,
|
||
return_bias=return_bias,
|
||
disable_tp=disable_tp)
|
||
|
||
self.input_is_parallel = input_is_parallel
|
||
self.reduce_results = reduce_results
|
||
|
||
# Divide the weight matrix along the last dimension.
|
||
self.tp_rank = get_tensor_model_parallel_rank()
|
||
self.tp_size = get_tensor_model_parallel_world_size()
|
||
self.input_size_per_partition = divide(input_size, self.tp_size)
|
||
assert self.quant_method is not None
|
||
|
||
self.scale_k = 1 # quant_block_k 128 需要除以 scale_k， 如设置为2 即 quant_block_k 是 64
|
||
self.scale_n = 1
|
||
self.scale_n_slice = 1
|
||
|
||
if quant_config is not None and hasattr(quant_config, "weight_block_size") and quant_config.weight_block_size is not None:
|
||
gcd_value = quant_config.weight_block_size[1]
|
||
import math
|
||
if self.input_size_per_partition % quant_config.weight_block_size[1]:
|
||
gcd_value = math.gcd(self.input_size_per_partition % quant_config.weight_block_size[1], quant_config.weight_block_size[1])
|
||
self.scale_k =self.scale_k * quant_config.weight_block_size[1] // gcd_value
|
||
if output_size % quant_config.weight_block_size[0]:
|
||
gcd_value = math.gcd(output_size % quant_config.weight_block_size[0], quant_config.weight_block_size[0])
|
||
self.scale_n = self.scale_n * quant_config.weight_block_size[0] // gcd_value
|
||
self.scale_n_slice = output_size // gcd_value
|
||
# N = 576, block = 128, n方向scale 扩充需要知道两个信息: 1.拷贝多少份 scale_n; 2. slice 有效的 scale_n_slice
|
||
# scale = [s0,s1,s2,s3,s4] 拷贝scale_n=2份
|
||
# scale = [s0,s0,s1,s1,s2,s2,s3,s3,s4,s4]，slice scale_n_slice=9份 =>[s0,s0,s1,s1,s2,s2,s3,s3,s4]
|
||
|
||
if self.quant_method.__class__.__name__ in ['GPTQLinearMethod']:
|
||
gcd_value = quant_config.group_size
|
||
import math
|
||
if self.input_size_per_partition % quant_config.group_size:
|
||
gcd_value = math.gcd(self.input_size_per_partition % quant_config.group_size, quant_config.group_size)
|
||
self.quant_method.scale_k = self.quant_method.scale_k * quant_config.group_size // gcd_value
|
||
|
||
self.quant_method.create_weights(
|
||
layer=self,
|
||
input_size_per_partition=self.input_size_per_partition,
|
||
output_partition_sizes=[self.output_size],
|
||
input_size=self.input_size,
|
||
output_size=self.output_size,
|
||
params_dtype=self.params_dtype,
|
||
scale_k = self.scale_k,
|
||
scale_n = self.scale_n,
|
||
weight_loader=(
|
||
self.weight_loader_v2 if self.quant_method.__class__.__name__
|
||
in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
|
||
if not reduce_results and (bias and not skip_bias_add):
|
||
raise ValueError("When not reduce the results, adding bias to the "
|
||
"results can lead to incorrect results")
|
||
|
||
if bias:
|
||
self.bias = Parameter(
|
||
torch.empty(self.output_size, dtype=params_dtype))
|
||
set_weight_attrs(self.bias, {
|
||
"output_dim": 0,
|
||
"weight_loader": self.weight_loader,
|
||
})
|
||
else:
|
||
self.register_parameter("bias", None)
|
||
|
||
def RowParallelLinear_weight_loader_v2_vacc(self, param: BasevLLMParameter,
|
||
loaded_weight: torch.Tensor):
|
||
# Special case for loading scales off disk, which often do not
|
||
# have a shape (such as in the case of AutoFP8).
|
||
if len(loaded_weight.shape) == 0:
|
||
assert loaded_weight.numel() == 1
|
||
loaded_weight = loaded_weight.reshape(1)
|
||
if self.quant_method.__class__.__name__ in ['Fp8LinearMethod', 'Fp8MoEMethod'] and torch.finfo(loaded_weight.dtype).bits > 8:
|
||
if self.scale_k > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_k, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,2,0).reshape([loaded_weight.shape[1], -1])
|
||
#[1,n,k] -> [scale_k,n,k] -> [n,k,scale_k] -> [n, k*scale_k]
|
||
|
||
if self.scale_n > 1 and len(loaded_weight.shape) == 2:
|
||
loaded_weight = loaded_weight.unsqueeze(0) #[1,n,k]
|
||
loaded_weight = loaded_weight.expand(self.scale_n, loaded_weight.shape[1], loaded_weight.shape[2]).permute(1,0,2).reshape([-1, loaded_weight.shape[2]])[:self.scale_n_slice]
|
||
#[1,n,k] -> [scale_n,n,k] -> [n,scale_n,k] -> [n*scale_n,k]
|
||
|
||
elif self.quant_method.__class__.__name__ in ['GPTQLinearMethod']:
|
||
# broadcast scale TODO: broadcast zero
|
||
if self.quant_method.scale_k > 1 and len(loaded_weight.shape) == 2 and loaded_weight.dtype in [torch.float16, torch.float32, torch.bfloat16]:
|
||
loaded_weight = loaded_weight.unsqueeze(1) #[k,1,n]
|
||
loaded_weight = loaded_weight.expand(loaded_weight.shape[0], self.quant_method.scale_k, loaded_weight.shape[2]).reshape([-1, loaded_weight.shape[2]])
|
||
#[k,1,n] -> [k,scale_k,n]] -> [k*scale_k, n]
|
||
|
||
param.load_row_parallel_weight(loaded_weight=loaded_weight)
|
||
|
||
class UnquantizedLinearMethod():
|
||
"""Linear method without quantization."""
|
||
|
||
def apply(self,
|
||
layer: torch.nn.Module,
|
||
x: torch.Tensor,
|
||
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||
if bias is not None:
|
||
from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
|
||
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
|
||
|
||
from vllm_vacc.vllm.model_executor.models.memory.memory_recycling import memory_recycler
|
||
parallel_embedding_output = None
|
||
if memory_recycler is not None:
|
||
if memory_recycler.EMBEDDING_OUT_BUFFER.size(0) == x.size(0):
|
||
parallel_embedding_output = memory_recycler.EMBEDDING_OUT_BUFFER
|
||
return torch.mm(x.view(-1, x.shape[-1]), layer.weight.transpose(1,0), out=parallel_embedding_output).view(*(x.shape[:-1]), layer.weight.shape[0]) |