TensorRT-LLM 检查点

概述

TensorRT-LLM的早期版本（0.8版本之前）是在非常紧迫的时间表下开发的。对于这些版本，重点并未放在定义统一的工作流程上。现在，TensorRT-LLM已经达到了一定程度的功能丰富性，开发团队决定投入更多精力来统一TensorRT-LLM的API和工作流程。本文档记录了围绕TensorRT-LLM检查点的工作流程以及用于生成检查点、构建引擎和评估引擎的CLI工具集。

工作流程中有三个步骤：

1. 将不同源框架的权重转换为TensorRT-LLM检查点。

2. 使用统一的构建命令将TensorRT-LLM检查点构建为TensorRT引擎。

3. 将引擎加载到TensorRT-LLM模型运行器，并使用不同的评估任务进行评估。

NeMo -------------
                  |
HuggingFace ------
                  |   convert                             build                    load
Modelopt ---------  ----------> TensorRT-LLM Checkpoint --------> TensorRT Engine ------> TensorRT-LLM ModelRunner
                  |
JAX --------------
                  |
DeepSpeed --------

准备TensorRT-LLM检查点

TensorRT-LLM 旨在支持不同的来源：

1. 来自NVIDIA NeMo、Microsoft DeepSpeed和JAX的训练模型

2. 来自NVIDIA Modelopt的量化模型

3. 来自HuggingFace的流行模型

TensorRT-LLM 定义了其自己的检查点格式。一个检查点目录包括：

1. 一个配置文件 json，其中包含多个模型超参数。

2. 一个或多个等级权重文件，每个文件包含一个张量（权重）的字典。 在多GPU（多进程）场景中，不同的文件由不同的等级加载。

配置

字段	类型	默认值
架构	字符串	强制性的
数据类型	字符串	强制性的
logits_dtype	字符串	‘float32’
词汇量大小	整数	强制性的
最大位置嵌入	整数	空值
隐藏大小	整数	强制性的
隐藏层数量	整数	强制性的
注意力头的数量	整数	强制性的
num_key_value_heads	整数	注意力头的数量
隐藏操作	字符串	强制性的
中间大小	整数	空值
norm_epsilon	浮点数	1e-5
位置嵌入类型	字符串	‘learned_absolute’
mapping.world_size	整数	1
映射.tp_size	整数	1
mapping.pp_size	整数	1
量化.量化算法	字符串	空值
量化.kv缓存量化算法	字符串	空值
量化组大小	整数	64
quantization.has_zero_point	布尔	假
量化前的比例	布尔	假
quantization.exclude_modules	列表	空值

mapping.world_size 表示 mapping 是一个包含 world_size 子字段的字典。

{
    "architecture": "OPTForCausalLM",
    "mapping": {
        "world_size": 1
    }
}

支持的量化算法列表：

• W8A16

• W4A16

• W4A16_AWQ

• W4A8_AWQ

• W4A16_GPTQ

• FP8

• W8A8_SQ_PER_CHANNEL

支持的KV缓存量化算法列表：

• FP8

• INT8

配置字段是可扩展的，模型可以添加其自己的特定配置字段。 例如，OPT模型有一个do_layer_norm_before字段。

以下是模型特定的配置列表：

字段	类型	默认值
OPT
do_layer_norm_before	布尔	假
猎鹰
偏差	布尔	真
新解码器架构	布尔	假
parallel_attention	布尔	假

排名权重

与PyTorch类似，张量（权重）名称是一个包含层次信息的字符串，它唯一映射到TensorRT-LLM模型的某个参数。

例如，OPT模型的每个变换器层包含一个Attention层，一个MLP层，以及两个LayerNorm层。

注意力权重

Attention 层包含两个 Linear 层，qkv 和 dense；每个 Linear 层包含一个权重和一个偏置。 总共有四个张量（权重），它们的名称是：

• transformer.layers.0.attention.qkv.weight

• transformer.layers.0.attention.qkv.bias

• transformer.layers.0.attention.dense.weight

• transformer.layers.0.attention.dense.bias

其中 transformer.layers.0.attention 是前缀名称，表示权重/偏置位于第0个transformer层的Attention模块中。

MLP权重

MLP 层还包含两个 Linear 层，fc 和 proj；每个 Linear 层包含一个权重和一个偏置。 总共有四个张量（权重），它们的名称是：

• transformer.layers.0.mlp.fc.weight

• transformer.layers.0.mlp.fc.bias

• transformer.layers.0.mlp.proj.weight

• transformer.layers.0.mlp.proj.bias

其中 transformer.layers.0.mlp 是前缀名称，表示权重/偏置位于第0个transformer层的MLP模块中。

LayerNorm 权重

两个LayerNorm层，即input_layernorm和post_layernorm，每个都包含一个权重和一个偏置。 总共有四个张量（权重），它们的名称是：

• transformer.layers.0.input_layernorm.weight

• transformer.layers.0.input_layernorm.bias

• transformer.layers.0.post_layernorm.weight

• transformer.layers.0.post_layernorm.bias

其中 transformer.layers.0.input_layernorm 和 transformer.layers.0.post_layernorm 是两个 layernorm 模块的前缀名称。

KV缓存量化缩放因子

如果我们对模型进行量化，将会产生不同的张量（取决于所应用的量化方法）。 例如，如果我们对KV缓存进行量化，Attention层将会有这个额外的缩放因子：

• transformer.layers.0.attention.kv_cache_scaling_factor

FP8 量化缩放因子

这里是attention.qkv线性层的FP8缩放因子：

• transformer.layers.0.attention.qkv.activation_scaling_factor

• transformer.layers.0.attention.qkv.weights_scaling_factor

AWQ量化缩放因子

这里是 mlp.fc 线性层的 AWQ 缩放因子：

• transformer.layers.0.mlp.fc.weights_scaling_factor

• transformer.layers.0.mlp.fc.prequant_scaling_factor

  注意

  TensorRT-LLM 检查点中的线性权重始终遵循 (out_feature, in_feature) 形状，而 TensorRT-LLM 中由插件实现的一些量化线性可能使用 (in_feature, out_fature) 形状。trtllm-build 命令添加了一个转置操作来对其进行后处理。

示例

让我们以OPT为例，使用张量并行2来部署模型：

cd examples/opt
python3 convert_checkpoint.py --model_dir ./opt-125m \
                --dtype float16 \
                --tp_size 2 \
                --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/

这里是检查点目录：

./opt/125M/trt_ckpt/fp16/1-gpu/
    config.json
    rank0.safetensors
    rank1.safetensors

这里是 config.json:

{
    "architecture": "OPTForCausalLM",
    "dtype": "float16",
    "logits_dtype": "float32",
    "num_hidden_layers": 12,
    "num_attention_heads": 12,
    "hidden_size": 768,
    "vocab_size": 50272,
    "position_embedding_type": "learned_absolute",
    "max_position_embeddings": 2048,
    "hidden_act": "relu",
    "mapping": {
        "world_size": 2,
        "tp_size": 2
    },
    "use_parallel_embedding": false,
    "embedding_sharding_dim": 0,
    "do_layer_norm_before": true,
}

将检查点构建到TensorRT引擎中

TensorRT-LLM 提供了一个统一的构建命令：trtllm-build。在使用它之前，您可能需要将其添加到 PATH 中。

export PATH=/usr/local/bin:$PATH

trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \
                --gemm_plugin float16 \
                --max_batch_size 8 \
                --max_input_len 924 \
                --max_seq_len 1024 \
                --output_dir ./opt/125M/trt_engines/fp16/2-gpu/

进行评估

mpirun -n 2 --allow-run-as-root \
    python3 ../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/2-gpu/ \
                        --batch_size 1 \
                        --test_trt_llm \
                        --hf_model_dir opt-125m \
                        --data_type fp16 \
                        --check_accuracy \
                        --tensorrt_llm_rouge1_threshold=14