acext使用指南 (v1.5.1)

1.1.1 fpAintB Gemm

其中fpA_intB_gemm.h提供fpAintB的混合精度量化，其中A是浮点型，B是整型，通过三个模板参数可以指定fpAintB量化算法中A和B具体的类型，以及需要使用的量化算法。

template <typename T, typename WeightType, cutlass::WeightOnlyQuantOp QuantOp>
class CutlassFpAIntBGemmRunner : public virtual CutlassFpAIntBGemmRunnerInterface
{
public:
    CutlassFpAIntBGemmRunner();
    ~CutlassFpAIntBGemmRunner();

    void gemm(const void* A, const void* B, const void* weight_scales, void* C, int m, int n, int k,
        tkc::CutlassGemmConfig gemmConfig, char* workspace_ptr, const size_t workspace_bytes,
        cudaStream_t stream) override;

    void gemm(const void* A, const void* B, const void* weight_scales, const void* weight_zero_points,
        const void* biases, void* C, int m, int n, int k, const int group_size, tkc::CutlassGemmConfig gemmConfig,
        char* workspace_ptr, const size_t workspace_bytes, cudaStream_t stream) override;

    // Returns desired workspace size in bytes.
    size_t getWorkspaceSize(const int m, const int n, const int k) override;

    tkc::CutlassGemmConfig getChosenConfig(const void*          A,
                                           const void* B,
                                           const void*          weight_scales,
                                           const void*          weight_zero_points,
                                           const void*          biases,
                                           void*                C,
                                           int               m,
                                           int               n,
                                           int               k,
                                           const int         group_size,
                                           char*             workspace_ptr,
                                           const size_t      workspace_bytes,
                                           cudaStream_t      stream) override;
};

template <typename ActType, WeightOnlyQuantType QType, cutlass::WeightOnlyQuantOp QuantOp>
std::shared_ptr<CutlassFpAIntBGemmRunnerInterface> createFpAIntBGemmRunner();

1.1.2 w8a8 Gemm

int8_gemm.h提供A和B都是整型的量化算法，模板参数T可以指定输出的数据类型。

template <typename T>
class CutlassInt8GemmRunner : public virtual CutlassInt8GemmRunnerInterface
{
public:
    CutlassInt8GemmRunner();
    ~CutlassInt8GemmRunner();

    void gemm(const void* A, const void* B, tk::QuantMode quantOption, const float* alphaCol, const float* alphaRow,
        void* C, int m, int n, int k, tkc::CutlassGemmConfig gemmConfig, char* workspacePtr,
        const size_t workspaceBytes, cudaStream_t stream) override;

    // Returns desired workspace size in bytes.
    size_t getWorkspaceSize(const int m, const int n, const int k) override;

    std::vector<tkc::CutlassGemmConfig> getConfigs() const override;

    tkc::CutlassGemmConfig getChosenConfig(const void* A, const void* B, tk::QuantMode quantOption,
        const float* alphaCol, const float* alphaRow, void* C, int m, int n, int k, char* workspacePtr,
        const size_t workspaceBytes, cudaStream_t stream);
};

1.1.3 Moe Gemm

moe_kernel.h提供了调用moe gemm所需要的API

template<typename T,          /*The type used for activations/scales/compute*/
         typename WeightType, /* The type for the MoE weights */
         cutlass::WeightOnlyQuantOp QuantOp,
         typename Enable = void>
class CutlassMoeFCRunner: public CutlassMoeFCRunnerInterface {
public:
    CutlassMoeFCRunner()           = default;
    ~CutlassMoeFCRunner() override = default;
    //for rtp-llm
    void runMoe(const void* input_activations_void, const float* gating_output, const float *gating_output_with_bias,
        const void* fc1_expert_weights_void, const void* fc1_scales_void, const void* fc1_zeros_void, const void* fc1_expert_biases_void,
        ActivationType fc1_activation_type, const void* fc2_expert_weights_void, const void* fc2_scales_void, const void* fc2_zeros_void,
        const void* fc2_expert_biases_void, const int num_rows, const int hidden_size, const int inter_size,
        const int num_experts, const int k, const int group_size, char* workspace_ptr, void* final_output_void, void* fc2_result_void,
        const bool* finished, const int active_rows, void* expert_scales_void,
        int* expanded_source_row_to_expanded_dest_row, int* expert_for_source_row, MOEParallelismConfig parallelism_config,
        MOEExpertScaleNormalizationMode normalization_mode, cudaStream_t stream) override;
    // for vllm
    void fused_moe(const void* input_activations_void, const float* gating_output, const float* gating_output_with_bias,
                   const void* fc1_expert_weights_void, const void* fc1_scales_void, const void* fc1_zeros_void, const void* fc1_expert_biases_void,
                   void* a1_scale_void, ActivationType fc1_activation_type, const void* fc2_expert_weights_void, const void* fc2_scales_void,
                   const void* fc2_zeros_void, const void* fc2_expert_biases_void, void* a2_scale_void, const int num_tokens, const int hidden_size,
                   const int inter_size, const int num_experts, const int k, const int group_size, char* workspace_ptr, void* final_output_void,
                   void* fc2_result_void, const bool* finished, const int active_rows, const void* expert_scales_void,
                   int* expanded_source_row_to_expanded_dest_row, const int* expert_for_source_row, MOEParallelismConfig parallelism_config,
                   MOEExpertScaleNormalizationMode normalization_mode, cudaStream_t stream) override;

1.2.1 Int8 weight-only(W8A16)

Int8 weight-only是最基础的量化算法，它的基本原理是在列的方向上对权重矩阵进行量化，最终权重数据会从浮点型转换为int8类型，并且每一列会得到一个缩放因子。假如以原始的权重数据类型为fp16，这种量化算法相对于原始权重内存减少为原来的一半，而且由于需要加载的权重数据减少，量化算法的性能在memory-bound的情况下相对于原来的fp16运算会有提升。

具体代码示例：

//准备数据act, weight, scales，以及结果out
......

// 创建A为half，B为Int8，量化算法为PER_COLUMN_SCALE_ONLY的FpAIntBGemmRunner
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int8b, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>();

// 对weight数据进行离线预处理重排
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::INT8_WEIGHT_ONLY);

// 准备Workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// 选择运行Kenrel需要的config
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, nullptr, nullptr,
                                    p_out, m, n, k, k, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// 运行量化算法，结果存储在out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.2 Int4 weight-only(W4A16)

Int4 weight-only和int8 weight-only的算法是一样的，区别是权重矩阵使用的是4bit来存储，精度比int8 weight-only要差很多。权重的内存进一步减少为原来的1/4，性能进一步提升。

具体代码示例：

//准备数据act, weight, scales，以及结果out
......

// 创建A为half，B为Int4，量化算法为PER_COLUMN_SCALE_ONLY的FpAIntBGemmRunner
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int4b, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>();

// 对weight数据进行离线预处理重排
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::PACKED_INT4_WEIGHT_ONLY);

// 准备Workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// 选择运行Kenrel需要的config
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, nullptr, nullptr,
                                    p_out, m, n, k, k, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// 运行量化算法，结果存储在out
runner->gemm(p_act, p_weight, p_scales, p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.3 Int4 Groupwise(W4A16)

Int4 weight-only虽然对内存减少和性能提升都非常的友好，但是它的精度较差（在数据表示范围只有4bit的情况下，量化误差相对于原来的fp16会差很多）。为了解决这一个问题，引入了groupwise的方式来避免精度太差而导致不可用。基本的原理是在原来weight-only算法的基础上，不是每一列共享一个scale，而是将每一列按照一定的大小分成不同的group，每个group共享一个scale，其中groupsize一般为64或者128。

在具体的实现中（AWQ/GPTQ），一般需要通过数据集对模型权重进行微调得到量化后的int4权重，同时假如是非对称量化也会有零点的偏移（zeros & bias）。

具体代码示例:

//准备数据act, weight, scales，以及结果out
......

// 创建A为half，B为Int4，量化算法为FINEGRAINED_SCALE_ONLY的FpAIntBGemmRunner
auto runner = createFpAIntBGemmRunner<half, WeightOnlyQuantType::Int8b, cutlass::WeightOnlyQuantOp::FINEGRAINED_SCALE_ONLY>();

// 对weight数据进行离线预处理重排
preprocess_weights_for_mixed_gemm(preprocessed_quantized_weight, row_major_quantized_weight, {k, n}, QuantType::PACKED_INT4_WEIGHT_ONLY);

// 准备Workspace
char* ws_ptr = nullptr;
const int ws_bytes = runner->getWorkspaceSize(m, n, k);
cudaMalloc(&ws_ptr, ws_bytes);

// 选择运行Kenrel需要的config
gemmConfig = runner->getChosenConfig(p_act, p_weight, p_scales, p_zeros, p_bias,
                                    p_out, m, n, k, groupsize, reinterpret_cast<char*>(ws_ptr),
                                    ws_bytes, stream);

// 运行量化算法，结果存储在out
runner->gemm(p_act, p_weight, p_scales, p_zeros, p_bias, p_out, m, n, k, groupsize, gemmConfig, ws_ptr, ws_bytes, stream);

1.2.4 SmoothQuant(W8A8)

fpAintB实现的还是混合精度量化算法，虽然权重是int类型的，但是在模型推理的过程中，需要先将权重反量化为浮点类型，然后在进行矩阵运算，最终调用的还是浮点的TensorCore。W8A8实现输入和权重都是int类型，这样可以调用设备上的整型TensorCore（算力更好）进行运算。

具体的实现中，使用到W8A8量化kernel的是SmoothQuant这个量化算法。https://github.com/mit-han-lab/smoothquant

其中对于A是否使用PerToken（每一行共享一组量化参数，默认是整个矩阵共享一组量化参数），B是否使用PerChannel（每一列共享一组量化参数，默认是整个矩阵共享一组量化参数），W8A8又可以区分为PerTensorPerTensor，PerTensorPerChannel，PerTokenPerTensor，PerTokenPerCHannel四种变种：

acext实现了这四种W8A8的量化kernel，具体的调用方式如下：

// 定义使用哪种W8A8的变种，可以从下面四种中任选一种：
QuantMode quant_mode = QuantMode::fromDescription(false, false, false, false);  // A_PerTensor_W_PerTensor
QuantMode quant_mode = QuantMode::fromDescription(false, false, false, true); // A_PerTensor_W_PerChannel
QuantMode quant_mode = QuantMode::fromDescription(false, false, true, false); // A_PerToken_W_PerTensor
QuantMode quant_mode = QuantMode::fromDescription(false, false, true, true); // A_PerToken_W_PerChannel

// 准备数据act, weight, scales, out
......

// 创建W8A8的Runner
auto runner = std::make_shared<acext::kernels::cutlass_kernels::CutlassInt8GemmRunner<half>>();

// 准备Workspace
int ws_bytes = runner->getWorkspaceSize(m, n, k);
char* ws_ptr = nullptr;
if (ws_bytes)
    cudaMalloc(&ws_ptr, ws_bytes);

// 选择kernel执行时需要的config
const auto gemmConfig = runner->getChosenConfig(p_act, p_weight, quant_mode,
                    reinterpret_cast<const float*>(p_scales_channels),
                    reinterpret_cast<const float*>(p_scales_tokens),
                    p_out, m, n, k, ws_ptr, ws_bytes, stream);

// 执行kernel，结果存储在out中
runner->gemm(p_act, p_weight, quant_mode,
                    reinterpret_cast<const float*>(p_scales_channels),
                    reinterpret_cast<const float*>(p_scales_tokens),
                    p_out, m, n, k, gemmConfig, ws_ptr, ws_bytes, stream);

2.1.1 AutoWOQ

使用量化算法，我们需要将大模型原本的权重从原本的fp16/bf16格式转换为int8格式。

这里我们使用我们开发的AutoWOQ来进行转换，WOQ具体组成部分如下：

AutoWOQForCausalLM为所有的大模型提供一个统一的入口，和Huggingface的Auto类作用相似，接受的参数是模型权重的路径，AutoWOQForCausalLM会自动识别当前是的大模型的类型，并且初始化不同的WOQForCausalLM实例。这些实例都是BaseWOQForCausalLM的子类。具体的量化流程会在BaseWOQForCausalLM中实现，主要做的事情就是调用acext中实现的对称量化算法。

使用AutoWOQ量化模型代码流程如下：

# 设置量化参数，支持8bit or 4bit的量化
quantize_config = BaseQuantizeConfig(
    bits=8,  # quantize model to 8-bit
)

# 加载未量化的权重（Huggingface格式），model_path即为模型权重的路径，默认会将模型加载到CPU上
model = AutoWOQForCausalLM.from_pretrained(model_path, quantize_config, trust_remote_code=args.trust_remote_code)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 模型量化，会根据模型类型不同，针对模型中的linear层进行对称量化
model.quantize()

# 保存量化后的模型和Tokenizer到指定位置
quant_path = args.output
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

2.1.2 支持WOQ量化

有了WOQ量化后的权重后，我们给VLLM增加一种weightonly的量化算法。用户在使用WOQ量化权重时，需要显示的指定量化算法为weightonly。也就是创建LLM对象时，我们需要初始化参数quantization为weightonly。

# Create an LLM.
llm = LLM(model=model_path, quantization="weightonly")

这样vllm在加载这个量化后的模型时，会自动将需要量化的linear层创建量化类型的参数Tensor，并且运行时调用混合量化算法。

如下图所示，LLM大模型中一个Transformer Block在VLLM会表示如下，其中绿色的部分是Linear层，使用WeightOnlyQuant量化算法。

这些linear层需要使用WeightOnlyLinearMethod来进行初始化，加载量化权重，并且运行时调用acext的量化算法。由于acext使用cutlass进行kernel加速，对权重有一定的要求，而且由于这个重排需要考虑卡数，目前是在系统启动时进行的一次重排。

class WeightOnlyLinearMethod(LinearMethodBase):
    # 设置量化配置
    def __init__(self, quant_config: WeightOnlyConfig):
        self.quant_config = quant_config

    # 初始化Weight时，创建Int类型的权重参数，并且加载从WOQ处理后的量化权重
    def create_weights(self, input_size_per_partition: int,
                       output_size_per_partition: int, input_size: int,
                       output_size: int,
                       params_dtype: torch.dtype) -> Dict[str, Any]:
        qweight = Parameter(
            torch.empty(
                input_size_per_partition,
                output_size_per_partition,
                device="cuda",
                dtype=torch.int8,
            ),
            requires_grad=False,
        )
        scales = Parameter(
            torch.empty(
                output_size_per_partition,
                device="cuda",
                dtype=params_dtype,
            ),
            requires_grad=False,
        )

        // 实例化一个acext的WeightOnlyQuantMatmul类，后续可以通过它来调用混合精度kernel
        try:
            self.weightOnlyMatmul = ops.WeightOnlyQuantMatmul()
        except AttributeError as e:
            raise AttributeError(
                "WeightOnlyQuantMatmul kernels could not be imported. If you built vLLM "
                "from source, make sure BUILD_WEIGHTONLY_KERNELS=1 env var "
                "was set.") from e
        woq_state = WOQState.UNINITIALIZED
        ......

    # 运行时，调用acext的量化算法进行矩阵运算
    def apply_weights(self,
                      weights: Dict[str, torch.Tensor],
                      x: torch.Tensor,
                      bias: Optional[torch.Tensor] = None) -> torch.Tensor:
        //在线的权重预处理，系统启动时会处理一次
        if weights["woq_state"] == WOQState.UNINITIALIZED:
            qweight_cpu = weights["qweight"].cpu()
            qweight_cpu = preprocess_weights_for_mixed_gemm(qweight_cpu, torch.int8, False, False)
            weights["qweight"] = qweight_cpu.cuda()
            weights["woq_state"] = WOQState.READY
        qweight = weights["qweight"]
        scales = weights["scales"]
        return self.weightOnlyMatmul.weightonly_gemm(x, qweight, scales, bias)

2.1.3 WOQ模型推理

使用下面的命令编译vllm，设置BUILD_WEIGHTONLY_KERNELS=1会在编译时支持WeightOnly量化算法。

# 编译vllm
cd vllm
BUILD_WEIGHTONLY_KERNELS=1 python setup.py develop

# 运行vllm大模型推理任务, 假设我们已经将模型进行量化
python examples/offline_inference.py --model /path/to/Llama-2-7b-hf-woq/ -q weightonly --trust-remote-code

运行结果：

2.2.1 优化awq/gptq量化推理

awq和gptq是vllm原生就支持的两种量化算法，他们都是int4的量化。

https://docs.vllm.ai/en/latest/quantization/auto_awq.html

但是目前VLLM中的AWQ/GPTQ量化实现，使用的不是一种高效的实现，它目前只是实现了降低模型的内存占用，但是性能可能还没有FP16的性能好。

我们在vllm上通过acext的混合精度量化，可以实现更高效的量化kernel，相对于原来的AWQ/GPTQ有性能加速。但是由于acext的混合量化，对权重的排列有一定的要求，所以对于原来的AWQ权重，在调用acext前，需要对权重做一次预处理。

也就是我们通过对vllm中增加对acext weightonly_groupwise_gemm的支持，可以加速awq的量化推理（使用的还是原来的AWQ权重），只是在权重第一次被使用前，需要做一次预处理。具体的流程如下所示：

正对GPTQ量化，也有同样的流程。

和WOQ一样，目前的权重处理是在线处理的，vllm在启动时会有较长的延迟，后面会增加离线处理的形式。

2.2.2 Fast AWQ/GPTQ推理

由于目前这个优化的流程还在持续的迭代中，我们在vllm的实现里面，并没有默认打开使用acext加速awq/gptq的路径。需要设置环境变量USE_FAST_AWQ=1/USE_FAST_GPTQ才会调用到acext的加速推理。假如不加环境变量，默认还是使用的vllm自带的量化算法。

使用下面的命令行，可以调用到FastAWQ的流程，使用acext对awq量化进行加速。

USE_FAST_AWQ=1 python offline_inference.py --model /ppusw/share/LLM_Weights/Llama-2-7b-hf-awq -q awq --trust-remote-code

可以看到AWQ和FastAWQ在推理结果上是一致的，但是FastAWQ的推理速度更快。

2.3.1 构造acext backend

在acext_moe_func.cu中， 我们增加了acext backend, 使用class AcextFusedMoe来调用acext库的各个API，并通过csrc/moe/moe_ops.h和csrc/torch_bindings.cpp将其注册为python API。

2.3.2 修改fused moe layer

在fused_moe.py文件中增加acext的调用，包含在fused_experts_impl的实现中。首先使用ops.get_acext_fusedmoe_status()判断当前模型配置是否适合acext backend，不适合的会调用triton backend。 如果返回值acext_fusedmoe_status = 0，通过ops.acext_fusedmoe_warpper() 调用acext的fused_moe实现。

2.3.3 增加W8A8Int8MoEMethod(w8a8-Int8使用)

在compressed_tensors_moe.py文件中增加CompressedTensorsW8A8Int8MoEMethod用来load w8a8-int8-channel模型的weight并执行对应的moe layer。

2.4.1 已知问题

2.5.1 详细性能数据

2.6.1 运行步骤

#假如setuptools版本太低，需要先升级
#pip3 install --upgrade pip setuptools
#又因为以下的原因，我们需要限制setuptools为69.5.1
#https://github.com/vllm-project/vllm/issues/4961

git clone https://github.com/EleutherAI/lm-evaluation-harness
# 假如你运行mmlu测试遇到问题，可能和版本有关系，可以尝试切换到这个我们测试过的commit:
# git reset --hard e5e5ee0cb629c9c88165292d1b4bf34623392d33
cd lm-evaluation-harness
pip install -e .

#使用Huggingface的代理
export HF_ENDPOINT=https://hf-mirror.com

#使用lm_eval进行mmlu推理
lm_eval --model vllm --model_args pretrained=/ppusw/datasets/huggingface/Llama-2-7b-hf --tasks mmlu --batch_size auto --num_fewshot 5

#假如使用了代理后，还是遇到了下载数据集的问题，这是由于下面这个文件写死了用Huggingface网站下载
#/root/.cache/huggingface/modules/datasets_modules/datasets/hails--mmlu_no_train/b7d5f7f21003c21be079f11495ee011332b980bd1cd7e70cc740e8c079e5bda2/mmlu_no_train.py
#我们相应的修改这个文件里面的URL指向hf-mirror

2.6.2 推理结果

通过设置num-fewshot可以控制使用fewshot的参数，设置fewshot越高，运行时间越慢，但是得到的分数可能会更高。

不加num_fewshot时，运行的比较快（minutes），结果为41.3。

加上num_fewshot=5时，运行比较慢（hours)，分数更高，结果为45.9

这个和网站上的结果是一致的。

2.6.3 不同量化方法MMLU精度

我们假定没有量化前的模型fp16推理得到的MMLU分数为baseline，通过设置推理时使用不同的量化方法，可以得到量化后模型的MMLU分数，并且通过跟baseline比较，得到MMLU的量化损失。

以下是不同的模型运行MMLU（few-shot 0）在不同的量化方法下得到的分数，以及相对于fp16的量化损失：

其中灰色为当前版本（PPU Release 1.3）不支持的量化方法，红色为推理结果存在Bug的量化方法。