PyTorch训练中如何优化cudaHostAlloc与cudaLaunchKernel耗时

1. 两个CUDA操作的具体作用 #

• cudaHostAlloc：是CUDA API中用于分配**页锁定主机内存（pinned memory）**的函数。普通可换页内存无法被CUDA直接访问，传输到GPU前需要先拷贝到页锁定内存，因此页锁定内存可以大幅提升CPU-GPU的数据传输效率，但分配操作本身的耗时远高于普通内存分配。你的profiling结果中31次调用就占了36%的CPU时间，符合该函数的特性。
• cudaLaunchKernel：是CUDA API中用于将GPU核函数提交到GPU任务队列的CPU侧函数。每次调用GPU算子（如卷积、矩阵乘）都需要CPU侧执行一次该函数提交任务，耗时为纯CPU侧开销，和GPU核函数本身的执行时间无关。你的结果中该函数调用量高达10861次，占近36%的CPU时间，说明当前任务存在严重的CPU提交瓶颈。

2. 耗时优化方案 #

针对cudaHostAlloc的优化 #

• 减少重复分配：尽量复用已经分配好的页锁定内存，不要每次推理都重新分配。如果是数据加载阶段的内存分配，直接使用PyTorch DataLoader的pin_memory=True参数，PyTorch会自动复用页锁定内存池，避免每次迭代都调用cudaHostAlloc。
• 控制页锁定内存分配量：仅给需要频繁传输到GPU的核心数据分配页锁定内存，非核心数据使用普通内存即可，降低单次分配的耗时。

针对cudaLaunchKernel的优化 #

该操作的调用量过高是当前最核心的性能瓶颈，可从以下方向优化：

• 算子融合：把多个小算子合并为一个大算子，减少核函数调用次数。可以直接使用PyTorch 2.0+的torch.compile功能自动完成算子融合，无需手动修改代码，通常可减少30%以上的核函数调用量。
• 减少不必要的CPU-GPU交互：检查模型推理逻辑中是否存在频繁把GPU数据拷回CPU、或者在CPU侧执行大量小计算的逻辑，这类逻辑会触发大量额外的核函数调用。
• 批量操作替代循环操作：如果代码中存在Python层面的循环，且每次循环都调用GPU小算子，将这类循环逻辑改为向量化的批量操作，用一次大算子提交替代多次小算子提交。
• 启用CUDA Graph：对于输入输出固定、计算逻辑稳定的推理场景，用CUDA Graph提前录制整个推理流程的核函数，后续推理仅需提交一次启动请求即可运行整个计算图的所有核函数，可以几乎消除cudaLaunchKernel的开销。

3. 易遗漏的标准优化操作 #

从你的profiling结果判断，以下标准优化大概率未执行：

• 未启用torch.compile优化推理流程
• 数据加载未开启pin_memory=True，或数据加载逻辑存在重复分配页锁定内存的问题
• 模型推理未适配静态输入shape，导致每次推理都需要重新分配内存、重新调度算子
• 模型未做量化、剪枝等优化，小算子数量过多，导致核函数调用量飙升

内容的提问来源于stack exchange，提问作者Nagabhushan S N