wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
1263 lines
50 KiB
Python
Executable File
1263 lines
50 KiB
Python
Executable File
# SPDX-License-Identifier: MIT AND Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
#
|
|
# Copyright (c) 2026 LightSeek Foundation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
|
|
import torch
|
|
from torch.nn.parameter import Parameter
|
|
|
|
from tokenspeed.runtime.distributed.comm_ops import all_gather, all_reduce
|
|
from tokenspeed.runtime.distributed.utils import divide, split_tensor_along_last_dim
|
|
from tokenspeed.runtime.layers.dense import (
|
|
Fp8LinearMethod,
|
|
Mxfp4LinearMethod,
|
|
Nvfp4LinearMethod,
|
|
UnquantizedLinearMethod,
|
|
W8A8Fp8LinearMethod,
|
|
)
|
|
from tokenspeed.runtime.layers.parameter import (
|
|
BaseWeightParameter,
|
|
BlockQuantScaleParameter,
|
|
PackedColumnParameter,
|
|
PackedWeightParameter,
|
|
PerTensorScaleParameter,
|
|
RowParallelWeightParameter,
|
|
)
|
|
from tokenspeed.runtime.layers.quantization import (
|
|
Fp8Config,
|
|
Mxfp4Config,
|
|
Nvfp4Config,
|
|
W8A8Fp8Config,
|
|
)
|
|
from tokenspeed.runtime.layers.quantization.base_config import (
|
|
QuantizationConfig,
|
|
QuantizeMethodBase,
|
|
)
|
|
from tokenspeed.runtime.layers.quantization.compressed_tensors.compressed_tensors import (
|
|
CompressedTensorsConfig,
|
|
)
|
|
from tokenspeed.runtime.layers.quantization.utils import (
|
|
should_exclude_quant_module,
|
|
should_ignore_quant_layer,
|
|
)
|
|
from tokenspeed.runtime.utils import get_colorful_logger, set_weight_attrs
|
|
|
|
logger = get_colorful_logger(__name__)
|
|
|
|
WEIGHT_LOADER_V2_SUPPORTED = [
|
|
"CompressedTensorsLinearMethod",
|
|
"AWQMarlinLinearMethod",
|
|
"AWQLinearMethod",
|
|
"GPTQMarlinLinearMethod",
|
|
"Fp8LinearMethod",
|
|
"BlockInt8LinearMethod",
|
|
"MarlinLinearMethod",
|
|
"QQQLinearMethod",
|
|
"GPTQMarlin24LinearMethod",
|
|
"TPUInt8LinearMethod",
|
|
"GPTQLinearMethod",
|
|
"IPEXAWQLinearMethod",
|
|
]
|
|
|
|
|
|
def adjust_marlin_shard(param, shard_size, shard_offset):
|
|
marlin_tile_size = getattr(param, "marlin_tile_size", None)
|
|
if marlin_tile_size is None:
|
|
return shard_size, shard_offset
|
|
|
|
return shard_size * marlin_tile_size, shard_offset * marlin_tile_size
|
|
|
|
|
|
def adjust_bitsandbytes_4bit_shard(
|
|
param: Parameter, qkv_offsets: dict[str, tuple[int, int]], loaded_shard_id: str
|
|
) -> tuple[int, int]:
|
|
"""Adjust the quantization offsets and sizes for BitsAndBytes sharding."""
|
|
|
|
total, _ = qkv_offsets["total"]
|
|
orig_offset, orig_size = qkv_offsets[loaded_shard_id]
|
|
|
|
quantized_total = param.data.shape[0]
|
|
quantized_offset = orig_offset * quantized_total // total
|
|
quantized_size = orig_size * quantized_total // total
|
|
|
|
return quantized_size, quantized_offset
|
|
|
|
|
|
def adjust_scalar_to_fused_array(param, loaded_weight, shard_id):
|
|
"""For fused modules (QKV and MLP) we have an array of length
|
|
N that holds 1 scale for each "logical" matrix. So the param
|
|
is an array of length N. The loaded_weight corresponds to
|
|
one of the shards on disk. Here, we slice the param based on
|
|
the shard_id for loading.
|
|
"""
|
|
qkv_idxs = {"q": 0, "k": 1, "v": 2}
|
|
|
|
if isinstance(shard_id, str):
|
|
shard_id = qkv_idxs[shard_id]
|
|
elif not isinstance(shard_id, int):
|
|
raise ValueError(f"Unknown Shard Id {shard_id}")
|
|
|
|
# AutoFP8 scales do not have a shape
|
|
# compressed-tensors scales do have a shape
|
|
if len(loaded_weight.shape) != 0:
|
|
assert loaded_weight.shape[0] == 1
|
|
loaded_weight = loaded_weight[0]
|
|
|
|
return param[shard_id], loaded_weight
|
|
|
|
|
|
class LinearBase(torch.nn.Module):
|
|
"""Base linear layer.
|
|
|
|
Args:
|
|
input_size: input dimension of the linear layer.
|
|
output_size: output dimension of the linear layer.
|
|
bias: If true, add bias.
|
|
skip_bias_add: If true, skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
override_kernel_name: Optional kernel name passed down to the
|
|
quant method's underlying ``tokenspeed_kernel.mm`` dispatch
|
|
(e.g. ``"cublaslt_mm_nvfp4"``). Lets the model force a
|
|
specific kernel for a particular layer.
|
|
interleave_linear_and_gate: If true, quantized post-load processing
|
|
prepares a 64-row linear/gate interleaved weight view for
|
|
fused GEMM+SwiGLU kernels.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
override_kernel_name: str | None = None,
|
|
interleave_linear_and_gate: bool = False,
|
|
):
|
|
super().__init__()
|
|
|
|
# Keep input parameters
|
|
self.input_size = input_size
|
|
self.output_size = output_size
|
|
self.skip_bias_add = skip_bias_add
|
|
if params_dtype is None:
|
|
params_dtype = torch.get_default_dtype()
|
|
self.params_dtype = params_dtype
|
|
self.prefix = prefix
|
|
self.override_kernel_name = override_kernel_name
|
|
self.interleave_linear_and_gate = interleave_linear_and_gate
|
|
|
|
self.quant_config = quant_config
|
|
if quant_config is None or should_ignore_quant_layer(
|
|
prefix=prefix, ignored_layers=quant_config.ignored_layers
|
|
):
|
|
self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod()
|
|
elif isinstance(quant_config, Nvfp4Config):
|
|
# For NVFP4, excluded layers use unquantized (bf16)
|
|
if should_exclude_quant_module(prefix, quant_config.exclude_modules):
|
|
self.quant_method = UnquantizedLinearMethod()
|
|
else:
|
|
self.quant_method = Nvfp4LinearMethod(quant_config)
|
|
elif isinstance(quant_config, Mxfp4Config):
|
|
if getattr(quant_config, "use_dynamic_mxfp4_activations", False):
|
|
self.quant_method = Mxfp4LinearMethod(quant_config)
|
|
else:
|
|
# Existing MXFP4 support applies to MoE weights; dense weights
|
|
# remain unquantized unless the checkpoint stores dense MXFP4.
|
|
self.quant_method = UnquantizedLinearMethod()
|
|
elif isinstance(quant_config, CompressedTensorsConfig):
|
|
self.quant_method = quant_config.get_quant_method(self, prefix)
|
|
else:
|
|
if isinstance(quant_config, Fp8Config):
|
|
self.quant_method = Fp8LinearMethod(quant_config)
|
|
if isinstance(quant_config, W8A8Fp8Config):
|
|
self.quant_method = W8A8Fp8LinearMethod(quant_config)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
raise NotImplementedError
|
|
|
|
|
|
class ReplicatedLinear(LinearBase):
|
|
"""Replicated linear layer.
|
|
|
|
Args:
|
|
input_size: input dimension of the linear layer.
|
|
output_size: output dimension of the linear layer.
|
|
bias: If true, add bias.
|
|
skip_bias_add: If true, skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
prefix: The name of the layer in the state dict, including all parents
|
|
(e.g. model.layers.0.qkv_proj)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
bias: bool = True,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
):
|
|
super().__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.quant_method.create_weights(
|
|
layer=self,
|
|
input_size_per_partition=self.input_size,
|
|
output_partition_sizes=[self.output_size],
|
|
input_size=self.input_size,
|
|
output_size=self.output_size,
|
|
params_dtype=self.params_dtype,
|
|
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 weight_loader(
|
|
self,
|
|
param: Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
shard_id=None,
|
|
begin_size=None,
|
|
):
|
|
# 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 begin_size is not None:
|
|
shard_size = loaded_weight.shape[0]
|
|
param[begin_size : begin_size + shard_size].data.copy_(loaded_weight)
|
|
elif shard_id is not None:
|
|
shard_size = loaded_weight.shape[0]
|
|
param[shard_id * shard_size : (shard_id + 1) * shard_size].data.copy_(
|
|
loaded_weight
|
|
)
|
|
else:
|
|
assert param.size() == loaded_weight.size()
|
|
param.data.copy_(loaded_weight)
|
|
|
|
def forward(
|
|
self, x: torch.Tensor, block_scale=None, output_dtype=None
|
|
) -> torch.Tensor:
|
|
bias = self.bias if not self.skip_bias_add else None
|
|
assert self.quant_method is not None
|
|
if block_scale is not None:
|
|
# Note: block_scale is not None means flashinfer reduce-scatter fusion is used for fp8 block quant
|
|
# in this case, the input_ is already quantized to a fp8 tensor
|
|
output = self.quant_method.apply(self, x, bias, block_scale, output_dtype)
|
|
else:
|
|
output = self.quant_method.apply(self, x, bias)
|
|
output_bias = self.bias if self.skip_bias_add else None
|
|
return output, output_bias
|
|
|
|
def extra_repr(self) -> str:
|
|
s = f"in_features={self.input_size}"
|
|
s += f", output_features={self.output_size}"
|
|
s += f", bias={self.bias is not None}"
|
|
return s
|
|
|
|
|
|
class ColumnParallelLinear(LinearBase):
|
|
"""Linear layer with column parallelism.
|
|
|
|
The linear layer is defined as Y = XA + b. A is parallelized along
|
|
its second dimension as A = [A_1, ..., A_p].
|
|
|
|
Args:
|
|
input_size: first dimension of matrix A.
|
|
output_size: second dimension of matrix A.
|
|
bias: If true, add bias.
|
|
gather_output: If true, call all-gather on output and make Y available
|
|
to all GPUs, otherwise, every GPU will have its output
|
|
which is Y_i = XA_i
|
|
skip_bias_add: This was added to enable performance optimizations where
|
|
bias can be fused with other element-wise operations. we
|
|
skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
output_sizes: list of output sizes packed into one output, like for QKV
|
|
the list would be size 3.
|
|
prefix: The name of the layer in the state dict, including all parents
|
|
(e.g. model.layers.0.qkv_proj)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
bias: bool = True,
|
|
gather_output: bool = False,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
output_sizes: list[int] | None = None,
|
|
prefix: str = "",
|
|
tp_rank: int | None = None,
|
|
tp_size: int | None = None,
|
|
tp_group: tuple[int, ...] | None = None,
|
|
use_presharded_weights: bool = False,
|
|
override_kernel_name: str | None = None,
|
|
interleave_linear_and_gate: bool = False,
|
|
):
|
|
super().__init__(
|
|
input_size,
|
|
output_size,
|
|
skip_bias_add,
|
|
params_dtype,
|
|
quant_config,
|
|
prefix,
|
|
override_kernel_name,
|
|
interleave_linear_and_gate,
|
|
)
|
|
self.gather_output = gather_output
|
|
self.use_presharded_weights = use_presharded_weights
|
|
assert self.quant_method is not None
|
|
|
|
if tp_rank is None:
|
|
assert tp_size is None
|
|
assert tp_group is None
|
|
tp_rank, tp_size = 0, 1
|
|
assert 0 <= tp_rank < tp_size
|
|
assert tp_size == 1 or tp_group is not None
|
|
self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group
|
|
|
|
self.output_size_per_partition = divide(self.output_size, self.tp_size)
|
|
if output_sizes is None:
|
|
self.output_sizes = [self.output_size]
|
|
self.output_partition_sizes = [self.output_size_per_partition]
|
|
else:
|
|
self.output_sizes = output_sizes
|
|
# If QKV or MergedColumn, use output size of each partition.
|
|
self.output_partition_sizes = [
|
|
divide(output_size, self.tp_size) for output_size in self.output_sizes
|
|
]
|
|
|
|
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,
|
|
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)
|
|
|
|
def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
|
|
output_dim = getattr(param, "output_dim", None)
|
|
|
|
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
|
|
|
|
param_data = param.data
|
|
# bitsandbytes loads the weights of the specific portion
|
|
# no need to narrow here
|
|
if output_dim is not None and not use_bitsandbytes_4bit:
|
|
shard_size = param_data.shape[output_dim]
|
|
start_idx = self.tp_rank * shard_size
|
|
if not self.use_presharded_weights:
|
|
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
|
|
|
|
# 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:
|
|
loaded_weight = loaded_weight.reshape(1)
|
|
|
|
assert param_data.shape == loaded_weight.shape
|
|
param_data.copy_(loaded_weight)
|
|
|
|
def 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)
|
|
|
|
from tokenspeed.runtime.layers.parameter import _ColumnParallelWeightParameter
|
|
|
|
if isinstance(param, _ColumnParallelWeightParameter):
|
|
param.load_column_parallel_weight(
|
|
loaded_weight,
|
|
tp_rank=self.tp_rank,
|
|
use_presharded_weights=self.use_presharded_weights,
|
|
)
|
|
else:
|
|
# Some DeepSeek V3 AWQ checkpoints still reach the generic column
|
|
# loader path instead of the _ColumnParallelWeightParameter specialization.
|
|
param.load_column_parallel_weight(loaded_weight)
|
|
|
|
def forward(self, input_, block_scale=None, output_dtype=None):
|
|
bias = self.bias if not self.skip_bias_add else None
|
|
|
|
# Matrix multiply.
|
|
assert self.quant_method is not None
|
|
if block_scale is not None:
|
|
# Note: block_scale is not None means flashinfer all-reduce fusion is used for fp8 block quant
|
|
# in this case, the input_ is already quantized to a fp8 tensor
|
|
output_parallel = self.quant_method.apply(
|
|
self, input_, bias, block_scale, output_dtype
|
|
)
|
|
else:
|
|
output_parallel = self.quant_method.apply(self, input_, bias)
|
|
if self.gather_output and self.tp_size > 1:
|
|
# All-gather across the partitions.
|
|
output = all_gather(output_parallel, self.tp_group, dim=-1)
|
|
else:
|
|
output = output_parallel
|
|
output_bias = self.bias if self.skip_bias_add else None
|
|
return output, output_bias
|
|
|
|
def extra_repr(self) -> str:
|
|
s = f"in_features={self.input_size}"
|
|
s += f", output_features={self.output_size_per_partition}"
|
|
s += f", bias={self.bias is not None}"
|
|
s += f", tp_size={self.tp_size}"
|
|
s += f", gather_output={self.gather_output}"
|
|
return s
|
|
|
|
|
|
class MergedColumnParallelLinear(ColumnParallelLinear):
|
|
"""Packed linear layers with column parallelism.
|
|
|
|
Similar to ColumnParallelLinear, but the weight matrix is concatenated
|
|
along the output dimension. When the weight matrix is loaded, the
|
|
different partitions are sharded separately.
|
|
|
|
Args:
|
|
input_size: input dimension of the linear layer.
|
|
output_sizes: list of output dimensions of the linear layer.
|
|
bias: If true, add bias.
|
|
gather_output: If true, call all-gather on output and make the output
|
|
available to all GPUs, otherwise, every GPU will have
|
|
its own output.
|
|
skip_bias_add: This was added to enable performance optimizations where
|
|
bias can be fused with other element-wise operations. we
|
|
skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
prefix: The name of the layer in the state dict, including all parents
|
|
(e.g. model.layers.0.qkv_proj)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_sizes: list[int],
|
|
bias: bool = True,
|
|
gather_output: bool = False,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
tp_rank: int | None = None,
|
|
tp_size: int | None = None,
|
|
tp_group: tuple[int, ...] | None = None,
|
|
use_presharded_weights: bool = False,
|
|
override_kernel_name: str | None = None,
|
|
interleave_linear_and_gate: bool = False,
|
|
):
|
|
super().__init__(
|
|
input_size=input_size,
|
|
output_size=sum(output_sizes),
|
|
bias=bias,
|
|
gather_output=gather_output,
|
|
skip_bias_add=skip_bias_add,
|
|
params_dtype=params_dtype,
|
|
quant_config=quant_config,
|
|
output_sizes=output_sizes,
|
|
prefix=prefix,
|
|
tp_rank=tp_rank,
|
|
tp_size=tp_size,
|
|
tp_group=tp_group,
|
|
use_presharded_weights=use_presharded_weights,
|
|
override_kernel_name=override_kernel_name,
|
|
interleave_linear_and_gate=interleave_linear_and_gate,
|
|
)
|
|
|
|
def weight_loader(
|
|
self,
|
|
param: Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
loaded_shard_id: int | None = None,
|
|
):
|
|
param_data = param.data
|
|
output_dim = getattr(param, "output_dim", None)
|
|
# Special case for AQLM codebooks.
|
|
is_metadata = getattr(param, "is_metadata", False)
|
|
# Special case for per-tensor scale to load scalar into fused array.
|
|
needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
|
|
|
|
if loaded_shard_id is None:
|
|
# Loaded weight is already fused on disk (qkv/mlp).
|
|
if output_dim is None:
|
|
if needs_scalar_to_array:
|
|
param_data, loaded_weight = adjust_scalar_to_fused_array(
|
|
param_data, loaded_weight, 0
|
|
)
|
|
|
|
assert param_data.shape == loaded_weight.shape
|
|
param_data.copy_(loaded_weight)
|
|
return
|
|
current_shard_offset = 0
|
|
shard_offsets: list[tuple[int, int, int]] = []
|
|
for i, output_size in enumerate(self.output_sizes):
|
|
shard_offsets.append((i, current_shard_offset, output_size))
|
|
current_shard_offset += output_size
|
|
packed_dim = getattr(param, "packed_dim", None)
|
|
for shard_id, shard_offset, shard_size in shard_offsets:
|
|
# Special case for Quantization.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
if packed_dim == output_dim:
|
|
shard_size = shard_size // param.pack_factor
|
|
shard_offset = shard_offset // param.pack_factor
|
|
# Special case for Marlin.
|
|
shard_size, shard_offset = adjust_marlin_shard(
|
|
param, shard_size, shard_offset
|
|
)
|
|
|
|
loaded_weight_shard = loaded_weight.narrow(
|
|
output_dim, shard_offset, shard_size
|
|
)
|
|
self.weight_loader(param, loaded_weight_shard, shard_id)
|
|
return
|
|
|
|
assert loaded_shard_id < len(self.output_sizes)
|
|
if output_dim is not None:
|
|
shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
|
|
shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
|
|
# Special case for quantization.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
packed_dim = getattr(param, "packed_dim", None)
|
|
if packed_dim == output_dim:
|
|
shard_size = shard_size // param.pack_factor
|
|
shard_offset = shard_offset // param.pack_factor
|
|
# Special case for Marlin.
|
|
shard_size, shard_offset = adjust_marlin_shard(
|
|
param, shard_size, shard_offset
|
|
)
|
|
|
|
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
|
|
if use_bitsandbytes_4bit:
|
|
shard_size = loaded_weight.shape[output_dim]
|
|
shard_offset = loaded_weight.shape[output_dim] * loaded_shard_id
|
|
|
|
param_data = param_data.narrow(output_dim, shard_offset, shard_size)
|
|
start_idx = self.tp_rank * shard_size
|
|
# bitsandbytes loads the weights of the specific portion
|
|
# no need to narrow here
|
|
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
|
|
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
|
|
# Special case for AQLM codebooks.
|
|
elif is_metadata:
|
|
# metadata indicates fixed size concatenated along dim 0
|
|
shard_size = loaded_weight.shape[0]
|
|
shard_offset = loaded_shard_id * shard_size
|
|
param_data = param_data.narrow(0, shard_offset, shard_size)
|
|
|
|
# Special case for per-tensor scales in fused case.
|
|
elif needs_scalar_to_array:
|
|
param_data, loaded_weight = adjust_scalar_to_fused_array(
|
|
param_data, loaded_weight, loaded_shard_id
|
|
)
|
|
|
|
else:
|
|
ignore_warning = getattr(param, "ignore_warning", False)
|
|
if not ignore_warning:
|
|
logger.warning(
|
|
"Loading a weight without `output_dim` attribute in "
|
|
"MergedColumnParallelLinear, assume the weight is "
|
|
"the same for all partitions."
|
|
)
|
|
|
|
assert param_data.shape == loaded_weight.shape
|
|
param_data.copy_(loaded_weight)
|
|
|
|
def _load_fused_module_from_checkpoint(
|
|
self, param: BaseWeightParameter, loaded_weight: torch.Tensor
|
|
):
|
|
"""
|
|
Handle special case for models where MLP layers are already
|
|
fused on disk. In this case, we have no shard id. This function
|
|
determmines the shard id by splitting these layers and then calls
|
|
the weight loader using the shard id.
|
|
|
|
An example of a model with these fused layers:
|
|
https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
|
|
"""
|
|
|
|
current_shard_offset = 0
|
|
shard_offsets: list[tuple[int, int, int]] = []
|
|
for i, output_size in enumerate(self.output_sizes):
|
|
shard_offsets.append((i, current_shard_offset, output_size))
|
|
current_shard_offset += output_size
|
|
|
|
for shard_id, shard_offset, shard_size in shard_offsets:
|
|
# Special case for Quantization.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
if (
|
|
isinstance(param, (PackedColumnParameter, PackedWeightParameter))
|
|
and param.packed_dim == param.output_dim
|
|
):
|
|
shard_size, shard_offset = param.adjust_shard_indexes_for_packing(
|
|
shard_size=shard_size, shard_offset=shard_offset
|
|
)
|
|
# Special case for block-wise quantization scales.
|
|
# The scale tensor is smaller than the weight tensor by a factor
|
|
# of block_n, so we need to adjust offset and size accordingly.
|
|
elif isinstance(param, BlockQuantScaleParameter):
|
|
weight_block_size = self.quant_method.quant_config.weight_block_size
|
|
block_n = weight_block_size[0]
|
|
shard_offset = (shard_offset + block_n - 1) // block_n
|
|
shard_size = (shard_size + block_n - 1) // block_n
|
|
|
|
loaded_weight_shard = loaded_weight.narrow(
|
|
param.output_dim, shard_offset, shard_size
|
|
)
|
|
self.weight_loader_v2(param, loaded_weight_shard, shard_id)
|
|
|
|
def weight_loader_v2(
|
|
self,
|
|
param: BaseWeightParameter,
|
|
loaded_weight: torch.Tensor,
|
|
loaded_shard_id: int | None = None,
|
|
):
|
|
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 (RowParallelWeightParameter, BaseWeightParameter):
|
|
param.load_merged_column_weight(loaded_weight=loaded_weight)
|
|
return
|
|
self._load_fused_module_from_checkpoint(param, loaded_weight)
|
|
return
|
|
|
|
assert loaded_shard_id < len(self.output_sizes)
|
|
|
|
if isinstance(param, BlockQuantScaleParameter):
|
|
weight_block_size = self.quant_method.quant_config.weight_block_size
|
|
block_n, _ = weight_block_size[0], weight_block_size[1]
|
|
shard_offset = (
|
|
(sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) // block_n
|
|
) // self.tp_size
|
|
shard_size = (
|
|
(self.output_sizes[loaded_shard_id] + block_n - 1)
|
|
// block_n
|
|
// self.tp_size
|
|
)
|
|
else:
|
|
shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
|
|
shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
|
|
|
|
param.load_merged_column_weight(
|
|
loaded_weight=loaded_weight,
|
|
shard_id=loaded_shard_id,
|
|
tp_rank=self.tp_rank,
|
|
shard_offset=shard_offset,
|
|
shard_size=shard_size,
|
|
use_presharded_weights=self.use_presharded_weights,
|
|
)
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
return None
|
|
|
|
|
|
class QKVParallelLinear(ColumnParallelLinear):
|
|
"""Linear layers for the attention's QKV transformation.
|
|
|
|
Linear layers for the linear transformation of the query, key, and value
|
|
vectors in the attention layer. The weight matrix is concatenated along
|
|
the output dimension. The layer is parallelized along the head dimension.
|
|
When the number of key/value heads is smaller than the number of query
|
|
heads (e.g., multi-query/grouped-query attention), the key/value head may
|
|
be replicated while the query heads are partitioned.
|
|
|
|
Args:
|
|
hidden_size: input hidden state size of the transformer.
|
|
head_size: size of each attention head.
|
|
total_num_heads: total number of attention query heads.
|
|
total_num_kv_heads: total number of attention key/value heads. If
|
|
None, assume total_num_kv_heads = total_num_heads.
|
|
bias: If true, add bias.
|
|
skip_bias_add: This was added to enable performance optimizations where
|
|
bias can be fused with other element-wise operations. we
|
|
skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
prefix: The name of the layer in the state dict, including all parents
|
|
(e.g. model.layers.0.qkv_proj)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
head_size: int,
|
|
total_num_heads: int,
|
|
total_num_kv_heads: int | None = None,
|
|
bias: bool = True,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
tp_rank: int | None = None,
|
|
tp_size: int | None = None,
|
|
tp_group: tuple[int, ...] | None = None,
|
|
load_presharded_attn: bool = False,
|
|
):
|
|
self.hidden_size = hidden_size
|
|
self.head_size = head_size
|
|
self.total_num_heads = total_num_heads
|
|
if total_num_kv_heads is None:
|
|
total_num_kv_heads = total_num_heads
|
|
self.total_num_kv_heads = total_num_kv_heads
|
|
|
|
tp_size = 1 if tp_size is None else tp_size
|
|
self.num_heads = divide(self.total_num_heads, tp_size)
|
|
|
|
if tp_size >= self.total_num_kv_heads:
|
|
self.num_kv_heads = 1
|
|
self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
|
|
else:
|
|
self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
|
|
self.num_kv_head_replicas = 1
|
|
input_size = self.hidden_size
|
|
output_size = (
|
|
(self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size
|
|
)
|
|
output_sizes = [
|
|
self.num_heads * self.head_size * tp_size, # q_proj
|
|
self.num_kv_heads * self.head_size * tp_size, # k_proj
|
|
self.num_kv_heads * self.head_size * tp_size, # v_proj
|
|
]
|
|
super().__init__(
|
|
input_size=input_size,
|
|
output_size=output_size,
|
|
bias=bias,
|
|
gather_output=False,
|
|
skip_bias_add=skip_bias_add,
|
|
params_dtype=params_dtype,
|
|
quant_config=quant_config,
|
|
output_sizes=output_sizes,
|
|
prefix=prefix,
|
|
tp_rank=tp_rank,
|
|
tp_size=tp_size,
|
|
tp_group=tp_group,
|
|
use_presharded_weights=load_presharded_attn,
|
|
)
|
|
|
|
def _get_shard_offset_mapping(self, loaded_shard_id: str):
|
|
shard_offset_mapping = {
|
|
"q": 0,
|
|
"k": self.num_heads * self.head_size,
|
|
"v": (self.num_heads + self.num_kv_heads) * self.head_size,
|
|
"total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size,
|
|
}
|
|
return shard_offset_mapping.get(loaded_shard_id)
|
|
|
|
def _get_shard_size_mapping(self, loaded_shard_id: str):
|
|
shard_size_mapping = {
|
|
"q": self.num_heads * self.head_size,
|
|
"k": self.num_kv_heads * self.head_size,
|
|
"v": self.num_kv_heads * self.head_size,
|
|
}
|
|
return shard_size_mapping.get(loaded_shard_id)
|
|
|
|
def _load_fused_module_from_checkpoint(
|
|
self, param: BaseWeightParameter, loaded_weight: torch.Tensor
|
|
):
|
|
"""
|
|
Handle special case for models where QKV layers are already
|
|
fused on disk. In this case, we have no shard id. This function
|
|
determmines the shard id by splitting these layers and then calls
|
|
the weight loader using the shard id.
|
|
|
|
An example of a model with these fused layers:
|
|
https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
|
|
"""
|
|
shard_offsets = [
|
|
# (shard_id, shard_offset, shard_size)
|
|
("q", 0, self.total_num_heads * self.head_size),
|
|
(
|
|
"k",
|
|
self.total_num_heads * self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
(
|
|
"v",
|
|
(self.total_num_heads + self.total_num_kv_heads) * self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
]
|
|
|
|
for shard_id, shard_offset, shard_size in shard_offsets:
|
|
# Special case for Quantization.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
if (
|
|
isinstance(param, (PackedColumnParameter, PackedWeightParameter))
|
|
and param.packed_dim == param.output_dim
|
|
):
|
|
shard_size, shard_offset = param.adjust_shard_indexes_for_packing(
|
|
shard_size=shard_size, shard_offset=shard_offset
|
|
)
|
|
|
|
if not self.use_presharded_weights:
|
|
loaded_weight_shard = loaded_weight.narrow(
|
|
param.output_dim, shard_offset, shard_size
|
|
)
|
|
self.weight_loader_v2(param, loaded_weight_shard, shard_id)
|
|
|
|
def weight_loader_v2(
|
|
self,
|
|
param: BaseWeightParameter,
|
|
loaded_weight: torch.Tensor,
|
|
loaded_shard_id: str | None = None,
|
|
):
|
|
if loaded_shard_id is None: # special case for certain models
|
|
if isinstance(param, PerTensorScaleParameter):
|
|
param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0)
|
|
return
|
|
elif type(param) in (RowParallelWeightParameter, BaseWeightParameter):
|
|
param.load_qkv_weight(loaded_weight=loaded_weight)
|
|
return
|
|
self._load_fused_module_from_checkpoint(param, loaded_weight)
|
|
return
|
|
|
|
assert loaded_shard_id in ["q", "k", "v"]
|
|
|
|
shard_offset = self._get_shard_offset_mapping(loaded_shard_id)
|
|
shard_size = self._get_shard_size_mapping(loaded_shard_id)
|
|
|
|
if isinstance(param, BlockQuantScaleParameter):
|
|
weight_block_size = self.quant_method.quant_config.weight_block_size
|
|
block_n, _ = weight_block_size[0], weight_block_size[1]
|
|
shard_offset = (shard_offset + block_n - 1) // block_n
|
|
shard_size = (shard_size + block_n - 1) // block_n
|
|
|
|
param.load_qkv_weight(
|
|
loaded_weight=loaded_weight,
|
|
num_heads=self.num_kv_head_replicas,
|
|
shard_id=loaded_shard_id,
|
|
shard_offset=shard_offset,
|
|
shard_size=shard_size,
|
|
tp_rank=self.tp_rank,
|
|
use_presharded_weights=self.use_presharded_weights,
|
|
)
|
|
|
|
def weight_loader(
|
|
self,
|
|
param: Parameter,
|
|
loaded_weight: torch.Tensor,
|
|
loaded_shard_id: str | None = None,
|
|
):
|
|
param_data = param.data
|
|
output_dim = getattr(param, "output_dim", None)
|
|
# Special case for AQLM codebooks.
|
|
is_metadata = getattr(param, "is_metadata", False)
|
|
|
|
# Special case for per-tensor scales in fused case.
|
|
needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
|
|
|
|
if loaded_shard_id is None:
|
|
# Loaded weight is already fused on disk (qkv/mlp).
|
|
if output_dim is None:
|
|
if needs_scalar_to_array:
|
|
param_data, loaded_weight = adjust_scalar_to_fused_array(
|
|
param_data, loaded_weight, 0
|
|
)
|
|
|
|
assert param_data.shape == loaded_weight.shape
|
|
param_data.copy_(loaded_weight)
|
|
return
|
|
shard_offsets = [
|
|
# (shard_id, shard_offset, shard_size)
|
|
("q", 0, self.total_num_heads * self.head_size),
|
|
(
|
|
"k",
|
|
self.total_num_heads * self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
(
|
|
"v",
|
|
(self.total_num_heads + self.total_num_kv_heads) * self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
]
|
|
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
|
|
|
|
packed_dim = getattr(param, "packed_dim", None)
|
|
for shard_id, shard_offset, shard_size in shard_offsets:
|
|
# Special case for Quantized Weights.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
if packed_dim == output_dim:
|
|
shard_size = shard_size // param.pack_factor
|
|
shard_offset = shard_offset // param.pack_factor
|
|
|
|
# Special case for Marlin.
|
|
shard_size, shard_offset = adjust_marlin_shard(
|
|
param, shard_size, shard_offset
|
|
)
|
|
|
|
if use_bitsandbytes_4bit:
|
|
orig_qkv_offsets = {
|
|
"q": (0, self.total_num_heads * self.head_size),
|
|
"k": (
|
|
self.total_num_heads * self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
"v": (
|
|
(self.total_num_heads + self.total_num_kv_heads)
|
|
* self.head_size,
|
|
self.total_num_kv_heads * self.head_size,
|
|
),
|
|
"total": (
|
|
(self.total_num_heads + 2 * self.total_num_kv_heads)
|
|
* self.head_size,
|
|
0,
|
|
),
|
|
}
|
|
|
|
shard_size, shard_offset = adjust_bitsandbytes_4bit_shard(
|
|
param, orig_qkv_offsets, shard_id
|
|
)
|
|
|
|
if not self.use_presharded_weights:
|
|
loaded_weight_shard = loaded_weight.narrow(
|
|
output_dim, shard_offset, shard_size
|
|
)
|
|
self.weight_loader(param, loaded_weight_shard, shard_id)
|
|
return
|
|
|
|
assert loaded_shard_id in ["q", "k", "v"]
|
|
|
|
# If output dim is defined, use the default loading process.
|
|
if output_dim is not None:
|
|
if loaded_shard_id == "q":
|
|
shard_offset = 0
|
|
shard_size = self.num_heads * self.head_size
|
|
elif loaded_shard_id == "k":
|
|
shard_offset = self.num_heads * self.head_size
|
|
shard_size = self.num_kv_heads * self.head_size
|
|
elif loaded_shard_id == "v":
|
|
shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size
|
|
shard_size = self.num_kv_heads * self.head_size
|
|
# Special case for Quantized Weights.
|
|
# If quantized, we need to adjust the offset and size to account
|
|
# for the packing.
|
|
packed_dim = getattr(param, "packed_dim", None)
|
|
if packed_dim == output_dim:
|
|
shard_size = shard_size // param.pack_factor
|
|
shard_offset = shard_offset // param.pack_factor
|
|
|
|
# Special case for Marlin.
|
|
shard_size, shard_offset = adjust_marlin_shard(
|
|
param, shard_size, shard_offset
|
|
)
|
|
|
|
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
|
|
if use_bitsandbytes_4bit:
|
|
orig_qkv_offsets = {
|
|
"q": (0, self.num_heads * self.head_size),
|
|
"k": (
|
|
self.num_heads * self.head_size,
|
|
self.num_kv_heads * self.head_size,
|
|
),
|
|
"v": (
|
|
(self.num_heads + self.num_kv_heads) * self.head_size,
|
|
self.num_kv_heads * self.head_size,
|
|
),
|
|
"total": (
|
|
(self.num_heads + 2 * self.num_kv_heads) * self.head_size,
|
|
0,
|
|
),
|
|
}
|
|
shard_size, shard_offset = adjust_bitsandbytes_4bit_shard(
|
|
param, orig_qkv_offsets, loaded_shard_id
|
|
)
|
|
|
|
param_data = param_data.narrow(output_dim, shard_offset, shard_size)
|
|
if loaded_shard_id == "q":
|
|
shard_id = self.tp_rank
|
|
else:
|
|
shard_id = self.tp_rank // self.num_kv_head_replicas
|
|
start_idx = shard_id * shard_size
|
|
|
|
# bitsandbytes loads the weights of the specific portion
|
|
# no need to narrow here
|
|
if not use_bitsandbytes_4bit and not self.use_presharded_weights:
|
|
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
|
|
|
|
# Special case for for AQLM codebooks.
|
|
elif is_metadata:
|
|
# metadata indicates fixed size concatenated along dim 0
|
|
shard_size = loaded_weight.shape[0]
|
|
shard_index = ["q", "k", "v"].index(loaded_shard_id)
|
|
param_data = param_data.narrow(0, shard_index * shard_size, shard_size)
|
|
# Special case for per-tensor scales in fused case.
|
|
elif needs_scalar_to_array:
|
|
param_data, loaded_weight = adjust_scalar_to_fused_array(
|
|
param_data, loaded_weight, loaded_shard_id
|
|
)
|
|
else:
|
|
ignore_warning = getattr(param, "ignore_warning", False)
|
|
if not ignore_warning:
|
|
logger.warning(
|
|
"Loading a weight without `output_dim` attribute in "
|
|
"QKVParallelLinear, assume the weight is the same "
|
|
"for all partitions."
|
|
)
|
|
|
|
assert param_data.shape == loaded_weight.shape
|
|
param_data.copy_(loaded_weight)
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
return None
|
|
|
|
|
|
class RowParallelLinear(LinearBase):
|
|
"""Linear layer with row parallelism.
|
|
|
|
The linear layer is defined as Y = XA + b. A is parallelized along
|
|
its first dimension and X along its second dimension as:
|
|
- -
|
|
| A_1 |
|
|
| . |
|
|
A = | . | X = [X_1, ..., X_p]
|
|
| . |
|
|
| A_p |
|
|
- -
|
|
Arguments:
|
|
input_size: first dimension of matrix A.
|
|
output_size: second dimension of matrix A.
|
|
bias: If true, add bias. Note that bias is not parallelized.
|
|
input_is_parallel: If true, we assume that the input is already
|
|
split across the GPUs and we do not split
|
|
again.
|
|
skip_bias_add: This was added to enable performance optimization where
|
|
bias can be fused with other element-wise operations.
|
|
We skip adding bias but instead return it.
|
|
params_dtype: Data type for the parameters.
|
|
quant_config: Quantization configure.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_size: int,
|
|
output_size: int,
|
|
bias: bool = True,
|
|
input_is_parallel: bool = True,
|
|
skip_bias_add: bool = False,
|
|
params_dtype: torch.dtype | None = None,
|
|
reduce_results: bool = True,
|
|
quant_config: QuantizationConfig | None = None,
|
|
prefix: str = "",
|
|
tp_rank: int | None = None,
|
|
tp_size: int | None = None,
|
|
tp_group: tuple[int, ...] | None = None,
|
|
use_presharded_weights: bool = False,
|
|
override_kernel_name: str | None = None,
|
|
interleave_linear_and_gate: bool = False,
|
|
):
|
|
super().__init__(
|
|
input_size,
|
|
output_size,
|
|
skip_bias_add,
|
|
params_dtype,
|
|
quant_config,
|
|
prefix,
|
|
override_kernel_name,
|
|
interleave_linear_and_gate,
|
|
)
|
|
|
|
self.input_is_parallel = input_is_parallel
|
|
self.reduce_results = reduce_results
|
|
assert self.quant_method is not None
|
|
self.use_presharded_weights = use_presharded_weights
|
|
|
|
if tp_rank is None:
|
|
assert tp_size is None
|
|
assert tp_group is None
|
|
tp_rank, tp_size = 0, 1
|
|
assert 0 <= tp_rank < tp_size
|
|
assert tp_size == 1 or tp_group is not None
|
|
self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group
|
|
|
|
self.input_size_per_partition = divide(input_size, self.tp_size)
|
|
|
|
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,
|
|
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, dtype=params_dtype))
|
|
set_weight_attrs(
|
|
self.bias,
|
|
{
|
|
"output_dim": 0,
|
|
"weight_loader": self.weight_loader,
|
|
},
|
|
)
|
|
else:
|
|
self.register_parameter("bias", None)
|
|
|
|
def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
|
|
input_dim = getattr(param, "input_dim", None)
|
|
use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)
|
|
|
|
param_data = param.data
|
|
# bitsandbytes loads the weights of the specific portion
|
|
# no need to narrow here
|
|
if (
|
|
input_dim is not None
|
|
and not use_bitsandbytes_4bit
|
|
and not self.use_presharded_weights
|
|
):
|
|
shard_size = param_data.shape[input_dim]
|
|
start_idx = self.tp_rank * shard_size
|
|
loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size)
|
|
|
|
# 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:
|
|
loaded_weight = loaded_weight.reshape(1)
|
|
|
|
assert (
|
|
param_data.shape == loaded_weight.shape
|
|
), f"{param_data.shape=} {loaded_weight.shape=}"
|
|
param_data.copy_(loaded_weight)
|
|
|
|
def weight_loader_v2(self, param: BaseWeightParameter, 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 isinstance(param, RowParallelWeightParameter):
|
|
# This `BaseWeightParameter` is defined in tokenspeed/runtime/layers/parameter.py,
|
|
# It supports additional parameters like tp_rank and use_presharded_weights.
|
|
param.load_row_parallel_weight(
|
|
loaded_weight,
|
|
tp_rank=self.tp_rank,
|
|
use_presharded_weights=self.use_presharded_weights,
|
|
)
|
|
else:
|
|
# Generic parameters do not support extra loading options.
|
|
param.load_row_parallel_weight(loaded_weight)
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
return None
|
|
|
|
def forward(self, input_, scale=None):
|
|
if self.input_is_parallel:
|
|
input_parallel = input_
|
|
else:
|
|
splitted_input = split_tensor_along_last_dim(
|
|
input_, num_partitions=self.tp_size
|
|
)
|
|
input_parallel = splitted_input[self.tp_rank].contiguous()
|
|
|
|
# Matrix multiply.
|
|
assert self.quant_method is not None
|
|
# Only fuse bias add into GEMM for rank 0 (this ensures that
|
|
# bias will not get added more than once in TP>1 case)
|
|
bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
|
|
|
|
if scale is not None:
|
|
output_parallel = self.quant_method.apply(
|
|
self, input_parallel, bias_, scale, torch.bfloat16
|
|
)
|
|
else:
|
|
output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
|
|
if self.reduce_results and self.tp_size > 1:
|
|
output = all_reduce(output_parallel, self.tp_group)
|
|
else:
|
|
output = output_parallel
|
|
|
|
output_bias = self.bias if self.skip_bias_add else None
|
|
|
|
return output, output_bias
|
|
|
|
def extra_repr(self) -> str:
|
|
s = f"input_features={self.input_size_per_partition}"
|
|
s += f", output_features={self.output_size}"
|
|
s += f", bias={self.bias is not None}"
|
|
s += f", tp_size={self.tp_size}"
|
|
s += f", reduce_results={self.reduce_results}"
|
|
return s
|