用于大规模并行序列生成的结构化循环混合器

在过去的二十年中，语言建模经历了从使用在训练和推理期间顺序处理词元的主要循环架构到在训练期间并行处理序列元素的非循环模型的转变，这导致了更高的训练效率和稳定性，但代价是较低的推理吞吐量。在这里，我们介绍了结构化循环混合器，这种架构允许在训练时的序列并行表示和推理时的循环表示之间进行代数转换，特别是不需要专门的内核或特定于设备的内存管理。我们通过实验证明，与其他线性复杂度模型相比，这种双重表示可以实现更高的训练效率、更高的输入信息容量以及更大的推理吞吐量和并发性。我们假设循环模型不太适合扩展序列长度缩放，以处理典型的语言中信息丰富的输入，但由于每个样本的内存恒定，因此非常适合样本（批量）维度的缩放。我们提供 SRM 的 Mojo/MAX 推理实现，与在 vLLM 上推理的类似功能强大的 Transformer 相比，吞吐量提高了 12 倍，并发性提高了 170 倍，增强了 Pytorch 实现的特性，从而使计算常数 GSM8k Pass@k 提高了 30%。最后，我们证明 SRM 是有效的强化学习训练候选者。