变量大小的输入在深度学习中很常见，但密集批处理需要为所有样本分配统一的形状，导致大量填充计算。填充成本会在每个可变轴上呈指数增长：显式配对状态会分配 $BN_{\max}^2$ 个位置，而实际只需要 $\sum_i N_i^2$。
打包（packing）可以去除这些浪费，但在打包后仍然难以组合操作，因为平坦缓冲区不再暴露逻辑轴和样本边界。我们提出 DanLing NestedTensor，这是一种基于 PyTorch 的张量抽象，把多层次不规则结构作为张量本身的属性。
打包后的数值携带张量支持的分区信息和逻辑维度顺序，因而广播会产生 ragged 轴，特征变换会保留这些轴，归约操作会消耗它们。相同的表示能够在 autograd、即时执行和编译执行之间无缝传递。
在 A100 GPU 上，对四种 BERT 规模的实验显示，几何平均加速比在即时模式下为 2.74×，在编译模式下为 3.39×；对四种 FCN 主干网络的即时加速为 1.97×。在四块 Pairformer 风格的工作负载中，使用原生 PyTorch 内核的即时执行比填充基准快 2.40‑4.32×，最高内存占用从 38.08 GiB 降至 5.41 GiB。
通过统一的张量接口，模型代码只需使用支持的算子即可组合高效的可变大小计算，无需在调用点手动管理偏移。代码将在论文发表后公开。
点评