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

首先，我们从工程的角度，分析并总结了一些大模型训练的相关经验，针对过去两年阿里安全工程团队对大模型训练加速这一部分做出分享。

NVIDIA NeMo 框架（Megatron-LM）是 NVIDIA 提供的一个端到端的云原生框架，无论是在本地还是在云上，用户可以灵活地构建、定制和部署生成式 AI 模型。它包含但不限于预训练模型、数据管护工具、模型对齐工具、训练和推理框架、检索增强工具和护栏工具包，为用户使用生成式 AI 提供了一种既方便又经济的方法，同时，NeMo 也支持多模态模型的训练，包括但不限于 Stable Diffusion、Vision Transformer 等。本文关注焦点在于大模型训练框架的速度对比，因此只聚焦 NVIDIA Megatron Core 部分。在使用 NeMo 进行大模型训练过程中，影响训练速度比较大的 feature 主要如下：

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

  
  

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

  
  

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

当然，NeMo 还有其他的一些 feature，比如计算和通信做 overlap，调用基于 cuDNN 实现的 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 倍。

NeMo 性能评测小结

• 由于训练过程中，模型并行以及梯度交换等需要在多个节点中通信，带宽会成为一个比较重要的瓶颈，在训练大模型的时候，网卡的带宽很容易成为瓶颈，建议使用 RDMA+ 多网卡进行大模型训练。

  
  

• 为了降低多节点之间模型的 OptimizerStates 的通信量，同时也保证训练的精度不会受到太大的影响，可以将 Distributed Optimizer dataType 设置为 FP32，同时将 grad_sync_dtype 设置为 BF16。

NeMo 使用问题总结

• 影响训练性能的三个重要参数：megatron_amp_O2x0；混合精度 O2 优化选项，设为 False 会降低训练性能；optimx0；.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，实现了业界领先的性能。经多方比较，TensorRT-LLM 被我们选为构建阿里安全大模型高性能推理的基石。在阿里安全的业务中，由于大模型服务比较多，为了让算法同学可以快速部署大模型，工程团队开发了一系列功能，让算法同学可以快速、高效、平稳地部署大模型。部署的具体流程图如下：

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

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

  
  

• 模型编译阶段：使用 TensorTR-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 中的其中一条输出的 output_len 很大，就可能会导致该 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 步；在实际的业务生产过程中，这个 step 的减少数远远大于案例中的 case。我们通过实践发现，在我们的业务中采用上述算法，服务的每秒查询率（QPS）一般有 2-3 倍以上的提升。

优化 prompt，提升服务性能

一般情况下，优化 prompt 可以提升模型的效果，案例如下：

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

上述案例是让大模型判断一段文字的正确性，如果错误则说明理由。在实际服务的生产过程中，对于回答是正确的 case，输出的 token 数是 1，对于回答错误的 case，输出的 token 数可能很大，如 Dynamic Batch 一节分析的那样，当输出的长度差别很大的时候，推理过程中产生的气泡会很大，这种情况将造成严重的算力浪费。为了解决该问题，我们可以将问题进行分解，例如对大模型提问两次，第一次让大模型判断问题的正确性，第二次针对事实错误的 case 问大模型原因。

输入：请判断这段文字是否正确，「古代帝王把帝位让给贤能的人称为禅让，是原始社会末期部落首领的选举制度」输出：对输入：请判断这段文字是否正确，「礼部为六部 之一，掌管礼仪、祭祀、土地、户籍等职事，部长官称为礼部尚书」输出：错
输入：这段有事实错误，请指出原因，「礼部为六部之一，掌管礼仪、祭祀、土地、户籍等职事，部长官称为礼部尚书」输出："礼部"，在明代有尚书一人，左、右侍郎各一人，下设四司，分别为：仪制清吏司，掌嘉礼、军礼以及管理学务、科举考试事；祠祭清吏司，掌吉礼、凶礼事物；主客清吏司，掌宾礼及接待外宾事物；精膳清吏司，掌筵飨等事务。

在实际业务实践中，对于这类模式的问题采用上述策略，部分业务的吞吐提升 10 倍以上。

模型推理量化

我们的模型量化大部分都是基于 NVIDIA TensorRT Model Optimizer（简称 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 是 Transformers 的计算瓶颈，从 Transformers 的 FLOPs 分析的最后计算公式中可以看出，Transformer 模型的主要计算是集中在 attention 计算上，过去两年，业内主要针对 Transformer 模型的优化也集中在这块，比如 FlashAttention、FlashAttention-2 等方法都是针对 attention 计算的优化。而 TensorRT-LLM 中针对不同模型的 attention，基于 flash 的思想，实现了更加丰富的功能及性能的支持。

  
  

• 大模型推理解决第 2-n 个推理的加速问题成为最迫切的问题，过去 1 年业内提出的 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`