阿里安全使用 NVIDIA NeMo 框架和 TensorRT-LLM 的大模型工程化落地实践

ChatGPT 的爆火，把大模型推向了前台，也让它们在各行各业的落地速度骤然加快。阿里安全涉足大模型应用，至今也已有两年多的时间。借这个机会，把团队在大模型工程领域摸爬滚打积累下来的经验做个梳理，分享出来。

实践下来，阿里安全主要借助了

框架和 大语言模型推理加速库，在训练和推理两方面都拿到了实打实的性能提升。具体来说：NeMo 在多卡环境下能实现 2-3 倍的训练加速；TensorRT-LLM 结合 SmoothQuant Int8 能做到领先的推理加速比；而动态批处理策略（Dynamic Batch）将计算步骤减少了 30%，实际 QPS 提升了 2-3 倍。值得一提的是，prompt 优化策略在特定业务中更是将吞吐量提升了 10 倍。总体来说，这些优化手段显著提升了模型的性能和业务的效率。

先从工程的角度出发，聊聊大模型训练。我们团队过去两年在这方面积累了一些经验，特别是关于训练加速的部分。

NVIDIA NeMo 框架（基于 Megatron-LM）是一个端到端的云原生框架。无论你在本地还是在云上，都能用它来灵活地构建、定制和部署生成式 AI 模型。它的功能模块覆盖了预训练模型、数据管护、模型对齐、训练推理、检索增强和护栏工具包等，为用上生成式 AI 提供了一种既方便又经济的方式。NeMo 也支持多模态模型的训练，比如 Stable Diffusion、Vision Transformer 等。不过，

。

在使用 NeMo 进行大模型训练时，以下几个 feature 对速度的影响最大：

• 模型并行
  ，包括张量并行和流水线并行。当然，现在还有更新的 feature，比如长序列下的 context parallel。

• Layer & Kernel Fusion
  ，Megatron Core 会把多个算子的计算融合在一起，放到一个 kernel 里计算，以此提升训练速度。

• Distributed Optimizer
  ，Megatron Core 会把 Optimizer States 分配到各个 device 中，减少显存占用。此外还有比较新的 FSDP，能支持梯度和参数的进一步分块，进一步降低显存占用。

当然，NeMo 还有其他的 feature，比如让计算和通信 overlap，以及调用基于

实现的 FlashAttention 等等。

在 Megatron-LM 的公开论文中可以看到（如图 1 所示），Megatron Core 能保证在 GPU 水平扩展时，单卡的 FLOPs 基本保持不变；而 DeepSpeed 框架则有较大的衰减。175B 模型在 1,536 卡的规模上，Megatron-LM 的性能是 DeepSpeed 的 3 倍多；530B 模型在 2,240 卡规模上，Megatron-LM 也是 DeepSpeed 的 3 倍多。

图 1. Megatron-LM 论文中性能对比数据

(https://people.eecs.berkeley.edu/~matei/papers/2021/sc_megatron_lm.pdf)

我们团队基于 Llama2-13B 模型也做了类似的实验，结论一致：NeMo 比 DeepSpeed 性能高。具体数据如下表所示：

无论在单机 8 卡，还是双机 16 卡的规模上，NeMo 的性能都是 DeepSpeed 的 2 倍。

• 训练过程中，模型并行以及梯度交换需要在多个节点间通信，此时带宽会成为比较重要的瓶颈。训练大模型时，网卡带宽很容易成为瓶颈，所以建议使用 RDMA+ 多网卡方案。
• 为了降低多节点间模型 Optimizer States 的通信量，同时保证训练精度不受太大影响，可以将 Distributed Optimizer 的 dataType 设为 FP32，同时将 grad_sync_dtype 设为 BF16。

• 影响性能的三个重要参数
  ：megatron_amp_O2x0（混合精度 O2 优化选项，设为 False 会降低训练性能）；optim.grad_sync_dtype（梯度同步精度，与显存占用相关，使用 FP16 可减少显存占用）；optim.optimizer_dtype（optimizer states 参数类型，设为 FP16 或 BF16 会降低模型精度）。

• SFT Chat 类模型
  ：需要将 NeMo 的默认 prompt 配置修改成与原模型一致的 prompt，否则会造成训练 loss 增大，影响最终效果。

• TP 和 PP 的参数配置
  ：实践中发现，对于 13B 左右的模型，当显卡规模小于 32，且有高效的 RDMA 网络时，通常只需开启 TP（TP>1），保持 PP=1；如果模型更大，比如 70B 模型，则需要同时开启 PP（PP>1）。

在大模型推理方面，NVIDIA 推出的 TensorRT-LLM 实现了业界领先的性能。经过多方比较，我们最终选定它作为阿里安全大模型高性能推理的基石。由于业务中涉及的大模型服务较多，为了让算法同学能快速部署，工程团队开发了一系列功能，让他们可以快速、高效、平稳地完成部署。具体流程如下图所示：

图 2. 大模型部署流程，图片来源于阿里安全

• 模型校验和标准化导出
  ：对用户提供的模型文件和包含相关参数的配置文件进行校验，并生成标准的数据格式，方便后续的模型编译。

• 模型编译
  ：使用 TensorRT-LLM 将模型编译成 engine 格式。

• 服务 DAG 编排
  ：将服务的各个模块定义为 Op，通过 DAG 的方式在可视化界面中展示，让算法同学可以自主编排服务逻辑。

• DAG 调试和构建
  ：用户在 DAG 中可自行构造服务输入数据，完成对整个 DAG 的实时调试。调试结束后，可将 DAG 固化，并构建成真实的服务配置。

• K8s 服务部署
  ：我们的模型服务都通过 K8s 进行部署，便于快速部署、快速复制和扩容。

• 每条请求的推理耗时分布非常不均匀。耗时短的，几百毫秒就能返回；耗时长的，可能需要几十秒。这种情况在非 LLM 的模型上基本不存在。
• 大模型是生成式的，实际应用中会依赖上文信息——之前的推理结果需要在后续的推理中作为输入再次传入。
• 部署的服务需要能实时监控内部的健康状态。因为大模型服务在使用 GPU 时，GPU 内部的错误（如数组越界）会导致后续所有计算都失败，所以推理框架需要能及时感知这类错误，并做到快速重启。

针对这些特点，我们采用异步服务进行部署：上游调用先提交一次请求，然后通过轮询或回调的方式获取计算结果。具体架构图如下：

图 3. 模型服务架构图，图片来源于阿里安全

TensorRT-LLM 在 In-Flight Batching 的基础上新增了调度机制，即动态批处理（Dynamic Batch）。为了进一步提升推理速度，我们仔细分析了推理过程中的计算。从公式来看，要想让推理过程中的计算密度变大，只能调大 batch_size；而 generate 过程中，每个样本的输入长度和输出长度都不相同，必然会导致算力浪费。具体推理过程如下图所示：

图 4. 大模型 generate 生成过程示意图，图片来源于阿里安全

从案例中可以看到，一共有 4 条样本需要推理。若 GPU 最多只能一次处理 3 条，则共需 9 步完成这 3 条样本的推理。其中，第 1 条样本在 step1 时就结束了，input2 在 step5 时结束，因此这 9 步中间出现很多 EOS token（我们称之为气泡）。气泡越多，算力浪费越严重；剩下的 input4 只能按 batch_size=1 进行推理，浪费更明显。

在实际的推理服务中，如果一个 batch 里的某条输出长度特别大，气泡比例甚至可能超过 80%，造成严重的算力浪费。为了解决这个问题，我们提出了一种新的推理算法：

1. 准备输入数据（n = max_batch_size*10），形成一个候选集合 S。
2. 如果集合 S 为空，退出；否则，对集合 S 中的文本序列进行排序。
3. 从集合 S 中取出 m 个文本进行推理，其中 m = find_max_batch_size(current_seq_len)。推理到一定 step（一般取 100），或者当本次 batch 中 70% 的样本已推理完成且未完成样本数大于 1/max_batch_size 时，推理中断退出。
4. 将上次未推理完成的样本拼接上已生成的部分文本，作为一个新文本插入候选集合 S，然后重复第 2 步。

算法的核心思想是：根据当前 seq_len 动态设置每次 batch 的大小（让 GPU 内存占满，不浪费算力），并在推理过程中动态检测是否需要提前终止，终止条件是大部分样本已推理完成或推理的 step 达到一定限度。针对图 4 中的案例，采用新调度思想后的效果如下图所示：

图 5. 启用新的调度策略后的推理，图片来源于阿里安全

采用新调度逻辑后，这个案例的 step 数从 20 步减少到 14 步；在实际生产中，步数减少的幅度远不止于此。经过验证，采用上述算法后，服务的 QPS 通常有 2-3 倍以上的提升。

通常情况下，优化 prompt 能提升模型效果。来看一个案例：

输入：请判断一下，这段文字是否正确，如果错误请说明理由。「古代帝王把帝位让给贤能的人称为禅让，是原始社会末期部落首领的选举制度」
输出：对

输入：请判断一下，这段文字是否正确，如果错误请说明理由。「礼部为六部之一，掌管礼仪、祭祀、土地、户籍等职事，部长官称为礼部尚书」
输出：错，"礼部"，在明代有尚书一人，左、右侍郎各一人，下设四司，分别为：仪制清吏司，掌嘉礼、军礼以及管理学务、科举考试事；祠祭清吏司，掌吉礼、凶礼事物；主客清吏司，掌宾礼及接待外宾事物；精膳清吏司，掌筵飨等事务。

上述案例让大模型判断一段文字的正确性，错误则说明理由。在实际服务中，回答正确的 case 输出 token 数为 1，而回答错误的 case 输出 token 数可能很大。如 Dynamic Batch 一节所述，当输出长度差异很大时，推理过程中产生的气泡会非常大，造成严重的算力浪费。为解决这个问题，我们可以将问题分解：先让大模型判断正确性，再针对错误 case 询问原因。

输入：请判断这段文字是否正确，「古代帝王把帝位让给贤能的人称为禅让，是原始社会末期部落首领的选举制度」
输出：对

输入：请判断这段文字是否正确，「礼部为六部 之一，掌管礼仪、祭祀、土地、户籍等职事，部长官称为礼部尚书」
输出：错

输入：这段有事实错误，请指出原因，「礼部为六部之一，掌管礼仪、祭祀、土地、户籍等职事，部长官称为礼部尚书」
输出："礼部"，在明代有尚书一人，左、右侍郎各一人，下设四司，分别为：仪制清吏司，掌嘉礼、军礼以及管理学务、科举考试事；祠祭清吏司，掌吉礼、凶礼事物；主客清吏司，掌宾礼及接待外宾事物；精膳清吏司，掌筵飨等事务。

在实际业务中，采用这种策略后，部分业务的吞吐量提升了 10 倍以上。

我们的模型量化大部分是基于

（ModelOpt，原名 AMMO）完成的。ModelOpt 提供了简明易用的接口，可以对各种第三方模型进行训练后量化（PTQ），并与 TensorRT-LLM 实现良好衔接。目前 TensorRT-LLM 已完美支持常见的各种量化方法，如 INT8 weight only（W8A16）、AWQ/GPTQ（W4A16 groupwise）、FP8（W8A8）等。

我们分别在 Baichuan2-13B 和 Qwen-14B 模型上使用了上述各种量化方法进行实验。量化后的性能结果如下：

图 6. 图片来源于阿里安全

模型推理基于 TensorRT-LLM 实现。实验表明，在 batch_size 和 seq_len 都相同的条件下，sq_int8 的推理速度最快。

以两个实际落地的业务来说明量化的效果：

图 7. 图片来源于阿里安全

结论是：业务 1 上 smoothQuant int8 表现最好，业务 2 上 int8 weight only 表现更好。实践中发现，并没有哪一种方法始终是最好的，但相对而言，smoothQuant 在大部分业务场景中表现都比较稳健。有些资料提到 smoothQuant int8 对模型精度会有一定影响，但在我们的应用场景下，ModelOpt 量化出来的 smoothQuant int8 精度令人满意。

• 前期是量化方法的选择问题
  ：我们会在标准的业务数据集上测试各种量化方法的效果，然后选取稳定的方法应用到实际业务中。

• 量化后期要注意精度损失
  ：理论上，量化后的模型精度不可能与未量化模型完全相同，但我们在实际业务中能将量化带来的损失控制在 1% 以内。如果精度损失过大，通常可以通过调整量化过程中的校准数据集来改善，校准样本一般需要 2,000-8,000 条，而对校准数据集的要求是，其分布要尽可能与实际业务的样本分布一致。

• Attention 是 Transformer 的计算瓶颈
  。从 Transformer 的 FLOPs 分析可以看出，其主要计算集中在 attention 上。过去两年，业内的优化也主要集中在这一点上，比如 FlashAttention、FlashAttention-2 等方法。而 TensorRT-LLM 针对不同模型的 attention，基于 flash 思想实现了更丰富的功能和性能支持。

• 解决第 2-n 个推理的加速问题，是当前最迫切的
  。过去一年，业内提出的 flash-decoding++ 就是解决这个问题的。在 TensorRT-LLM 中对此部分也有特定的性能优化，在此基础上，通过调大 batch_size 可以进一步提升吞吐量。

• 优化 prompt、减少推理中的气泡也是性能优化的重要手段
  。需具体问题具体分析，总的原则是，在一个 batch 的推理过程中要避免 token 的无效计算。

• prompt 工程是大模型实际应用中非常重要的一环。prompt 设计的好坏，不仅影响模型的业务效果，也极大地影响服务的吞吐量。prompt 是个经验工程，需要在实践中不断尝试和总结。
• 另外，prompt 的 few-shot 中 example 的顺序也会影响模型推理的结果。

• 为避免大模型输出预期之外的结果
  ，建议开发者设计多套 prompt，逐步引导大模型正确输出，并在 prompt 中设计结束符，避免模型出现幻觉。

• 对于结构化的 prompt
  ，不推荐直接使用换行或空格来分开各个部分。建议使用 markdown 或 xml 语法，这可能是因为大模型在预训练阶段使用了大量 xml 和 markdown 语料。

• prompt 设计要考虑边界情况
  。例如，可以在 prompt 中添加类似这样的语句：

if you cannot fetch any information, output `no answer`