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

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

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

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

• 一个 embedding 层

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

• Attention 部分

• MLP 部分

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

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

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

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

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

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

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

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

• 如果不用 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，所以显存主要占用之一就是优化器。

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

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

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

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

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

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

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

• 影响训练性能的三个重要参数：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）。

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

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

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

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

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

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

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

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

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

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

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

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

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

• 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 的过程，避免算力的浪费。

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

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

• R 表示量化前的浮点数

• Q 表示量化后的定点数

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

• Z（Zero）表示零点的数值

• 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 的性能的同时达到精度损失较小的目标。

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

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

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

• 量化后期针对模型在实际业务中存在精度损失问题： 理论上，量化后的模型肯定做不到和未量化的模型的精度相同，但我们在实际业务中能将量化的损失控制在 1% 以内，如果精度损失过大，一般可以调整量化过程中的校准数据集，而校准样本一般 2000-8000 条，而对校准数据集的要求是分布尽可能和实际业务的样本的分布一致。

• 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 中可以添加这样的语句：

• 为了进一步提升大模型的 generate 的速度，后续将会尝试 medusa decode 等解码手段。

• 后续会考虑尝试使用 Attention Sink 相关 skills 在长文本中尝试。

• 过去 1 年一直没有在 FP8 上进行尝试，后续将会用 FP8 在实际业务中实践。

• RAG 的应用，目前大模型能力有一定限制，后续将会尝试使用 RAG（检索增强）的方式提升大模型在业务中效果。

• 持续关注 TensorRT-LLM 最新的进展和 feature，TensorRT-LLM 这个开源项目更新迭代的频率还是挺快的，基本 2-3 个月就会有一个大版本出来。