TensorRT-LLM的内存使用情况

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

了解推理时GPU内存使用情况

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

1. 权重大小

权重大小取决于模型大小、所选权重精度和并行化策略。使用较低的精度（如INT8或FP8）可以减少权重大小。当使用张量并行或管道并行时，每个等级只存储部分权重。例如，当使用8路张量并行或8阶段管道并行时，每个等级通常只使用模型权重的1/8。

2. 激活大小

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

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

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

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

变压器网络中的大多数张量与输入令牌的数量呈线性关系，因此激活大小将接近max number of input tokens * some constant factor，常数因子取决于网络结构和TRT内部优化。输入令牌的最大数量来自构建时的参数，可以更改提供给prepare_inputs函数的参数，如PretrainedModel.prepare_inputs，以影响内存使用，或者可以更改示例中使用的trtllm-build命令的命令行选项。

当使用打包的张量格式并且指定了max_num_tokens时，减少其值也会减少激活内存的大小。

当使用填充的张量格式时，输入令牌的最大数量等于max_batch_size*max_input_len，因此减少max_batch_size和max_input_len几乎可以线性减少激活内存的大小。

推荐使用打包的张量格式，因为它既节省内存又节省计算资源。

当将张量范围传递到TensorRT时，光束宽度将被折叠到批量大小维度中，因此减少max_beam_width也可以减少内存使用。
开启上下文FMHA

当使用GPT注意力插件时，开启插件的context_fmha_type将显著减少内存占用。详情请参见上下文阶段。当context_fmha_type设置为禁用时，插件的工作空间大小将随序列长度呈二次方依赖。
张量并行和管道并行

TensorRT 会尽可能在层之间重用内存，举一个典型的例子，给定一个变压器网络中的 N 个解码器块，TRT 不会为每个块分配 N 份激活内存，因为在第一个块中的张量内存在执行后可以被释放，内存可以被后续块重用，只需要一个块的内存。

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

3. I/O 张量

3.1 运行时和解码器缓冲区（不包括KV缓存张量）

C++ 运行时

在分配KV缓存块之前，C++运行时会预先分配一定量的GPU内存，用于存储TensorRT引擎和分离的动态解码器的I/O张量。它是根据运行时的max_batch_size和max_seq_len进行分配的，以便在确实有那么多请求被调度时避免OOM。

3.2 KV缓存张量

C++ 运行时

当启用分页KV缓存时

TensorRT-LLM 运行时在初始化期间为配置的块数预分配 KV 缓存张量，并在运行时分配它们。

KV缓存张量在创建Executor时基于KVCacheConfig对象进行分配。如果既未指定maxTokens也未指定freeGpuMemoryFraction，KV缓存将默认分配剩余GPU内存的90%。当指定了maxTokens或freeGpuMemoryFraction时，将使用指定的值来计算KV缓存的内存大小。如果两者都指定了，首先使用freeGpuMemoryFraction来计算KV缓存中的令牌数量，然后使用此计算出的令牌数量与maxTokens之间的最小值。

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

如果在GptSession（已弃用）中使用分页KV缓存而没有进行飞行批处理，如果分页KV缓存不足以容纳整个批次，TensorRT-LLM可能会报告OOM错误，并显示消息“无法分配新块。没有剩余的空闲块”。
当分页KV缓存被禁用时（不推荐且仅允许用于已弃用的GptSession）

C++ 运行时为每一层分配形状为 [batch size, 2, heads, max seq length, hidden dimension per head] 的 KV 缓存张量，其中 max seq length 在创建 GptSession 时由 GptSession::Config::maxSequenceLength 指定。

Python 运行时（不推荐使用）

Python 运行时根据 GenerationSession.setup 函数的参数分配 KV 缓存张量，KV 缓存大小与 batch_size 和 max_context_length+max_new_tokens 线性相关。注意：这可能会在未来发生变化，因为 C++ 运行时的 Python 绑定可能会在未来取代当前的 Python 运行时。C++ 运行时的 Python 绑定行为类似于 C++ 运行时。

内存池

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

然而，nvidia-smi 在内存返回到 CUDA 驱动程序的存储池后，可能仍会显示高内存占用。这不应引起担忧，并且是预期的行为。可以通过 BufferManager::memoryPoolReserved()) 和 BufferManager::memoryPoolFree()) 分别检查池中保留和空闲的内存量。

已知问题

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

常见问题解答

如何调试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网络的权重。

为什么在运行时使用了小批量大小和序列长度，内存大小仍然很大？

如上所述，激活内存大小是根据TensorRT引擎构建时的最大张量形状计算的，尝试减少引擎构建时间参数，如max_num_token，详情请参见激活大小。
为什么可以生成引擎，但在运行时推理会耗尽内存（OOM）？

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

TensorRT-LLM 提供了一个 check_gpt_mem_usage 实用函数，用于检查给定引擎的内存大小的上限，以及相关的批量大小、I/O 序列长度等，当上限检查超过 GPU 物理内存大小时，将打印警告信息。
对于管道并行，构建时的最大批次大小是微批次大小的限制吗？

是的，在管道并行模式下，TensorRT-LLM 运行时会将请求批次分割成微批次，并按顺序将这些微批次排入 TRT 引擎。

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

例如，如果你有4阶段的管道并行性，并打算使用微批次大小为2运行引擎，并在一次generate调用中运行16个微批次（总批次大小为32）。

你可以在构建时将max_batch_size设置为2，而不是32。将构建时的max_batch_size设置为32将占用几乎16倍的激活内存。