87a19d2d00
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
234 lines
12 KiB
Python
234 lines
12 KiB
Python
import os
|
|
|
|
import torch
|
|
|
|
def w8a16_prepare_quantize(self, quant_params={}):
|
|
# We need do nothing in here
|
|
pass
|
|
|
|
|
|
def w8a16_export_quantized_weights(self, save_path, quant_params={}):
|
|
gb_per_file = quant_params.get("filesize_limit", None)
|
|
int8_min = -127
|
|
|
|
def w8a16_quantization(weight):
|
|
# all weights should be [output,input], otherwise, we may get an wrong weight and scale...
|
|
scale = torch.abs(weight).max(dim=-1)[0] / 127.0
|
|
int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous()
|
|
scale = scale.view(1,-1).contiguous()
|
|
return int8_weight, scale
|
|
|
|
|
|
model = self.model_runner.model
|
|
|
|
from vllm.distributed.communication_op import (
|
|
tensor_model_parallel_all_gather,
|
|
tensor_model_parallel_all_reduce,
|
|
)
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
MergedColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
|
|
for name, m in model.named_modules():
|
|
if isinstance(m, VocabParallelEmbedding):
|
|
# weight shape: [vocab_size // tp, embedding_dim]
|
|
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
|
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
|
print(f"merged: {name}, shape={m.weight.shape}")
|
|
|
|
elif isinstance(m, ParallelLMHead):
|
|
# weight shape: [vocab_size // tp, embedding_dim]
|
|
# bias shape: [vocab_size // tp]
|
|
if m.bias is not None:
|
|
bias = tensor_model_parallel_all_gather(m.bias, dim=0)
|
|
bias = bias[:m.org_vocab_size].contiguous()
|
|
if self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
|
|
|
|
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
|
weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
|
|
print(f"merged: {name}, shape={m.weight.shape}")
|
|
|
|
elif isinstance(m, QKVParallelLinear):
|
|
# weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
|
|
# bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
|
|
if self.parallel_config.world_size > 1:
|
|
total_q_hidden_size = m.total_num_heads * m.head_size
|
|
partial_q_hidden_size = m.num_heads * m.head_size
|
|
total_kv_hidden_size = m.total_num_kv_heads * m.head_size
|
|
partial_kv_hidden_size = m.num_kv_heads * m.head_size
|
|
|
|
if m.bias is not None:
|
|
bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
|
|
q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
|
|
k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
|
|
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
|
v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
|
|
[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
|
|
|
q_bias[:] = m.bias[:partial_q_hidden_size]
|
|
k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
|
|
v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
|
|
|
|
bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
|
|
if self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
|
|
|
|
weight_tensor = m.weight.new_zeros(
|
|
total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size
|
|
)
|
|
|
|
q_tensor = weight_tensor[:total_q_hidden_size, :]
|
|
q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
|
|
|
|
k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size]
|
|
k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
|
|
|
v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :]
|
|
v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
|
|
|
|
q_tensor[:, :] = m.weight[: partial_q_hidden_size, :]
|
|
k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :]
|
|
v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :]
|
|
|
|
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
|
else:
|
|
weight_tensor = m.weight
|
|
if m.bias is not None and self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
|
|
|
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
|
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
|
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
|
print(f"Quantized: {name}")
|
|
|
|
elif isinstance(m, MergedColumnParallelLinear):
|
|
# weight shape: [intermediate_size // tp * 2, hidden_size]
|
|
# bias shape: [intermediate_size // tp * 2]
|
|
if self.parallel_config.world_size > 1:
|
|
output_sizes = m.output_sizes
|
|
output_size = sum(output_sizes)
|
|
partial_output_sizes = [
|
|
i // self.parallel_config.world_size for i in output_sizes
|
|
]
|
|
|
|
if m.bias is not None:
|
|
index_start = 0
|
|
partial_index_start = 0
|
|
bias_tenosr = m.bias.new_zeros(output_size)
|
|
for i in range(len(output_sizes)):
|
|
index_out = index_start + output_sizes[i]
|
|
sub_bias_tensor = bias_tenosr[index_start:index_out]
|
|
partial_size = partial_output_sizes[i]
|
|
sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
|
|
|
|
index_start += output_sizes[i]
|
|
partial_index_start += partial_size
|
|
bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
|
|
if self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
|
|
|
|
weight_tensor = m.weight.new_zeros(output_size, m.input_size)
|
|
|
|
index_start = 0
|
|
partial_index_start = 0
|
|
for i in range(len(output_sizes)):
|
|
index_out = index_start + output_sizes[i]
|
|
sub_weight_tensor = weight_tensor[index_start:index_out]
|
|
partial_size = partial_output_sizes[i]
|
|
sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size]
|
|
|
|
index_start += output_sizes[i]
|
|
partial_index_start += partial_size
|
|
weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
|
|
else:
|
|
weight_tensor = m.weight
|
|
if m.bias is not None and self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
|
|
|
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
|
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
|
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
|
print(f"Quantized: {name}")
|
|
|
|
elif isinstance(m, ColumnParallelLinear):
|
|
# weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
|
|
# bias shape: [some_dim // tp]
|
|
if m.bias is not None:
|
|
bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
|
|
if self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
|
|
|
|
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
|
|
|
|
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
|
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
|
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
|
print(f"Quantized: {name}")
|
|
|
|
elif isinstance(m, RowParallelLinear):
|
|
# weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
|
|
# bias shape: [hidden_size]
|
|
if m.bias is not None:
|
|
bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
|
|
if self.is_driver_worker:
|
|
m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
|
|
|
|
weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1)
|
|
int8_weight, weight_scales = w8a16_quantization(weight_tensor)
|
|
|
|
if self.is_driver_worker:
|
|
m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
|
|
m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
|
|
print(f"Quantized: {name}")
|
|
else:
|
|
pass
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
# save weights
|
|
if self.is_driver_worker:
|
|
from safetensors.torch import save_file
|
|
|
|
tensors = {}
|
|
saved = False
|
|
count = 0
|
|
size_in_bytes = 0
|
|
|
|
for name, weight in model.named_parameters():
|
|
if "lm_head" in name and model.config.tie_word_embeddings:
|
|
continue
|
|
size_in_bytes += weight.numel() * weight.element_size()
|
|
tensors[name] = weight
|
|
saved = False
|
|
if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
|
|
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
|
save_file(tensors, weight_path)
|
|
print(f"The quantified weights were successfully saved in {weight_path}.")
|
|
tensors.clear()
|
|
saved = True
|
|
count += 1
|
|
size_in_bytes = 0
|
|
|
|
if not saved:
|
|
weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
|
|
save_file(tensors, weight_path)
|
|
print(f"The quantified weights were successfully saved in {weight_path}.")
|