- 选题：本次复赛选题为 NVIDIA TensorRT Hackathon 2023 选题3: 用TensorRT-LLM 优化 examples 目录下的 GPT 模型
- 优化效果：
- 在Docker里面代码编译、运行步骤的完整说明
- 请做到只要逐行运行你给的命令，就能把代码跑起来
请在这一节里总结你的工作难点与亮点。
- 如果使用 TensorRT 进行优化，请介绍一下在模型在导出时、或用polygraphy/trtexec解析时，或在使用TensorRT中，遇到了什么问题并解决了。换句话说，针对这个模型，我们为什么需要额外的工程手段。
- 如果使用 TensorRT-LLM 进行优化，描述以下方面可供选手参考：如果搭建了新模型， 请介绍模型结构有无特别之处，在模型的搭建过程中使用了什么算子，有没有通过plugin支持的新算子。如果支持新feature，请介绍这个feature具体需要修改哪些模块才能实现。如果优化已有模型，请介绍模型性能瓶颈以及解决方法。另外还可以包含工程实现以及debug过程中的难点。
在进行优化之前，首先需要定位出优化点。这个时候需要跑一些 benchmark 以及通过 nsys 对整体耗时及具体 kernel 耗时进行一些分析。
GPT 的计算过程可以将第一次计算（context）和后续自回归解码（generation）区别开来。第一次计算时为整个 input_length 个 tokens 作为输入，input_length 的大小会影响计算耗时。后续自回归解码每个请求输入为 1 个 token（此处取 beam 为 1），循环 output_length 次。用 Tctx 来标记第一次计算耗时，Tgen 来标记自回归解码一个 token 的耗时。不考虑KV 累积导致 Tgen 轻微增加的情况，则整个解码的耗时 T = Tctx + output_length * Tgen
使用 benchmarks/benchmark.py 跑一下 gpt_350m 模型在不同参数下的性能，
- 运行如下命令：
python benchmark.py -m gpt_350m --mode plugin --batch_size "1;8;16;32;64" --input_output_len "64,20;128,20;64,120;128,120" --output_dir gpt2
- 结果如下表所示：
| batch | input_length | output_length | tokens_per_sec | latency(ms) |
|---|---|---|---|---|
| 1 | 64 | 20 | 386.56 | 51.74 |
| 8 | 64 | 20 | 2593.23 | 61.70 |
| 16 | 64 | 20 | 4384.88 | 72.98 |
| 32 | 64 | 20 | 6840.90 | 93.56 |
| 64 | 64 | 20 | 8639.60 | 148.16 |
| 1 | 128 | 20 | 373.96 | 53.48 |
| 8 | 128 | 20 | 2215.36 | 72.23 |
| 16 | 128 | 20 | 3505.70 | 91.28 |
| 32 | 128 | 20 | 4747.70 | 134.80 |
| 64 | 128 | 20 | 5722.28 | 223.69 |
| 1 | 64 | 120 | 393.02 | 305.33 |
| 8 | 64 | 120 | 2764.54 | 346.88 |
| 16 | 64 | 120 | 5002.83 | 383.78 |
| 32 | 64 | 120 | 8524.22 | 450.48 |
| 64 | 64 | 120 | 11932.40 | 643.63 |
| 1 | 128 | 120 | 382.07 | 314.08 |
| 8 | 128 | 120 | 2588.72 | 370.84 |
| 16 | 128 | 120 | 4532.15 | 423.64 |
| 32 | 128 | 120 | 7178.17 | 534.96 |
| 64 | 128 | 120 | 9635.28 | 797.07 |
从表中我们可以看出如下信息：
- batch 从 1 增加到 64 的过程中 tokens_per_sec 一直在增加，即服务的吞吐在增加。batch 由 32 增加到 64 时延迟并没有翻倍，说明整个系统并没有达到 compute-bound 的状态。原因可能在于自回归解码时计算量还不够大。上表中 beam 为 1，当 beam 增大时可能会出现 compute-bound 的情况。
- batch 从 1 变为 8 的时候，延迟只有轻微的增加。进一步说明了在 batch 很小的时候，整个系统处于一种 memory-bound 的状态。
- 在 output_length 为 20 的时候（自回归循环次数少的时候），input_length 从 60 变为 128 的时候，延迟的变化还与 batch 相关。
- 在 batch 较小时，input_length 从 60 变为 128 对延迟影响不大。说明 batch 较小时
Tctx部分的计算还是 memory-bound，主要耗时都花在读取权重上了。 - 在 batch 较大时，input_length 从 60 变为 128 对延迟影响不大。说明此时
Tctx部分的计算已经开始转变为 compute-bound。同时也说明Tctx在T中占比较大。
- 在 batch 较小时，input_length 从 60 变为 128 对延迟影响不大。说明 batch 较小时
- 在 output_length 为 120 的时候，input_length 从 60 变为 128 对延迟的影响很小。这也说明此时即使
Tctx变为 compute-bound，但Tctx在T中占比并不高。
可以使用 nsys 工具来验证我们得到的结论，通过可视化的 timeline 及各 kernel 的耗时来进一步定位出可优化的空间。
- 运行如下命令：
nsys profile -o gpt2_bench python benchmark.py -m gpt_350m --mode plugin --batch_size "1;8;64" --input_output_len "64,20;128,20" --engine_dir gpt2
- 观察不同参数下 timeline 上耗时情况
| (batch, input_len, output_len) | timeline |
|---|---|
| (1, 64, 20) |  |
| (8, 64, 20) |  |
| (64, 64, 20) |  |
| (1, 128, 20) |  |
| (8, 128, 20) |  |
| (64, 128, 20) |  |
观察 output_len 为 20 时不同 batch、input_len 下的 timeline，基本符合之前的猜想。随着 batch 变大，Tctx 在整体耗时中占比越来越大，而 Tgen 的耗时并没有随着 batch 线性增长。
在自回归的时候，每个 step 计算结束时都会检查解码结果来判断是否需要退出循环，导致每个 step 都会有一次显式的流同步，会打断 cpu 和 gpu 的异步计算，观察 timeline 可以发现每个 step 结束之后 gpu 都会有一小段空闲。为了减少 gpu 的空闲时间，在代码中已经采取了双 context 的方案，即在同步之前提前设置好下一个 step 的运行信息，这样能有效减少 gpu 的空闲时间。
同步的目的是为了判断是否需要跳出自回归解码的循环，但仔细想一想就会发现，我们并不需要每一个 step 都判断一次是否需要跳出循环，而是可以间隔几个 step 判断一次。在此优化中，我们添加两个参数来控制同步频率：
skip_step: 表示每次解码的前 skip_step 此循环不检查是否退出，这个参数可以统计一下线上平均输出 tokens 长度或者平均解码 step 次数来设置。由于在 batch 解码时所有的请求都需要按照 batch 中输出最长的 step 来计算，在 batch 比较大的时候 skip_step 往往可以设置得更大一些。check_step: 表示出了 skip_step 之外，后续每循环 check_step 次才检查一次是否要退出循环。如果设置为 1 表示每个 step 都会检查。当然这个设置也可能会导致解码循环的次数变多，极端情况下会导致多计算check_step - 1次。比如某个请求输出包含 41 个 tokens，在 skip_step = 0, check_step = 4 这样的设置下，会循环解码 44 次，当然这只会导致计算次数变多，不会导致解码结果出错。此时需要权衡合适的 check_step 设置，一般来说 output_len 越大加速效果越明显。
主要修改的代码 DynamicDecodeOp 的 forward 计算，原始代码中会将是否结束解码的标志保持在 should_stop (pinned memory) 中，这也导致了每个 step 都需要做流的同步。我们将逻辑改为将解码结束的请求数写到一个 finished_sum (gpu memory) 中，在 python 的循环逻辑中根据 skip_step 和 check_step 的设置来判断是否需要进行同步计算 should_stop，代码逻辑如下：
if step == self.max_new_tokens - 1 \
or (step >= scfg.skip_step and (step - scfg.skip_step) % scfg.check_step == 0):
should_stop = self.finished_sum.cpu().item() == batch_size
else:
should_stop = False通过 nsys 来观察 skip_step=8,check_step==4 时的 timeline，可以发现 gpu 空闲有明显减少，同步次数也变少。
加速效果
| (batch, input_len, output_len) | base | skip_8_check_4 | 加速比 |
|---|---|---|---|
| (1, 64, 20) | 51.72 | 50.03 | 1.034 |
| (8, 64, 20) | 61.36 | 59.60 | 1.030 |
| (64, 64, 20) | 146.05 | 144.64 | 1.010 |
| (1, 128, 20) | 53.32 | 51.63 | 1.033 |
| (8, 128, 20) | 71.47 | 69.90 | 1.022 |
| (64, 128, 20) | 222.42 | 220.10 | 1.011 |
| (1, 64, 120) | 305.35 | 295.67 | 1.033 |
| (8, 64, 120) | 345.49 | 337.53 | 1.024 |
| (64, 64, 120) | 634.63 | 628.59 | 1.010 |
| (1, 128, 120) | 312.75 | 303.46 | 1.031 |
| (8, 128, 120) | 367.25 | 359.55 | 1.021 |
| (64, 128, 120) | 787.35 | 779.07 | 1.011 |
从表中可以看到，此方案减少的延迟与 output_len 相关，output_len 相同时无论 batch 多大，减少的延迟都是一致的。最终的效果是 output_len 越大减少的延迟越多，batch 越小加速比越大。
首先使用 nsys 来统计下在整个 pipeline 中不同 kernel 耗时占比。
nsys profile -o gpt2_raw_v2 python benchmark.py -m gpt_350m --mode plugin --batch_size "1;8;64" --input_output_len "64,20;64,120" --engine_dir gpt2从图中可以看出
- 矩阵乘法占比较大，分散调用了不同的 kernel
- attention 部分计算耗时占比我哦 21.6% 左右
- tensorrt jit 生成的 elementwise（sum 及 gelu）操作占比 为 8.9%
- layer norm 耗时占 6.4%
其中 layer norm 在 tensorrt_llm 中是通过插件来实现的，比较方便我们来进行修改优化。
通过分析 layer norm kernel 源码可以发现有如下可以优化的点：
- layer norm 的 block_size 参数设置有错误，源码中设置为
dim3 block(min(hidden_dim, 1024));，但实际上在能做向量化计算时，应该设置为dim3 block(min(hidden_dim / vec_size, 1024));。GPT-medium 的 hidden_dim 为 1024。由于 Compute Capability 8.6 的卡一个 sm 最多能同时运行 1536 个线程，如果 block 中线程数为 1024，则一个 sm 最多能同时运行 1 个 block。此时如果 block 中线程数能正确设置为 512，则一个 sm 最多能同时运行 3 个 block。A10 有 72 个 sm，在 tokens 数较多的 Tctx 阶段，错误的线程数设置对性能影响较大。 - layer norm 的通用实现在 hidden dim 较小时性能不够好。layer norm 中使用一个 CTA 来计算一行的结果，并考虑了是否使用 shared memory。但在 hidden dim 为 1024 的场景，可以使用一个 warp 来计算一行的结果，直接使用寄存器来做缓存。
- 在加载 gamma 和 beta 参数的时候可以使用 __ldg 来做 cache。
基于此，我们针对 hidden_dim <= 1024 这种情况专门实现了一个优化版本的 kernel，使用一个 warp 来计算一行，并使用寄存器来缓存数据。
加速效果 (skip=0, check=1)
| (batch, input_len, output_len) | base | layer_norm_opt | 加速比 |
|---|---|---|---|
| (1, 64, 20) | 51.72 | 51.79 | 0.998 |
| (8, 64, 20) | 61.36 | 60.55 | 1.013 |
| (64, 64, 20) | 146.05 | 141.47 | 1.032 |
| (1, 128, 20) | 53.32 | 53.30 | 1.000 |
| (8, 128, 20) | 71.47 | 70.45 | 1.014 |
| (64, 128, 20) | 222.42 | 213.79 | 1.040 |
| (1, 64, 120) | 305.35 | 305.31 | 1.000 |
| (8, 64, 120) | 345.49 | 345.42 | 1.000 |
| (64, 64, 120) | 634.63 | 631.26 | 1.005 |
| (1, 128, 120) | 312.75 | 313.06 | 0.999 |
| (8, 128, 120) | 367.25 | 366.26 | 1.003 |
| (64, 128, 120) | 787.35 | 776.98 | 1.013 |
可以看出，优化版本主要是优化了 Tctx 部分 batch 较大时的 layer norm 计算，在 Tgen 部分由于 tokens 较少，使用 warp 的方式计算时起的 block 数较少，相对使用 shared memory 时加速效果并不明显。
从统计到的耗时可以看出，layer norm 计算最大耗时从之前的 158us 降为 33us。
在 FasterTransformer 中可以看到将各类 elementwise 操作进行了合并，但这样的合并在 tensorrt-llm 中并不容易集成。
观察上图中 GPT 一个 block 的计算逻辑调用的 kernel，可以发现 TensorRT 以及自动合并了一些 elementwise 操作，更进一步的，我们能否将这些操作和 GEMM 合并起来呢。通过观察，我们发现在 GPT 计算中包含两种计算模式：
- 图中橙色圈: GEMM + AddBias + Residual
- 图中绿色圈: GEMM + AddBias + GeLU
通过合并这些计算，有可能进一步减少 elementwise 操作的时间。在探索过程中，我们调研了 cutlass 和 cudnn-frontend 两种方案。在 cutlass 中通过添加 EpilogueOp 来在矩阵乘法之后做一些 elementwise 的操作。但在单测时发现有一些尺寸，不管设置怎样的 ThreadblockShape、WarpShape、InstructionShape 都没有拆分成两个 kernel 快。基于此选择了实测性能更好的 cudnn-frontend 方案。我们开发了两个 TensorRT 插件：
- GemmBiasAct：实现 GEMM + AddBias + GeLU 逻辑，由于 cudnn-frontend 中
MatMul操作不支持设置 transpose，所以在模型导出和加载时mlp.c_fc.weight需要做一些转置的修改。另外 cudnn-frontend 需要先 build 生成 plan，执行时再用具体的数据来执行 plan，由于 build 耗时比较长，在代码中需要缓存不同参数的 plan。 - GemmBiasRes：实现了 GEMM + AddBias + Residual 逻辑。按照 GemmBiasAct 的思路来实现，但在实测中发现性能比 base 版本慢了很多，原因在于 base 版本在计算小尺寸 GEMM 的时候采用了 splitK 方案，比 cudnn 中的 GEMM 快很多。基于此我们将
MatMul操作替换成了一个 1*1 的Conv操作，测试下来速度比 base 版本更快。cudnn 中并没有支持 Conv 与 GeLU 的融合，所以 GemmBiasAct 中继续采用MatMul。
在多机或者多卡环境下，GemmBiasAct 中并没有 allgather 或者 allreduce 操作，所以不需要做任何改变。GemmBiasRes 操作 base 版本的计算逻辑为 allreduce(Gemm) + Bias + Res，合并之后将计算逻辑修改为 allreduce(Gemm + beta * Bias + beta * Res)，其中 beta 为 1.0 / tp_size，通过设置 beta 参数能做到在多机多卡环境下和 base 版本计算逻辑一致。
对比上图中 block 内的计算与 base 版本的计算可以发现优化后的版本能有效减少 elementwise 操作的开销。
加速效果 (skip=0, check=1, 关闭 layer norm 加速)
| (batch, input_len, output_len) | base | cudnn-frontend | 加速比 |
|---|---|---|---|
| (1, 64, 20) | 51.72 | 48.81 | 1.060 |
| (8, 64, 20) | 61.36 | 57.74 | 1.063 |
| (64, 64, 20) | 146.05 | 137.44 | 1.063 |
| (1, 128, 20) | 53.32 | 50.27 | 1.061 |
| (8, 128, 20) | 71.47 | 67.19 | 1.064 |
| (64, 128, 20) | 222.42 | 209.82 | 1.060 |
| (1, 64, 120) | 305.35 | 288.35 | 1.059 |
| (8, 64, 120) | 345.49 | 323.76 | 1.067 |
| (64, 64, 120) | 634.63 | 617.25 | 1.028 |
| (1, 128, 120) | 312.75 | 296.34 | 1.055 |
| (8, 128, 120) | 367.25 | 346.82 | 1.059 |
| (64, 128, 120) | 787.35 | 764.64 | 1.030 |
这一部分介绍你的工作在云主机上的运行效果。如果是优化模型，需要分两部分说明：
本工作针对 GPT-Medium 模型在 tensorrt-llm 上实现了三个优化点：
- 优化点1: 减少显式同步的次数
- 优化点2: 优化 layer norm 的计算
- 优化点3: 合并计算减少 elementwise 计算开销
测试命令
python benchmark.py -m gpt_350m --mode plugin --batch_size "1;8;64" --input_output_len "64,20;128,20;64,120;128,120" --engine_dir gpt2 --skip_step=8 --check_step=4开启三个优化后，原始模型与优化模型的 Rouge score 差异(GPT-Medium)。此处延迟高是因为没有做 warmup，包含了 build plan 的时间。
TensorRT-LLM (total latency: 66.98257255554199 sec)
TensorRT-LLM beam 0 result
rouge1 : 20.735151996347536
rouge2 : 6.1022639415477835
rougeL : 16.40108348647125
rougeLsum : 18.184358801291715
Hugging Face (total latency: 14.055776119232178 sec)
HF beam 0 result
rouge1 : 18.182978950152904
rouge2 : 5.166241888544473
rougeL : 14.851620358520162
rougeLsum : 16.95757748412272
| (batch, input_len, output_len) | base | opt | 加速比 |
|---|---|---|---|
| (1, 64, 20) | 51.72 | 46.75 | 1.106 |
| (8, 64, 20) | 61.36 | 55.95 | 1.097 |
| (64, 64, 20) | 146.05 | 132.19 | 1.105 |
| (1, 128, 20) | 53.32 | 48.25 | 1.105 |
| (8, 128, 20) | 71.47 | 64.79 | 1.103 |
| (64, 128, 20) | 222.42 | 202.71 | 1.097 |
| (1, 64, 120) | 305.35 | 278.99 | 1.094 |
| (8, 64, 120) | 345.49 | 318.35 | 1.085 |
| (64, 64, 120) | 634.63 | 604.63 | 1.050 |
| (1, 128, 120) | 312.75 | 287.21 | 1.089 |
| (8, 128, 120) | 367.25 | 339.14 | 1.083 |
| (64, 128, 120) | 787.35 | 750.78 | 1.049 |
通过优化在不同的参数下获得 4.9%～10.6% 的加速。
请注意：
- 相关测试代码也需要包含在代码仓库中，可被复现。
- 请写明云主机的软件硬件环境，方便他人参考。
- 硬件: A10 gpu, 16 核 Intel(R) Xeon(R) Platinum 8369B
- 软件镜像: registry.cn-hangzhou.aliyuncs.com/trt-hackathon/trt-hackathon:final_v1
- TensorRT 在转换包含 Trilu 的 onnx 模型时可能会获得错误的结果
- tensorrt_llm 中 layer norm 插件的 USE_DIFF_OF_SQUARES 实现可能导致结果出现 nan
A: chef before moving to London in the early
Q: python3 summarize.py --engine_dir trt_engine/gpt2/fp16/1-gpu --test_hf --batch_size 1 --test_trt_llm --hf_model_location=gpt2 --check_accuracy --tensorrt_llm_rouge1_threshold=14
A:
TensorRT-LLM (total latency: 2.840153455734253 sec)
TensorRT-LLM beam 0 result
rouge1 : 21.869322054781037
rouge2 : 6.258925475911645
rougeL : 16.755771650012953
rougeLsum : 18.68034777724496
Hugging Face (total latency: 14.237422227859497 sec)
HF beam 0 result
rouge1 : 18.182978950152904
rouge2 : 5.166241888544473
rougeL : 14.851620358520162
rougeLsum : 16.95757748412272
欢迎在这里总结经验，抒发感慨。