TensorRT-LLM 的内存使用#

本文档总结了 TensorRT-LLM 的内存使用情况，并解答了用户报告的常见问题。

理解推理时 GPU 内存使用#

在推理时，对于从 TensorRT-LLM 模型生成的给定 TRT 引擎，GPU 内存使用主要有 3 个贡献者：权重、内部激活张量和 I/O 张量。对于 I/O 张量，主要的内存占用来自 KV cache 张量。

1. 权重大小#

权重大小取决于模型大小、权重的选定精度和并行策略，是固定的。使用 INT8 或 FP8 等较低精度可以减小权重大小。使用张量并行或流水线并行时，每个 rank 只存储一部分权重。例如，使用 8 路张量并行或 8 阶段流水线并行时，每个 rank 通常只使用模型权重的 1/8。

2. 激活大小#

TensorRT 可以通过根据实时分析和张量大小复用内存来优化内存使用。为了避免运行时内存不足错误并减少切换优化配置文件和更改形状的运行时开销，**TensorRT 在构建时会预先计算激活张量的内存需求**。内存需求是基于优化的 TensorRT 图计算得出的，一个配置文件的内存使用是使用最大张量形状计算的，一个引擎的内存需求是根据不同配置文件之间的最大大小计算的。有外部和内部因素会影响 TensorRT 返回的激活大小，例如网络结构、内核融合、操作调度等。

一旦 TensorRT 引擎构建完成，该引擎的激活内存大小就**无法更改**，可以通过 API trt.ICudaEngine.device_memory_size_v2 进行查询。

实际上，对于给定的模型、指定的精度和并行策略，可以通过调整最大批处理大小、最大输入长度、最大 beam width、最大令牌数、填充移除开关标志、上下文 FMHA 开关标志来调优激活内存使用。以下是一些关于这些值如何影响内存的解释：

  1. 减少构建时最大输入令牌数 (max_num_tokens)

    transformer 网络内部的大多数张量与输入令牌数呈线性关系，因此激活大小将接近 最大输入令牌数 * 某个常数因子，常数因子取决于网络结构和 TRT 内部优化。最大输入令牌数由构建时参数确定，可以通过更改提供给 prepare_inputs 函数的参数（例如 PretrainedModel.prepare_inputs）来影响内存使用，也可以更改示例中使用的 trtllm-build 命令的命令行选项。

    使用 packed tensors 格式并指定 max_num_tokens 时，减小其值也会减小激活内存大小。

    使用 padded tensors 格式时，最大输入令牌数等于 max_batch_size*max_input_len，因此减小 max_batch_size 和 max_input_len 几乎可以线性减小激活内存大小。

    建议使用 packed tensors 格式，因为它可以节省内存和计算量。

    在将张量范围传递给 TensorRT 时，beam width 将折叠到批处理大小维度中，因此减小 max_beam_width 也可以减少内存使用。

  2. 开启上下文 FMHA

    使用 GPT attention 插件时，开启插件的 context_fmha_type 将显着减少内存占用。详情请参阅上下文阶段。当 context_fmha_type 设置为禁用时，插件的工作空间大小将与序列长度呈平方关系。

  3. 张量并行和流水线并行

    TensorRT 会尽可能地在层之间复用内存，例如，在一个 transformer 网络中有 *N* 个解码器块，TRT 不会为每个块分配 *N* 份激活内存副本，因为第一个块中张量的内存在执行后可以释放，内存可以用于后面的块，只需要 1 个块的内存。

    使用张量并行时，一些张量会被分割成更小的块，每个 rank 只持有张量的一个块，每个 rank 的激活内存大小将小于在单个 GPU 上执行网络时的激活内存大小。使用流水线并行时，每个 rank 执行几个解码器块，并且所有张量都是完整大小的张量，因此激活内存大小等于 1 个块的内存大小。因此，当所有其他参数相同时，张量并行通常比流水线并行具有更高的内存效率。

3. I/O 张量#

3.1 运行时和解码器缓冲区（KV cache 张量除外）#

C++ 运行时#

在分配 KV cache 块之前，C++ 运行时会预先分配一定量的 GPU 内存，用于存储 TensorRT 引擎和解耦动态解码器的 I/O 张量，它是根据运行时 max_batch_size 和 max_seq_len 分配的，这样当实际调度了这么多请求时可以避免 OOM。

3.2 KV cache 张量#

C++ 运行时#

  • 当启用分页 KV cache 时

    TensorRT-LLM 运行时在初始化期间会为配置数量的块预先分配 KV cache 张量，并在运行时进行分配。

    创建 Executor 时，KV cache 张量是根据 KVCacheConfig 对象分配的。如果既未指定 maxTokens 也未指定 freeGpuMemoryFraction，KV cache 将默认分配剩余可用 GPU 内存的 90%。当指定了 maxTokens 或 freeGpuMemoryFraction 时，将使用指定的值来计算 KV cache 内存大小。如果两者都指定了，则首先使用 freeGpuMemoryFraction 计算 KV cache 中的令牌数，然后使用该计算结果与 maxTokens 之间的最小值。

    在进行中批量处理中，只要有足够的 KV cache 空间可用，调度器就可以自动调度请求（具体行为取决于调度器策略）。

    如果在未进行中批量处理的情况下在 GptSession（已弃用）中使用分页 KV cache，如果分页 KV cache 对于整个批处理不够大，TensorRT-LLM 可能会报告 OOM 错误，并显示消息“无法分配新块。没有剩余的空闲块”。

  • 当禁用分页 KV cache 时（不推荐，只允许用于已弃用的 GptSession）

    C++ 运行时为每一层分配形状为 [批处理大小, 2, heads, 最大序列长度, 每个头的隐藏维度] 的 KV cache 张量，其中 最大序列长度 由创建 GptSession 时的 GptSession::Config::maxSequenceLength 指定。

内存池#

TensorRT-LLM C++ 运行时使用流有序内存分配器来分配和释放缓冲区，参见 BufferManager::initMemoryPool，它使用 CUDA 驱动程序管理的默认内存池。当 GptSession 对象被销毁时，内存会被返回到内存池，可以被下一个 GptSession 对象实例复用。如果需要进行其他内存分配，内存将从内存池中释放。

然而，即使内存已返回到 CUDA 驱动程序的内存池，nvidia-smi 仍可能显示高内存占用。这不是问题，是预期行为。池中保留和空闲的内存量可以分别通过 BufferManager::memoryPoolReserved()) 和 BufferManager::memoryPoolFree()) 查看。

已知问题#

使用 FP8 GEMM 时，激活内存可能大于理论上的优化内存大小，这将在未来的版本中得到改进。

常见问题#

  1. 如何调试 TensorRT-LLM 的内存使用？

    使用 info 日志级别时，TensorRT 和 TensorRT-LLM 将打印内存使用详细信息。以下是运行时 info 日志级别的一部分日志示例：

    [TensorRT-LLM][INFO] Loaded engine size: 6695 MiB
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 1134.01 MiB for execution context memory.
[TensorRT-LLM][INFO] [MS] Running engine with multi stream info
[TensorRT-LLM][INFO] [MS] Number of aux streams is 1
[TensorRT-LLM][INFO] [MS] Number of total worker streams is 2
[TensorRT-LLM][INFO] [MS] The main stream provided by execute/enqueue calls is the first worker stream
[TensorRT-LLM][INFO] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +0, now: CPU 0, GPU 6678 (MiB)
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 43.29 MB GPU memory for runtime buffers.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 180.30 MB GPU memory for decoder.
[TensorRT-LLM][INFO] Memory usage when calculating max tokens in paged kv cache: total: 79.10 GiB, available: 70.48 GiB
[TensorRT-LLM][INFO] Number of blocks in KV cache primary pool: 4060
[TensorRT-LLM][INFO] Number of blocks in KV cache secondary pool: 0, onboard blocks to primary memory before reuse: true
[TensorRT-LLM][INFO] Max KV cache pages per sequence: 32
[TensorRT-LLM][INFO] Number of tokens per block: 64.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 63.44 GiB for max tokens in paged KV cache (259840).

    您可以看到，运行时发生了几次以 [MemUsageChange] 关键字开头的 GPU 内存分配。

    显示“Total Weights Memory”的行表示权重内存大小，显示“Total Activation Memory”的行表示激活内存大小。

    通常，权重内存大小接近 TensorRT 引擎大小，因为引擎中的大部分内容来自 LLM 网络的权重。

  2. 即使在运行时使用了小的批处理大小和序列长度，为什么内存大小仍然很大？

    如上所述，激活内存大小是根据 TensorRT 引擎构建时的最大张量形状计算的，请尝试减小引擎构建时的参数，例如 max_num_token，详情请参见激活大小。

  3. 为什么可以生成引擎，但在运行时推理会耗尽内存 (OOM)？

    在引擎构建时，TensorRT 会逐层调整内核选择，它不一定分配运行整个引擎所需的所有内存。如果运行单个层所需的激活张量很小，而运行引擎所需的 I/O 张量（如 KV cache）大小很大，构建可能会成功，因为它可能不需要分配大型 I/O 张量，而运行时可能会在分配大型 I/O 张量时因 OOM 错误而失败。

    TensorRT-LLM 提供了一个 check_gpt_mem_usage 实用函数，用于检查给定引擎以及相关的批处理大小、I/O 序列长度等的内存大小上限，当上限检查超出 GPU 物理内存大小时，将打印警告消息。

  4. 对于流水线并行，构建时最大批处理大小是否是微批处理大小的限制？

    是的，在流水线并行模式下，TensorRT-LLM 运行时会将请求批次拆分成微批次，然后按顺序将这些微批次排队到 TRT 引擎中。

    构建时的 max_batch_size 意味着一次引擎入队调用的批处理大小应小于它。拆分成微批次之前的总批处理大小可以大于构建时的 max_batch_size。

    例如，如果您有 4 阶段流水线并行，并打算使用微批处理大小 2 运行引擎，并在一次 generate 调用中运行 16 个微批次（总批处理大小 32）。

    您只需将构建时的 max_batch_size 设置为 2，而不是 32。将构建时的 max_batch_size 设置为 32 将占用近 16 倍的激活内存。