ringattention

收录理由

Ring Attention 在论文里看起来简单，实际实现却有不少细节要处理。这个仓库正是《Ring Attention with Blockwise Transformers》和《Blockwise Parallel Transformer》两篇论文所对应的 JAX 实现，把注意力计算和前馈网络都做了分块处理，并让通信与计算重叠进行，因此可训练的上下文长度会随加速器数量线性扩展，而不是被单块芯片的显存卡死。项目通过 pip 安装后，可以直接导入 ringattention 和 blockwise_feedforward 两个函数，仓库里还给出了用 shard_map 做多设备切分的示例，方便接入自己的模型。不过要说明的是，这属于研究代码，使用前需要针对自己的训练流程做适配，不能当作开箱即用的库。如果你正在 TPU 或多 GPU 的 JAX 环境里探索超长上下文模型，这里是最直接的起点。