专栏名称: 芯东西

芯东西专注报道芯片、半导体产业创新，尤其是以芯片设计创新引领的计算新革命和国产替代浪潮；我们是一群追“芯”人，带你一起遨游“芯”辰大海。

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

芯东西 · 公众号 · · 2024-05-29 18:32

正文

阿里安全采用 NVIDIA NeMo 框架和 TensorRT-LLM 大语言模型推理加速库，显著优化了模型训练与推理性能。

作者 | 刘彬（花名：慧原）

阿里安全算法工程平台工程师

彭伟（花名：又可）

阿里安全算法工程平台负责人

前言

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

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

01 .

Transformer 模型 FLOPs

目前市场上主流大模型以 Transformer 网络结构为主，作为阿里安全的工程落地团队来说，全面分析这个模型的结构以和计算其 FLOPs 十分必要。本文首先重温 Transformer 模型的网络结构，结构如图一所示：

▲图一：Transformer 模型的网络结构

Transformer 模型主要包含三个部分：

Embedding 层，模型的输入层，主要功能是将输入的 input_ids 转成 token 的 embedding, 主要是通过查表的方式实现。

DecodeLayer, 主要功能是将输入的向量，经过 attention，MLP 的相关计算，求出 next_token 相关的向量。

OutputLayer, 主要功能是通过一层的 linearLayer 计算，求出 next_token 对应的每个 word 的 logit。

从上述的模型结构中可以看出，主要的计算量在 decodeLayer, 接下来，我们对该网络（以 Llama2 为例）的计算做详细的分析。

符号定义

矩阵乘法 FLOPs：矩阵乘法为乘加过程，用浮点数运算次数 (FLOPs, floating point operations) 表示计算量的大小：

M [n,m] x N[m,p] => 2mnp

上面的矩阵 M 和 N 相乘的 FLOPs 计算量。

Transformer FLOPs 具体计算过程

decode-only Transformer 模型架构包含：

一个 embedding 层

transformer layer 层，每层的构造相同，包括：

Attention 部分

MLP 部分

一个 output layer 输出层，将 hidden states 转化为词表输出概率

因此模型的总 FLOPs 为：

FLOPs = (attention + mlp) * layers + output_layer

其中输入 embedding 层不涉及 FLOPs 计算，因为 embedding 层做的事情只是根据输入 token 选择对应行。

总 FLOPs

以下是 FLOPs 计算结果：

FLOPs 分析

由于l、h、f、V都是模型的参数，这几个变量是固定的，影响模型计算的 FLOPs 的就剩下B、s、s0这三个变量，可以总结如下规律：

batch_size 越大，模型的计算密度越大

seq_len 越长，first_token 的推理计算密度越大

推理阶段 first_token 是 compute_bound, 后续的 token 推理是 memory_bound

当s0和s相等的时候，第一个 token 的推理计算量等于后续所有的 token 的计算之和，因此从这里可以看出，在推理过程中，kv_cache 的 prefill 的代价是比较高的，因此推理的过程中需要考虑到这个因素，才能更加高效的利用 GPU。

02 .

大模型训练

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

为什么选择 NVIDIA NeMo 框架（Megatron-LM）训练

目前开源市场使用人数最多的是 DeepSpeed 和 NVIDIA NeMo 框架 (Megatron-LM) 这两种，本文主要从工程角度（训练速度的角度）来分析训练框架的特点，接下来分别介绍这两个框架的主要特征（feature）。

DeepSpeed 是微软开发的一款十分受欢迎的大模型训练框架，feature 有很多，对于训练速度提速这块，本文重点介绍其模型并行策略，主要是 ZeRO 相关 feature。

ZeRO1：将优化器状态（optimizer states）切分到不同的 GPU 上，从而降低单个 GPU 的显存占用，这块节省的显存是最大的。

ZeRO2：优化器状态 (optimizer states) + 梯度 (gradients) 切分到不同的 GPU 上。

ZeRO3：把优化器状态 (optimizer states) + 梯度 (gradients) + 模型参数 (parameters) 切分到不同的 GPU 上。

▲图二：DeepSpeed ZeRO 模型并行【1】

假如 GPU 卡数为 N=64，Ψ 是模型参数，假设 Ψ=7.5B，使用 Adam 优化器，K 是优化器的超参，在 64 个 GPU 下 K=12，则：

如果不用 ZeRO，需要占用 120GB 的显存，一般目前的 GPU 单卡放不下

如果用 ZeRO1，则占用 31.4GB，显存大于 40GB 的显卡即可直接启动训练，并且单机多卡或多机多卡训练的通信量不变

如果用 ZeRO2，则占用 16.6GB，通信量同样不变

如果用 ZeRO3，则占用 1.9GB，但是通信量会变为 1.5 倍

备注：

优化器状态一般包含 FP32 Gradient、FP32 Variance、FP32 Momentum、FP32 Master Parameters。

梯度和模型参数一般是 FP16，所以显存主要占用之一就是优化器。

NVIDIA NeMo 框架 (Megatron-LM) 是 NVIDIA 提供的一个端到端的云原生框架，无论是在本地还是在云上，用户可以灵活地构建、定制和部署生成式 AI 模型。它包含但不限于预训练模型、数据管护工具、模型对齐工具、训练和推理框架、检索增强工具和护栏工具包，为用户使用生成式 AI 提供了一种既方便、又经济的方法，同时，NeMo 也支持多模态模型的训练，包括但不限于 Stable Diffusion, Vision Transformer 等。

本文关注焦点在于大模型训练框架的速度对比，因此只聚焦 Megatron-Core 部分。在使用 NeMo 进行大模型训练过程中，影响训练速度比较大的 feature 主要如下：

张量并行（Tensor Parallelism）, 将矩阵乘法中的 X x A 分别切分成 X x [A1, A2] 或者 B x Y x1b 切分成从而降低每个GPU的显存占用，并且将计算平均切分到各个GPU上。如下图三所示：

▲图三：张量并行（TP, Tensor Parallelism）【2】

流水线并行（Pipeline Parallelism），该方法是将 Transformers 中多个 layer 分别分配到不同的 device 上，从而降低每个 device 的显存占用和计算量，如图四所示，为了进一步降低流水线并行过程中 GPU 气泡，可以将流水线并行过程拆分成更小的 stage。

▲图四：流水线并行（PP, Pipeline Parallelism）【3】

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

Distributed Opitimizer， Megatron-Core 也和 DeepSpeed 一样，把 OptimizerStates 分配到各个 device 中，减少显存的占用。

当然，NeMo 还有其他的一些 feature，比如计算和通信做 overlap，调用基于 CuDNN 实现的 FlashAttention 等等。

训练性能对比

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

▲图五：Megatron-LM 论文中性能对比数据

我们团队在 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）。

03 .

大模型推理

大模型快速部署流程

针对大模型推理，NVIDIA 推出了 TensorRT-LLM，实现了业界领先的性能。经多方比较，TensorRT-LLM 被我们选为构建阿里安全大模型高性能推理的基石。在阿里安全的业务中，由于大模型服务比较多，为了让算法同学可以快速部署大模型，工程团队开发了一系列功能让算法同学可以快速、高效、平稳的部署大模型。部署的具体流程图如下：

▲图六：大模型部署流程

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

模型编译阶段：使用 TensorTR-LLM 将模型编译成 engine 格式。

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

DAG 调试和构建阶段：用户在 DAG 中可以自己构造服务的输入数据，完成对整个 DAG 进行调试。同时调试是实时的，用户可以在不部署的服务的情况下调试模型服务，调试结束后，可以将 DAG 固化，最后构建成真实的服务配置。

K8s 服务部署阶段：我们的模型服务都是 K8s 进行部署，便于快速部署，快速复制扩容等。

大模型服务在线推理架构

大模型推理场景有什么特点

每条请求推理耗时的分布是非常不均匀，耗时短的请求可能在几百毫秒能返回，而耗时长的请求可能需要几十秒才能返回，这种情况在非 LLM 的模型上一般是不存在。

目前的大模型是生成式的模型，在实际应用过程中，会依赖上文的信息，即之前的推理结果需要在后续的推理过程中作为输入再次传进来。

部署的服务要求能做到实时监控内部的健康状态，因为大模型服务在使用 GPU 过程中，GPU 内部的错误（例如，数组越界等）会导致后续所有的计算都会失败，因此推理框架需要能及时感知这种错误，并且让框架做到快速重启。

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

▲图七：模型服务架构图

HttpServer 层

Httpserver 层是服务的前端，它负责服务对外的交互，输入处理和输出处理等事情。

SocketQueue 和 socket 缓冲队列，将所有接收的请求放到缓冲队列中。

Flowlimit 是限流模块，会根据当前服务的负载，以及堆积的请求的状况，对服务进行限流，避免服务堆积过多的请求。

SessionManager 模块的功能主要是管理每个请求的状态，由于异步服务的请求在提交后，需要让外界知晓其当前的处理状态，该模块主要是负责该功能。

Input&Output 模块主要是负责将输入的请求转换成服务内部的结构化对象以及将结果序列化返回给上游。

ContextManager 模块管理服务的上下文，例如，上游可能会传一些额外的字段或者信息，服务的内部处理的时候，可能需要使用到这些额外信息。

ComputeWorker 层

该层是模型服务的核心计算层，服务主要的计算都在这层中完成。

JobManager，后台计算的时候，会将若干个请求当成一个 job 进行处理，JobManager 主要负责对这些 job 状态进行管理。

CallbackManager, 对于需要进行回调返回的请求，该模块在 job 处理完成后，会回调相应的接口对结果进行回调。

LogManager, 本模块对服务的各种处理日志进行管理。

HealtyChecker, 该模块主要负责监控当前服务的健康状态，如果发现服务处理了异常，或者内部某些功能出现了不可恢复的问题后，会向 k8 集群汇报当前服务是不健康的，让集群对本节点进行重启。

ServiceDag, 服务中 op 的调度，负责执行 DAG 中的所有的 op。

ConfigManager, 服务的配置管理。

DynamicBatch, 对于大模型 batch 数据，进行动态调度管理，并且充分发挥 GPU 算力，不仅让 GPU 能够满负载的运行，而且可以避免有冗余的计算，具体的算法逻辑会在后续介绍。

LLMLogits, 该模块主要是注册一些钩子函数到 TensorRT-LLM 中，根据业务的需要是否要对模型中的 logits 进行处理。

TextLLM，对文本大模型一些业务逻辑的封装。

MultiModalLLM，对多模态大模型的一些业务逻辑的封装。

TrtLLMOp, 主要是对 TensorRT-LLM 的接口封装。

StopStrategy，注册一些钩子函数到 TensorRT-LLM 中，可以及时识别出是否要及时中断 generate 的过程，避免算力的浪费。

Basic Component 层

本层是服务中一些基础的模块依赖。

AsyncData，对服务中各种数据层的组件的功能封装。

TensorRT-LLM，大模型的服务底层核心的推理计算依赖于 TensorRT-LLM 的推理，TensorRT-LLM 支持 PagedKVCache、in-flight batching、attentionPlugins、Quant（量化）、TP&PP（tensor 并行和 pipeline 并行）、Rope（rope 位置编码）等 feature。

以上是服务的主要架构，采用该架构基本能 GPU 在服务中做到满负载的工作（凑满 batch 计算），同时可以保障服务的稳定运行。

04 .

大模型服务性能优化篇

动态批处理（Dynamic Batch）

为了进一步提升模型推理的速度，我们仔细分析了大模型在推理过程中一些计算，从上述公式中我们可以看到，要想让推理过程中的计算密度变大，只能调大 batch_size；而 generate 过程中，每个样本的输入长度和输出长度都不太一样，必然会导致算力有浪费的情况，具体的推理过程如下图所示：

▲图八：大模型 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 到达一定限度。针对图九中的案例，我们使用新的调度思想之后的效果如下图所示：

▲图九：启用新的调度策略之后的推理

采用新的调度逻辑后，这个案例的 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 实现良好衔接。

量化原理

大模型的参数相对较大，占用的显存较多，并且大模型在 generate 过程中 KVCache 也需要占用大量的显存，如果对模型的参数及 KVCache 进行量化，可以显著节省显存，进而在实际服务推理过程中增大服务的 batch_size, 提升服务的吞吐量。而量化的核心是获得一个对应的缩放系数，具体公式如下：

Q = R/S + Z, (Quantize)

R = S * (Q-Z), (Dequantize)

R 表示量化前的浮点数

Q 表示量化后的定点数

S（Scale）表示缩放因子的数值

Z（Zero）表示零点的数值

如上述公式所表达的那样，模型量化的本质是将一个 FP32 或者 FP16 的数据映射到一个更低 bit 数值区间，进而减少了数据在 memory 间传输的耗时，然后基于不同的算法使用不同数据精度矩阵乘的硬件（如 Int8 的 Tensor Core），最后反量化回原来的数据精度。

量化方法

在大模型量化实践中，我们主要尝试了对称量化的方法。

Int8 weight only，该方法是对一个权重矩阵求出 S (scale) 和 Z (zero) 后，然后应用上述量化公式进行量化和反量化，S 的求法为： S = max(Wi) ，其中 Wi 是 W 的列向量。

Int4 weight only，该方法和 int8 weight only 类似，只不过量化的数据类型改成 int4。

SmoothQuant int8，该算法认为：常见模型的 X (或者是activation) 存在 outlier 的现象。如果使用 int8 weight only 等方式进行量化，只能利用 FP16 的 Tensor Core，不能使用 int8 的 Tensor Core，也就意味着不能使用同一硬件下的更强算力。

‍ SmoothQuant 的方法是找出一个平滑因子，将模型计算改成：，通过使用平滑因子 s，解决掉 outlier 的数值问题，进而利用 8bit 的 TensorCore 或者 CudaCore 提升 GEMM 的性能的同时达到精度损失较小的目标。

▲图十：SmoothQuant 原理【4】

AWQ量化，作者通过观察发现，模型的权重不是等价重要的，在大部分场景，大概 1% 的显著权重就可以让模型保持住比较好的精度。实现方式是重要权重不量化。

GPTQ 量化，该量化思想来源于 OBQ，利用的二阶导数 (hession 矩阵) 来逐步量化模型的权重 W。

模型量化后的性能表现

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

正文

请到「今天看啥」查看全文