导读 本文将分享 NVIDIA 在大语言模型领域的解决方案。
包括三部分内容：1. 第一部分是关于 NeMo Framework。NeMo Framework 是 NVIDIA 大语言模型的全栈解决方案，它通过不同组件完成生成式 AI 各个环节的任务，包括数据预处理、分布式训练、模型微调、模型推理加速及部署（TensorRT-LLM及Triton）、检索增强生成（RAG）、护栏技术(Guardrails)等等。
2. 第二部分是关于 TensorRT-LLM。TensorRT 从 2016 年推出以来，一直在不断发展。随着进入 AI 2.0 大语言模型时代，NVIDIA 隆重推出了 TensorRT-LLM。在 TensorRT 的基础上，TensorRT-LLM 对 LLM 部分进行了相应的封装和增强。一个令人振奋的消息是，TensorRT-LLM 是开源的，采用 Apache-2.0 开源协议，对用户非常友好。
3. 第三部分是关于 RAG（Retrieval Augmented Generation），即检索增强生成。在使用大语言模型时，存在一个问题，即模型在某些未掌握领域可能出现胡言乱语。为了解决这个问题，NVIDIA 引入了 RAG 技术，用于增强和辅助实际大模型的应用。
分享嘉宾｜Jack He （何成杰） NVIDIA 资深解决方案架构师
编辑整理｜吴叶国
内容校对｜李瑶
出品社区｜DataFun
01
1. NeMo Framework Overview
这一轮生成式 AI 浪潮涵盖的范围非常广泛，除了大家熟知的大语言模型之外，还包括文生图、科学研究，和最近备受关注的多模态及不同模态切换等方面。生成式AI 不仅仅局限于大语言模型，还涵盖了与深度神经网络相关的部分以及更广泛的内容。
聚焦于大语言模型领域，NVIDIA 推出的 NeMo
Framework 涵盖了大语言模型的开发、部署以及使用的全流程。整体包括六大部分，如上图所示，其中前三部分涉及模型开发，后三部分涉及企业应用部署。- 第二，分布式训练，其中涉及诸多技术，后文中会详细讨论。
- 第三，模型定制化，针对不同的场景或行业进行相应的适配工作。
- 第四，推理加速。在这方面，NVIDIA 有两个非常重要的产品，一个是用于单个模型优化加速的 TensorRT-LLM，另一个是用于直接进行模型部署的 Triton。
- 第五，RAG（Retrieval Augmented Generation），召回和增强生成部分。涉及知识库、向量库以及 Agent，与大模型的最终落地息息相关。
- 第六，Guardrails，主要充当守门员的角色。对于一些常见的问题，比如对时政等问题首先进行过滤。
NeMo Framework 不仅仅用于大模型的训练和推理过程，也可以适配其他模态，比如语音转文字等不同场景。2. NeMo Framework Key
Components
接下来介绍 NeMo Framework 的核心组件。
首先是数据处理部分，这是一个关键环节。在实际场景中，原始数据集通常需要经过多次处理。典型的处理步骤包括去重、清除一些低质量的数据等。对于去重，通常使用基于规则的方法。例如，如果发现某条数据中同一个词语重复出现超过一定的频率，我们就会将其视为不合适或低质量的数据，这就是基于规则的过滤方法。此外，我们还可以通过各种手段，包括训练模型的方式，来区分低质量和高质量的数据。通过这一系列手段对数据进行处理的过程称为数据预处理。这一部分非常重要，因为数据的质量直接影响到最终模型的训练效果，我们必须注意确保数据预处理的质量。
如上图所示，经过一些处理，比如去重、Quality
filtering，通过对数据进行高质量的处理，模型效果得到了显著的提升。右侧梯度图表示对数据进行了不同处理后的效果，例如，经过去重和基于规则的数据清洗，训练效果有了明显的提升。目前，NeMo 提供了一个专门的工具 NeMo Data Curator，让用户评估并提高数据集的质量。
NeMo 底层使用了 NVIDIA Megatron Core 技术，对训练进行了很多加速操作。因为在大语言模型的预训练阶段，性能要求相当高，训练所需的资源和时间都是相当庞大的，加速训练非常必要。在加速方面，有三种方法：- 第一种是 Tensor 和 Pipeline 的并行。这意味着将模型的每一层 Tensor 做划分，并让每个 GPU 分别计算其中一部分 Tensor，以加速计算过程。
- 第二种是 Sequence 的并行，因为 Dropout 和 LayerNorm 并不依赖完整 reduce 后的结果，所以可以拆分到不同 GPU 上运行，以加速计算过程。
- 第三种是选择性重计算或称为选择性激活重计算。对于激活函数并不需要每次都进行重计算。激活层占用大量显存，这一方法通过精心的选择性计算一部分激活函数，可以显著减少显存使用和重复计算，提高效率。
这些方法都是 Nemo 框架在训练中用来提升效率的手段。
NeMo Framework 还提供了一个自动选择参数的工具，叫Auto-Configurator。这个工具的使用非常简便，用户只需要提供一些训练时的限制，比如训练的时间限制、训练规模的设定等。将这些限制写入一个文件中作为输入，通过 Auto-Configurator 工具能够自动生成合理的训练参数表。用户可以使用输出参数表进行后续的训练，因此这个工具非常实用。
大模型训练可以分为不同的类型。在上图中，从左到右分别是 Prompt Engineering、Prompt Learning、Parameter Efficient Fine-Tuning、Fine Tuning。这几种训练类型在资源消耗和对参数的改变等方面都存在很大差异，效果也有很大的不同。NeMo Framework 对这四种类型的 Tuning 都是支持的。这里特别提一下，SFT 和 RLHF 会对模型所有参数都进行改变。在这四种微调中，SFT 和 RLHF 所需要的资源是最大的。上图中提到了 LoRA 技术。该技术通过矩阵的分解，将一个大的参数矩阵分解成两个矩阵相乘。其中大矩阵保持不变，而小矩阵在训练时进行改变。这样就能以最小的参数代价达到最佳效果。3. NeMo Framework
Deployment Practice
下面介绍在使用 NeMo Framework 时的一些最佳实践。
第一步，使用 NeMo Data
Curator 对数据集进行清洗。第二步，配置参数。配置好参数后，可以直接使用 NeMo
Framework Launcher。这个 Launcher 实际上是由许多脚本或者 E2E recipes 组成的。这些 recipes 中，训练参数可以进行自定义设置，比如需要运行多少步，每多少步保存一次 checkpoint 等。配置好这些参数后，只需启动相应的脚本，训练就会开始。模型训练过程，可参见上图。中间的框图描述了 NeMo 在整个训练过程中的步骤。尤其是最耗时的 Pre-Train 过程。第三步，Pre-Train 结束后会进入模型效果对齐（aligner）阶段，比如与场景相关的 align，以及一些类似强化学习的模型调整，NeMo 平台也会提供支持。不管是 SFT、PEFT、Reward Model 还是 RLHF，NeMo 都可以通过 NeMo Framework
Launcher 启动不同的脚本，使用不同的参数选择不同的阶段来完成训练或者微调。这里提到了 NeMo Training
Container，NeMo Framework Launcher 实际上是一套脚本，通过参数配置控制 NeMo Training Container 容器进行训练。在容器化后，可以很方便地进行单节点多卡、单节点多 Pod，甚至多节点多卡等扩缩容。尤其在扩容方面非常方便，因为训练通常需要大量资源，可以由 NeMo Framework Launcher 进行调配组织和协调容器资源。当训练过程结束后，就涉及到 Inference 阶段。对于 Inference，NeMo 也提供了 Inference Container。这个容器将运行 Triton 和 TensorRT-LLM，完成推理工作。通常情况下，NeMo Training
Container 运行在 Slurm 或者是 K8S 集群上。K8S 目前也是支持的，但可能会有一些限制，Slurm 集群是比较主流的选择。但推理部署通常是基于 K8S 的。
通常来说，训练可以分为三种类型。首先是预训练（Pre-training），其次是监督微调（Supervised Fine Tuning，SFT），最后是提示学习（Prompt Learning）。这三种训练需要的资源是大相径庭的。- 预训练（Pre-training）是一种从头开始训练的方式 Training from scratch，因此它需要的资源是最多的。通常情况下，它涉及到多个节点，每个节点可能有 8 张显卡，使用大规模集群来进行训练。
- 监督微调（Supervised Fine Tuning，SFT）相对来说资源需求较少，通常情况下不需要多节点，但会使用多张显卡进行训练。
- 提示学习（Prompt Learning），只需要较少的资源，一般单卡的或者较小规模的资源就能够完成训练。
TensorRT-LLM
1. TensorRT-LLM Optimizing
LLM Inference训练完成后，将进入一个非常重要的步骤，即推理部分。下面就来介绍TensorRT-LLM。
TensorRT-LLM 的作用主要在于推理加速。简而言之，推理加速主要关注两个方面：延迟和吞吐。延迟指的是输入一定数量的字符后，返回字符所需的时间；吞吐则涉及一次性输入多个句子并返回推理结果的能力。因此，推理加速的目标是尽量减少延迟，同时提高吞吐能力。
目前，TensorRT-LLM 已经在 Github 上开源，我们也提供了详尽的文档，大家可以通过相关链接学习和使用。
在 NeMo 生态中，Triton Inference Server 主要负责整个模型的部署，而TensorRT-LLM 主要负责模型推理的加速，使模型推理能够更加高效。
这里介绍下 TensorRT-LLM 和 TensorRT 之间的关系。TensorRT-LLM 实际上是基于 TensorRT 的，它对 LLM（语言模型）相关的一些操作进行了一些优化，但是很多 CUDA kernel 仍然来自于 TensorRT。KV Caching，每次计算中，KV Caching 始终是一个较大的部分，因为这部分有很多无需进行重复计算的内容，需要将之前的计算结果保存在 Caching 中，这是一个非常重要的优化。其次，对于 MHA（Multi-Head Attention）kernels，也就是多头注意力的 CUDA kernel，TensorRT-LLM 也做了增强。除此之外，TensorRT-LLM 有 Inflight Batching。语言模型接收的每句话长度都不同，对于比较短的话，很快就结束了，是否可以在运行时插入后面一句话进行推理呢？这是语言模型的一个特点，与视觉不同，视觉推理的每个请求基本上都具有相对固定的维度，正因为语言模型存在每句话的长度都不同的特点，因此需要 Inflight Batching 特性。对于大型语言模型，即使在推理阶段，也可能涉及到多个 GPU、多个节点等，因此 TensorRT-LLM 还拥有 Multi-GPU、Multi-Node 等功能。除了这些特性之外，TensorRT-LLM 剩下的部分与 TensorRT 是一致的，它也会借用到非常多 TensorRT 原有的 CUDA kernel。以矩阵运算为例，MHA kernel 实际上最终会退化到矩阵运算层，因此，最终它仍然会使用 TensorRT 中的GEMM 运算，因为这些都是通用的。
首先 TensorRT-LLM 将整个网络模型以及参数加载进来，然后进行 engine 的构建。构建 engine 的过程实际上是为了选择最快的计算过程，这包括对模型图的重构，比如层的融合以及各种图操作等。然后选择运算最快的 CUDA kernel。最终，将最优 CUDA kernels 固定为一个新的 engine，以后每次运行这个模型都会通过这种方式来加载运行。整个流程与 TensorRT 是一致的。
一般都会按照以下步骤进行：首先是构建，包括加载模型和加载权重；接着，建立一个 engine；然后尝试运行这个构建好的 engine，并测试效果。
TensorRT-LLM 的工作流程借鉴了之前我们所采用的 Fast Transformer 的流程，随着时间推移，这些 FT kernel 的数量会逐渐增多，包括我们自己编写的一些 plugin，所有这些都会被集成到 TensorRT-LLM 中。此外，在 TensorRT 中有一个编译器，它会从 pattern 的角度来进行 kernel 优化。
因为对于大型语言模型，可能会牵涉到卡间通信，甚至节点间通信的问题，因此也需要将 NCCL 部分纳入考虑。
首先来看一下 Inflight Batching。在我们的语言模型中，每一句话的长度可能各不相同，为了解决这个问题，我们采取了一些措施，以确保整个计算过程不会因为长度的不同而导致计算不充分或浪费资源。我们会将后续的内容插入其中，使得整个计算块都能够被充分利用，如上图所示。Inflight Batching 在运行过程中实施了动态的 Batching，因此被称为 Inflight Batching。
TensorRT-LLM 对多种 attention类型提供支持。我们知道，Multi-Head Attention 实际上还有其他类型的变体，例如有 Multi-Query Attention、Group-Query
Attention 等，这些都是为了减少 attention 计算量而进行的一些优化，但这些方法也导致了 attention 计算过程的差异，因此针对这些不同的计算过程，需要有一些 CUDA
kernel 来进行优化支持。这里的 optimize attention 主要指的是对不同类型的 attention 进行支持的优化。
TensorRT-LLM 目前已支持多种模型的量化过程，上图中显示了对不同模型的支持情况。
目前，TensorRT-LLM 正在迅速加入多 GPU 以及多节点的支持。一些尚未支持的部分正在不断扩展支持，在这方面需要处理的问题包括如何进行 Tensor 并行以及各种并行操作、如何处理卡间通信比如 NCCL 相关的问题等。RAG
(Retrieval-Augmented Generation)
下面介绍 RAG
(Retrieval-Augmented Generation)。
RAG 的主要目标是处理大型语言模型由于专业领域知识不足以及其他一些原因导致的幻觉问题。大语言模型不可能一直进行训练和微调，因为那样的代价相当大，所以训练和微调通常是周期性而不是一直进行的。在这些周期之间，如何确保大语言模型能够迅速适应专业领域知识呢？RAG 就是用来解决这个问题的。这里通过一个小视频来演示 NeMo 提供的解决方案。首先将我们的 Knowledge Base，也就是与专业领域相关的整个文献，输入到系统中，作为其知识输入。这样，我们就为大语言模型提供了专业领域的知识。然后，输入查询（query），经过一系列的处理过程，在处理完毕后，再输入我们的问题。可以输入一个相对专业的问题，比如询问 NVIDIA 的 Grace 芯片有多少个 CPU core，显然，对于一个预训练模型来说，它的回答肯定是不准确的，因为它没有这方面的知识。但由于刚刚输入了 Knowledge Base 对其进行了增强，现在大语言模型就能够更清楚地回答这个问题了，它有 200 个核心。这个示例直观地展示了 RAG 的作用。
上图展示了应用 RAG 的过程，图的右边展示了 RAG 过程中涉及的模块或参与的组件。左侧展示了 RAG 的过程，这里将 Knowledge Base 文件作为输入，经过预处理模块进行处理，主要包括特征提取等操作，我们使用了 LlamaIndex，当然也有很多其它工具可以在这个步骤中发挥作用。处理完的数据进入 Embedding
model，在这个产品中我们使用了 E5 模型。将刚才提到的 Knowledge
Base 作为模型的输入，进行微调。微调完成后，将 embedding 结果作为向量数据库的输入。向量数据库使用的是 Milvus，而 RAFT 是用于加强当前数据库搜索过程的工具，主要引入了 GPU 加速的功能。接下来是大语言模型 LLM，这里使用的是 Llama 2 模型。通过这样的流程来回答我们的问题，相当于在经过大语言模型之前，先在 prompt 之外添加了一个上下文，使得大语言模型的反馈更准确。这就是 RAG 的过程。
上图完整展示了大语言模型 LLM 的服务过程以及 RAG 在其中的位置，以便大家更好地理解 RAG。系统首先通过 Nvidia 相应的工具来进行 RAG 的处理，然后结合大型模型进行推理。接着，将结果反馈给用户的查询或提示，这使得我们的大语言模型能够真正朝着实用方向发展，不会出现幻觉或者在处理专业领域知识时无法精确匹配的问题。
- 第一步，对输入的 Knowledge Base 进行分块分片，因为输入的数据通常有一定的格式，需要将其处理成类似于问答对或其他固定格式长短的数据，以便进一步处理。
- 第二步，对 Embedding 模型基于刚才输入的 Knowledge Base 进行 Fine-tune，并进行其他处理，如索引等。
- 第四步，进行 Top K 的选择，选出最符合条件的多条 chunk。
- 第五步，将这些 chunk 作为输入，与 prompt 一起传递给大语言模型，得到更加精准且匹配的反馈和回答。
这样大语言模型就能够真正实现落地，更贴近于当前的使用场景。
分享嘉宾
INTRODUCTION
Jack He （何成杰）
NVIDIA 资深解决方案架构师
熟悉 NVAIE 软件栈，特别对集群部署和运维，基于大模型的训练和推理有深入经验。