生成式人工智能/大语言模型

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

前言

随着 ChatGPT 的一夜爆火,大模型如今越来越广泛的应用到各种业务领域中,阿里安全的业务领域对大模型技术的应用也已经 2 年有余。本文对阿里安全在大模型工程领域积累的实践经验做出总结和分享。

在大模型实际应用实践的过程中,阿里安全采用 NVIDIA NeMoTM 框架TensorRT-LLM 大语言模型推理加速库,显著优化了模型训练与推理性能。其中 NeMo 在多卡环境可实现 2-3 倍的训练加速,TensorRT-LLM 结合 SmoothQuant Int8 可实现领先的推理加速比,动态批处理策略 (Dynamic Batch) 将计算步骤减少 30%,实际 QPS 增益 2-3 倍。Prompt 优化策略在特定业务中提升吞吐高达 10 倍。整体优化成果显著增强了模型性能与业务效率。

扫描下方二维码、点击文章底部阅读原文,或复制链接(https://www.nvidia.cn/webinars/sessions/?session_id=240516-32936&ncid=so-wech-741757-vt04),登陆后即可观看相关演讲,并下载讲义。

大模型训练

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

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 论文中性能对比数据

我们团队在 Llama2-13B 的模型做了类似的实验,得出的结论也是 NeMo 比 DeepSpeed 性能高,具体的数据如下标所示:

 gpusgbsmbsiter timeTFLOPS/GPUpercent of peak flops
NeMo (TP=2, PP=1)8128128.719964%
DeepSpeed (ZeRO3)8128158.4810032%
NeMo (TP=2, PP=1)16128114.719563%
DeepSpeed (ZeRO3)16128130.459631%

无论在单机 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 步完成这三条样本的推理,而第 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 到达一定限度。针对图九中的案例,我们使用新的调度思想之后的效果如下图所示:

图 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`

未来计划

  • 为了进一步提升大模型的 generate 的速度,后续将会尝试 Medusa Decode 等解码手段。
  • 后续会考虑尝试使用 Attention Sink 相关 skills 在长文本中尝试。
  • 过去 1 年一直没有在 FP8 上进行尝试,后续将会用 FP8 在实际业务中实践。
  • RAG 的应用,目前大模型能力有一定限制,后续将会尝试使用 RAG(检索增强)的方式提升大模型在业务中效果。
  • 持续关注 TensorRT-LLM 最新的进展和 feature,TensorRT-LLM 这个开源项目更新迭代的频率还是挺快的,基本 2-3 个月就会有一个大版本出来。

感谢

本文最后感谢我的主管和团队,以及 NVIDIA 解决方案技术团队和 GPU 计算专家团队对本文的指导。

Tags