论文

DanLing NestedTensor:可组合的多轴不等长张量

DanLing NestedTensor: Composable Multi-Ragged Tensors for Deep Learning

AI 基础设施模型服务框架

摘要

变长输入在深度学习中很常见,但稠密批处理分配统一的外包尺寸,并在填充位置上消耗计算。多个变长轴使成本相乘:显式成对状态分配BN_max²个位置,而实际只需Σ_i N_i²。打包可消除浪费,但组合打包运算仍需要扁平缓冲区不再显式提供的逻辑轴和样本边界。DanLing NestedTensor是一种PyTorch张量抽象,使多个不等长轴成为张量自身的属性。打包值携带由张量表示的分区与逻辑维度顺序,因此广播可以创建不等长轴、特征变换保留这些轴、归约消去这些轴。同一表征贯穿自动微分、即时执行和编译执行。在A100上,相对同执行模式的填充实现,四种BERT规模的几何平均加速在即时模式为2.74倍、编译模式为3.39倍;四种FCN骨干的即时模式加速为1.97倍。四块Pairformer式负载在即时执行的多种平方长度条件下,比使用原生PyTorch内核的填充参考快2.40—4.32倍,高变异批次峰值分配从38.08降至5.41 GiB。由受支持算子构建的模型可借助该接口组合高效变长计算,无需在调用处管理偏移量。代码计划在发表时公开。