利用 FP8 量化加速 Llama-3-70B 推理

本文针对在 Amazon P5 (H100 GPU) 上部署 Llama-3-70b FP8 精度模型的两种方案选型（Trion 及 LMI – Large Model Inference 容器）进行了探索，同时提供了基于 FMBench 的性能评估的最佳实践，TensorRT-LLM 引擎的优化建议，以及快速上线生产应用的最佳实践。

Llama-3 概要

Meta 基于 Llama 2 架构推出了四款新一代 Llama 3 开放型大语言模型，分为 8B 和 70B 两种参数规模，每种规模提供预训练基础版和指令微调版，上下文长度达 8k Token。最显著的变化是采用新 Tokenizer 将词汇表扩大至十余万个 Token，能更高效编码文本并提升多语种处理能力。同时该版本还引入高效的分组查询注意力（GQA）以处理更长上下文。Llama 3 模型在数万卡的 GPU 集群上训练，并使用了规模超过 15 万亿 Token 的新公共在线数据集，更大规模且更细致的数据规划可能是其性能提升的重要因素。针对对话应用的 Llama 3 Instruct 版本则融合了千余万条人工标注数据，通过监督微调、拒绝采样、邻近策略优化和直接策略优化等技术进行优化训练，以提升对话交互质量（参考 HuggingFace Llama 3 Blog）。

FP8 在推理场景的优势

Nvidia 在 Hopper（H100 on Amazon P5）和 Ada Lovelace 架构 GPU（L4 on Amazon G6）上推出了 FP8 Tensor Core。FP8 是一种 8 位浮点数据格式，与同为 8 位的 INT8 量化相比，具有比 INT8 更高的动态范围，使其非常适合量化大型语言模型（LLM）的多个组件。同时，相比于 16 位的浮点表示如 FP16 或 BF16，FP8 在数据表示位数降低，进一步减少显存占用及通信带宽要求，因此带来了更大的吞吐和更高的计算性能（参考 FP8 New Frontier）。

H100 on Amazon P5

Amazon P5 实例由最新的 NVIDIA H100 Tensor Core GPU 提供支持，为深度学习（DL）和高性能计算（HPC）应用程序提供极高性能。同时可以使用 Amazon SageMaker 托管的 P5 实例来快速搭建不同类型的训练集群（如 Training Job 及 HyperPod）及推理集群（Endpoint）。H100 GPU 除去新增了 FP8 Tensor Core 之外，性能上也有大幅性提升。以下通过不同方式对 Llama-3-70B 的推理部署，展示基于 P5 实例进行 FP8 精度模型部署的最佳实践。

模型部署

以 Llama-3-70B 基于 TensorRT-LLM FP8 量化及其在 Triton 以及 Large Model Inference（LMI）推理容器的部署为例。

TensorRT-LLM 介绍

TensorRT-LLM 是一个易于使用的 Python API，用于定义大型语言模型（LLM）并构建包含优化的 TensorRT 引擎，以在 NVIDIA GPU 上高效执行推理。TensorRT-LLM 包含用于创建执行 TensorRT 引擎的 Python 和 C++ 运行时的组件。它还包括一个用于与 NVIDIA Triton 推理服务器集成的后端，这是一个用于服务 LLM 的生产质量系统。使用 TensorRT-LLM 构建的模型可以在从单个 GPU 到多个节点多个 GPU（使用张量并行和/或流水线并行）的广泛配置上执行。同时，TensorRT-LLM 也针对 H100 GPU 进行了大量性能优化，详情可参考 TensorRT-LLM 加速 H100 推理，H00vsA00。

一、基于 Triton 部署

Triton 是 Nvidia 发布的一个高性能推理服务框架，可以帮助开发人员高效轻松地在云端、数据中心或者边缘设备部署高性能推理服务。Triton Server 可以提供 HTTP/gRPC 等多种服务协议。同时支持多种推理引擎后端，如：TensorFlow, TensorRT, PyTorch, ONNXRuntime 等。Server 采用 C++ 实现，并采用 C++ API调用推理计算引擎，保障了请求处理的性能表现。整体架构示意如下：

环境准备

使用预置 Deep Learning AMI with PyTorch 启动 P5 实例后，并进入预置的 PyTorch 环境：

source activate pytorch
git clone -b v0.X.0 https://github.com/NVIDIA/TensorRT-LLM.git
pip install -r TensorRT-LLM/examples/llama/requirements.txt

使用如 huggingface-cli 等工具将模型下载至本地，如 P5 实例的 32TB 高性能 NVMe SSD 的默认挂载路径 /opt/dlami/nvme 。

模型量化

TensorRT-LLM 同时具备 TP 及 PP 的模型切分方式，可以直接通过配置的形式进行指定

python3 TensorRT-LLM/examples/quantization/quantize.py --model_dir $hf_model_path \
            --output_dir $trt_ckpt_path \
            --dtype bfloat16 \
            --qformat fp8 \
            --kv_cache_dtype fp8 \
            --calib_size N \
            --pp_size 2 \ *
            --tp_size 4

** 仅对 pp_size 的设置进行展示，生产场景需根据实际的性能表现确定模型切分策略*

Engine 编译

trtllm-build --checkpoint_dir $trt_ckpt_path \
            --output_dir $trt_engine_path \
            --gpt_attention_plugin bfloat16 \
            --gemm_plugin bfloat16 \
            --max_num_tokens N \ *
            --use_custom_all_reduce \ **
            --use_fused_mlp \
            --enable_xqa enable \
            --use_fp8_context_fmha enable ***

**use_custom_all_reduce 为 NVLink 相关优化参数，非 NVLink 的 GPU 类型无法使用。

* max_num_tokens 最大 Token 数的估计值

* **use_fp8_context_fmha enable 仅限 Hopper 架构 GPU，其他 GPU 类型需要 disable

tensorrtllm_backend 配置

git clone -b v0.X.0 https://github.com/triton-inference-server/tensorrtllm_backend.git

clone tensorrtllm_backend 项目到本地之后，进行如下一系列的配置。

model_path_name=Meta-Llama-3-70b-Instruct
modelserve_name=inflight_batcher_llama3_fp8opt

# 拷贝一个部署模板
cp -r tensorrtllm_backend/all_models/inflight_batcher_llm tensorrtllm_backend/all_models/$modelserve_name

# 存入前序编译生成的engine
cp $trt_engine_path/ tensorrtllm_backend/all_models/$modelserve_name/tensorrt_llm/1/

# 拷贝tokenizer
mkdir -p tensorrtllm_backend/$model_path_name/
cp $model_path_name/*.json tensorrtllm_backend/$model_path_name/

在目标的推理模板路径中，需要对 ensemble、preprocessing、postprocessing、tensorrt_llm_bls、tensorrt_llm 等路径下的 config.pbtxt 文件进行编辑，并将所需调整的 placeholder 替换为所需的取值。如：

triton_max_batch_size:1024
tokenizer_dir:$HF_model_dif
preprocessing_instance_count:1
tokenizer_type:auto

除上述通用配置参数之外，tensorrt_llm/config.pbtxt 中需进一步指定模型部署的核心配置，示例如下：

engine_dir:$trt_engine_path
kv_cache_free_gpu_mem_fraction:0.95
exclude_input_in_output:True
enable_kv_cache_reuse:False *
batching_strategy:inflight_fused_batching
max_queue_delay_microseconds:0

** 参数 enable_kv_cache_reuse = True 需要在模型编译时配置参数 use_paged_context_fmha enable 配合生效

Triton Server 拉起

使用标准 tritonserver 容器，并基于前序的 serving 配置启动服务。

docker run --rm -it \
    --net host --shm-size=2g \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 --gpus all \
    -v /opt/dlami/nvme/tensorrtllm_backend:/tensorrtllm_backend \
    nvcr.io/nvidia/tritonserver:24.04-trtllm-python-py3 bash

python3 /tensorrtllm_backend/scripts/launch_triton_server.py \
    --world_size=8 \
    --model_repo=/tensorrtllm_backend/all_models/$modelserve_name

其中，/opt/dlami/nvme 为 P5 实例上 32TB NVMe SSD 的默认挂载路径。

小型 Python Client 搭建及模型调用

import requests

url = "http://<YOUR_PUBLIC_HOST_NAME>:8080/v2/models/ensemble/generate"

def send_request(payload):
    response = requests.post(url, 
                headers={"content-type": "application/json"}, 
                json=payload)
    return response.json()
    
payload= {
    "inputs": "What is Deep Learning? ",
    "parameters" : {
            "do_sample": True,
            "max_new_tokens": 256,
            "details": False,
            "temperature": 0.2
        },
}

send_request(payload)

Triton 同时支持 HTTP 及 gRPC 的请求形式，参考 Inference Request 接口。

生产环境下的 Client 构建

在实际生产环境中，由于客户端位于 Internet 的另一侧，客户端是否是合法访问，什么时候有多少访问量等都是未知的。不合法的访问，容易造成信息泄漏等安全隐患，同时这些不合法的访问还会消耗服务端的算力，在服务端算力有限的情况下，容易造成正常的访问会由于算力不足被拒绝。并且，突然的高并发访问，也会导致服务端消耗完全部算力而无法正常提供服务。因此，需要在客户端访问服务端时进行鉴权、验证等，以确保本次访问是来源于合法的客户端，并且未被“中间人”篡改信息。

对于高并发场景，在客户端一般是采用异步调用以及重试、指数回退等方式，在服务端一般是采用队列的方式，将积压的访问放在队列里，并由服务端进行队列中请求的编排处理。因此，需要在客户端构建重试、指数回退等能力配合服务端的队列机制，以保障在突发的高并发场景下，不会出现大面积的客户端无法访问，服务端由于算力不足而崩溃。

二、基于 LMI（Large Model Inference）容器部署

LMI

LMI 容器是一组专为大型基础模型推理而打造的高性能 Docker 容器。使用 LMI 可以快速利用高性能的开源推理库，如 LMI-Dist 自研推理引擎、vLLM、TensorRT-LLM 和 Transformers NeuronX。借助 LMI 内置的推理处理程序和统一配置，只需更改配置即可在 Amazon SageMaker 端点上部署大语言模型。LMI 容器将模型服务器与开源推理引擎捆绑在一起，提供包括针对 Llama、Mixtral 等流行模型架构优化的推理性能，连续批处理，以最大化高并发下的吞吐量；以及对多种量化选型，及 Multi-Lora 的支持。同时仍在快速的性能提升迭代中。模型部署示例见示例 Repo。

SageMaker

Amazon SageMaker 是一个全生命周期，全托管的机器学习平台，为算法工程师和数据科学家提供了端到端的解决方案。它不仅消除了管理维护底层基础设施的复杂性，并提供了不同类型的平台环境（如 Jupyter 集成开发环境，按需及固定集群、同步及异步推理集群等），使得模型开发、训练和部署变得高效。另外，其全面的监控和日志功能则确保了模型的持续优化和维护。通过这些特性，SageMaker 使得算法专家能够专注于模型的创新和优化上，极大缩短了机器学习项目的开发周期。以下展示了基于 LMI 容器内置的 TensorRT-LLM Engine 在 SageMaker 推理端点上部署 Llama-3-70B 的示例。

部署模型到 SageMaker 的推理端点

LMI 容器所需的 serving.properties 配置参数：

engine=MPI
option.model_id=nousresearch/meta-llama-3-70b # AWS S3 URI or HF MODEL_ID
option.tensor_parallel_degree=8
option.rolling_batch=trtllm
# Adjust the following based on model size and instance type
option.max_num_tokens=100000
option.quantize=fp8
option.use_fp8_context_fmha=true

部署至 SageMaker 托管的 H100 （ml.p5.48xlarge）服务器 / 集群：

model = Model(image_uri=<DJL-LMI Latest Image URI>, 
              model_data=<S3 URI contains above setting>
              )
model.deploy(initial_instance_count=<N>, instance_type='ml.p5.48xlarge')

推理端点调用

使用 SageMaker Python SDK 对推理端点调用：

predictor = sagemaker.Predictor(
    endpoint_name=<ENDPOINT_NAME>,
)

predictor.predict({"inputs": <YOUR_PROMPT>, 
            "parameters": {"temperature":0.1, "max_new_tokens": 1024}}
         )

可以看出，利用LMI 容器可以通过非常简化的 Low-code 形式进行配置，简便地使用 LMI 自研的 LMI-Dist 自研推理引擎，vLLM 或前序介绍的 TensorRT-LLM（支持 AOT 及 JIT 编译，平衡部署复杂性及推理集群的扩缩容性能）等流行的推理引擎；并可以使用 SageMaker Python SDK（或 boto3 Python SDK）快速进行模型推理。

推理端点的安全集成及集群托管

基于 AWS SDK 的推理调用，提供了多种安全功能，以帮助开发人员构建安全的应用程序。例如：默认启用 HTTPS/TLS 加密，确保数据传输安全；自动为API请求添加签名，使用 AWS Signature Version 4 (SigV4) 算法，确保请求的完整性和身份验证，不再需要再单独考虑签名、鉴权等安全问题；与 Amazon CloudTrail 集成，支持对API调用进行审计和日志记录，有助于安全分析和合规性要求。

此外，AWS SDK 底层集成了完善的重试和指数回退机制，以帮助应用程序更好地处理临时故障和网络问题。AWS SDK 使用截断的指数回退算法来计算重试延迟，并添加了抖动（jitter）来避免多个客户端同时重试造成的“惊群效应”。

集群管理方面，可直接借助 SageMaker 的集群托管能力，快速实现对推理集群的 Auto-scaling，智能负载均衡，灰度发布等能力，更加高效地搭建生产级别的模型服务及应用并快速投产。

最佳实践

TensorRT-LLM 编译优化项

以上的示例中，在 Trition 及 LMI 容器中，均使用了 TensorRT-LLM 作为模型的推理引擎。TensorRT-LLM 在近期 版本发布中涉及到的重点编译优化项如下：

--max_num_tokens 最大 Token 数的预估值，同时影响 KV Cache 显存分配及 GPU 利用率

--use_fused_mlp 启用矩阵乘法融合

--enable_xqa enable 启用 XQA kernel，优化 MQA 及 GQA 性能

--use_fp8_context_fmha enable 基于 FP8 Context 的融合注意力优化，仅限 Hopper 架构 GPU，如 Amazon P5 实例

--gemm_swiglu_plugin fp8（Preview) 使用 NVIDIA cuBLASLt 加速库进行矩阵乘法，仅限 Hopper 架构 GPU

同时，如果需要在 Triton Server 的推理配置中使用到 enable_kv_cache_reuse，需要对应开启如下选项：

另外，对于开启 FP8 精度推理下，结合其他编译优化参数的性能最佳实践，参考 TensorRT-LLM 性能调优最佳实践。

基于 FMBench 的推理性能基准测试

实际生产场景中，模型的推理环境及上下文的差异可能导致不可忽略的性能差异。因此建议基于真实场景的流量情况，对大模型的推理进行压测及评估。使用 FMBench（Foundation Model Benchmark Tool）能够帮助我们快速进行针对基础模型的评估，帮助基于推理延迟、吞吐量和成本来综合评估推理性能，从而为给定工作负载提供最佳性价比的资源组合（可参考 FMBench 在 Meta 的使用案例）。FMBench 默认采用了 LongBench 数据集，此外，仍可以通过自定义数据集等形式，使用实际业务数据进行验证。示例如下：

1. 自定义 dataset，并组装为 .jsonline 格式，key 与模板文件中的待填充变量一致

{
    "question": "`Summarize the discussion about the working design of the remote control.`",
    "context": "`Project Manager: I hope you're ready for this uh functional design meeting . Marketing: Of course...`"
}

2. 自定义模板文件

<|begin_of_text|><|start_header_id|>user<|end_header_id|>

You are an assistant for question-answering tasks. Use the following pieces of context in the section demarcated by "```" to answer the question. 

```
{context}
```

Question: {question} 

<|eot_id|><|start_header_id|>assistant<|end_header_id|>

3. 同时在测试配置的 .yml 文件中，使用调整 filters 及 inference_parameters 中的推理参数来进一步对输入输出的长度进行控制。

运行完成后，获取不同并发请求下的关键性能指标如 TTFT， prompt_token_throughput，completion_token_throughput ，latency，TPM 等。

Kernel 调用分析

针对推理过程的性能分析，除了利用 FMBench 等工具获取客户端到模型服务器的端到端性能之外，可以进一步通过 profiling 来分析 GPU 上的 op 或 kernel 的消耗情况。H100 GPU 在针对上一代的全面能力提升之外，也新增了对于 FP8 Tensor Core 的支持。针对 TensorRT-LLM FP8 精度下模型推理的 profiling metric，列举了以下几个 Top GPU 消耗（不含 kernel launch）的关键内核差异。

通过对关键执行内核的对比，可以看出在 H100 GPU 上，针对 gemm 的常见的 Warp Matrix Multiply Accumulate (wmma) 变为了新的 xmma。同时对于各类注意力计算，TensorRT-LLM 均提供了特定的内核供执行，如不同精度下的 mmha。值得注意的是，对于 FP8 量化模型的 fmha 计算，从其执行内核的命名可以看出其默认仍使用 bf16 精度，与非量化版本的模型中使用的内核一致。而仅当编译选项中增加了特定参数后，从该内核的命名看出其使用了 e4m3 版本的 FP8 精度参与计算。

总结