趋近智
现代机器学习 (machine learning)运行时系统的一个主要功能是管理在执行开始前其维度无法完全确定的张量。与提前(AOT)编译场景中张量形状通常固定(例如, 图像,批次大小为 32)不同,许多应用涉及动态形状。例子包括处理变长序列的自然语言处理模型、处理不同分辨率图像的对象检测模型,或处理不同批次大小请求的推理 (inference)服务器。与静态形状执行相比,这种动态性带来了显著的挑战,主要影响内存管理、内核选择以及整体性能的可预测性。
考虑一个形状为 的张量。在静态情况下, 是在编译时已知的常量。这使得编译器和运行时能够精确地预先规划内存分配,为这些确切维度选择或生成高度优化的内核,并高效地安排操作。然而,如果一个或多个维度是符号化的或直到运行时才可知(例如,),运行时必须采用特定的策略来处理这种不确定性。
一种方法是运行时系统动态执行形状推断。随着计算图中的操作执行,运行时可以传播具体的形状信息。例如,如果一个操作接受两个已知形状为 和 的输入张量并产生一个输出,运行时可以推断出输出形状将是 。这依赖于运行时能够访问图结构以及每个算子的形状转换函数。通常,编译器使用的中间表示(IR)(如 MLIR)会编码这些形状函数,或提供表示未知维度(通常表示为 ? 或 -1)的机制。虽然在许多情况下有效,动态形状推断可能无法解决所有歧义,尤其是在复杂的控制流情况下,或输出形状取决于输入值(而不仅仅是形状)的操作。
当操作执行前确切形状已知时,运行时可以选择一个针对这些特定维度优化的预编译内核,甚至可以触发即时(JIT)编译过程来动态生成一个专用内核。这对于计算密集型操作尤其重要,如卷积或矩阵乘法,因为其性能对维度高度敏感。
示意图,展示了运行时形状检查和内核选择/JIT编译的简化流程。
这种策略为实际计算保持高性能,但可能引入延迟。JIT编译需要时间,并且管理专用内核缓存会增加复杂性。形状多态性等技术旨在通过生成能够高效处理有限范围输入形状的内核来缓解此问题,从而减少过度专用化的需求。运行时必须权衡专用化的成本(编译时间、缓存管理)与使用特定维度代码所带来的性能提升。
一种常用的施加规整性的方法是填充。如果一个维度可以变化到某个最大大小(例如,序列长度最大为 512),运行时可以基于这个最大大小分配缓冲区,并填充较小的输入以使其适应。
填充:
分桶: 为减少填充的开销,可以根据输入的形状特征将它们分组到“桶”中。例如,长度为 1-64 的序列可能分到一起,65-128 的分到另一组,以此类推。然后填充仅应用于每个桶的范围内。
比较填充到固定最大值与使用离散桶的内存分配/处理大小。分桶减少了较小输入的浪费。
分桶要求运行时根据形状动态地对输入进行排序或批处理,这增加了调度复杂性,但通常比简单的填充能带来更好的资源利用。
更高级的运行时可能包含符号形状计算。它们不是立即要求具体的维度,而是操作表示形状的符号表达式(例如,N 表示批次大小)。操作在这些符号形状上执行,直到需要具体值进行内存分配或内核执行时。这可以延迟专用化或分配决策。通常,这与上限设定结合使用,即编译器或用户为动态维度提供可能的最大值。运行时随后可以使用这些上限进行初始内存预留(通常在一个竞技场分配器内),同时可能根据后续遇到的具体形状来优化实际使用的部分。
动态形状从根本上使内存管理复杂化。所有缓冲区地址和大小在离线确定的静态内存规划不再完全适用。运行时内存管理器必须处理那些大小仅在执行期间才已知的缓冲区请求。这通常需要动态内存分配器(如为高效处理不同大小而设计的竞技场分配器)以及减少碎片化的策略。未能高效管理内存可能导致内存不足错误,或由于分配开销或数据局部性差而导致显著的性能下降。
同样,调度也变得更为复杂。操作的执行时间会因输入形状而显著不同。这使得调度器更难预测成本、有效重叠计算和数据移动,或在异构设备间平衡负载。调度器可能需要根据观察到的特定形状的执行时间进行动态调整。
处理动态形状涉及灵活性和性能可预测性之间的权衡。
最佳策略很大程度上取决于具体的模型架构、形状变化的程度和性质、目标硬件能力(例如 JIT 编译速度、内存容量)以及应用的延迟要求。分析工具对于理解动态形状如何影响内存使用、内核执行时间以及给定运行时系统中潜在的 JIT 开销变得必不可少。有效的处理通常涉及结合多种技术,例如使用上限进行初始分配、执行运行时形状推断,并采用结合分桶的形状专用内核来高效管理可变性。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•