折叠张量和序列并行性，实现高效内存 Transformer 训练和推理

我们提出了张量和序列并行（TSP），这是一种将张量并行和序列并行折叠到单个设备轴上的并行执行策略。在传统的多维并行布局中，张量并行（TP）对模型权重进行分片，而序列并行（SP）对词元进行分片，从而分别减少每个设备的参数或激活内存。传统上，每个方案都分配有自己的网格尺寸。相反，TSP 为每个排名分配一个权重分片和一个序列分片，从而减少沿同一设备轴的参数和激活内存。我们用两个运行时计划来实现这个设计。为了引起注意，对广播参数分片进行排名迭代，并通过按顺序的键/值交换重建上下文。对于门控 MLP，权重分片在环中循环，而部分输出在本地累积。通过在同一设备上分片权重和激活，TSP 以额外的通信量换取减少的内存开销。我们提供了理论通信和内存分析，描述了我们对 TSP 注意力和门控 MLP 块的实现，并针对 TP、SP 和 TP+SP 对 TSP 进行了基准测试。这些结果将 TSP 定位为长上下文和内存受限的 模型训练 的硬件感知替代方案，并作为与现有并行方案（例如密集和专家混合模型的管道和专家并行）相结合的可行并行轴。