如何在 NVIDIA TensorRT-LLM 中支持 Qwen 模型
背景介绍
大型语言模型正以其惊人的新能力推动人工智能的发展，扩大其应用范围。然而，由于这类模型具有庞大的参数规模，部署和推理的难度和成本极高，这一挑战一直困扰着 AI 领域。此外，当前存在大量支持模型部署和推理的框架和工具，如 ModelScope 的 Model Pipelines API，和 HuggingFace 的 Text Generation Inference 等，各自都有其独特的特点和优势。然而，这些工具往往未能充分发挥 GPU 的性能。
为了解决这些问题，NVIDIA 推出了一种全新的解决方案——TensorRT-LLM。这是一款高度优化的开源计算框架，它将 NVIDIA TensorRT 的深度学习编译器、FasterTransformer 的优化内核、预处理和后处理，以及多 GPU/多节点通信等功能封装在一个简单的开源 Python/C++ API 中，同时与硬件进行了一体化优化，形成了一种产品级的大模型推理解决方案。NVIDIA TensorRT-LLM 具有多项突出的特性，包括支持新的 FP8 数据格式，这使得模型可以在更低的精度下运行，从而减少内存消耗，同时保持模型的准确性。它还支持一种名为'In-flight batching'的新调度技术，可以更有效地处理动态负载，提高 GPU 利用率。
此外，TensorRT-LLM 还支持模型的并行化和分布式推理，利用张量并行性进行模型并行化，使模型可以在多个 GPU 之间并行运行，从而实现大型模型的高效推理。最重要的是，TensorRT-LLM 极大地简化了开发流程，使得开发者无需深入了解底层的技术细节，也无需编写复杂的 CUDA/C++ 代码。它提供了一个易用、开源和模块化的应用编程接口，使开发者能够轻松定义、优化和执行新的大型语言模型架构和增强功能。总的来说，TensorRT-LLM 让用户可以专注于模型的设计和优化，而将底层的性能优化工作交给 TensorRT 来完成，大大提高了开发效率和生产效率，真正实现了大模型推理的易用性和高效性。
阿里云的通义千问开源模型 Qwen-7B，拥有 70 亿参数，在一系列全方位的评估中展示了其在自然语言理解与生成、数学问题求解、代码生成等领域的优秀能力。这些评估涵盖了多个数据集，包括 MMLU、C-Eval、GSM8K、HumanEval 以及 WMT22 等。在这些评测中，Qwen-7B 不仅超越了同等规模的其他大型语言模型，甚至在某些方面超过了参数规模更大的模型。因此，对于 TensorRT-LLM 来说，支持 Qwen 系列模型具有重要的意义。
开发与优化过程
我们是社区开发者，通过阿里云天池举办的 NVIDIA TensorRT Hackathon 2023 接触到了 NVIDIA TensorRT-LLM，并为它贡献了代码。以下是我们的开发记录，供其他开发者参考。
基础功能支持
- 分析示例代码：首先初步分析了 examples/llama 代码，以深化对 trt-llm 基本流程的理解。在 llama 项目的 weight.py 中，存在一个 load_from_meta_llama 函数，该函数包含 tensorrt_llm.models.LLaMAForCausalLM，此部分定义了 TensorRT 的模型结构。复制 examples/llama 并将其重命名为 examples/qwen，同时将 LLaMAForCausalLM 复制并创建新的 mode.py 文件，将相关内容粘贴至此。在这个过程中，所有包含'llama'的模型都被替换为'qwen'。
- 权重对齐：接下来，我们对项目中的 weight.py 的 load_from_hf_qwen 函数进行修改，目的是逐步将 HuggingFace 的权重名称与 TensorRT 的权重名称对齐。执行 build.py 后，虽然编译成功，但执行 run.py 的结果并未如预期。
- 调试与修正：参照 TensorRT-LLM 的 docs/source/2023-05-19-how-to-debug.md 文档，我们对模型进行了详细的调试，从外到内打印模型层的数值，观察 mean/sum/shape，并与原版进行对比。经过排查，我们发现 attention 部分已经包含了 rope 计算，通过调整 gpt attention plugin 的参数，最终使得输出的 logits 正常。
- 生成逻辑适配：再次优化 run.py，将 HuggingFace 原版的 qwen_generation_utils.py 中的 make_context 函数迁移到 utils/utils.py 中，并导入该函数。这个函数被用来构造一个 chat 版的 prompt 输入，同时我们调整 eos 和 pad token 为 qwen 专属的 <|im_end|> 或者 <|endoftext|>，最终 run.py 输出也正常。
增加功能：Weight Only 量化
在 FP16 对齐成功，并且 run.py 以及 summarize.py 文件均能正常运行之后，我们开始探索实现 weight only int8/int4 量化。这只需要在 build.py 文件中对 weight only int8/int4 分支进行轻微调整，包括 shape 的修改，以及保持权重名称与 FP16 一致。接下来，我们进行编译测试，发现这一过程顺利完成，且工作量并未超出预期，这部分工作基本无需投入大量人力资源。
