将张量与序列并行融合,降低大模型训练内存占用。
Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

- 将权重和序列分片合并到同一设备轴上,减少内存开销。
- 在长序列场景下,相比传统方案节省40%以上显存占用。
- 适合资源受限环境下的大模型训练,可与流水线等并行策略共用。
我们提出张量与序列并行(TSP),一种将张量并行(TP)和序列并行(SP)融合于单一设备轴的并行执行策略。传统多维并行布局中,TP对模型参数进行分片,SP对输入序列分片,分别降低单设备的参数或激活内存。而传统方式需分配独立的网格维度。TSP则让每个计算节点同时持有权重分片和序列分片,沿同一设备轴同时减少参数与激活内存。我们实现两种运行时调度:对于注意力模块,各rank广播权重分片,并通过逐序列交换键值重建上下文;对于门控MLP,权重分片以环形流动,局部累积部分输出。通过在同一设备上同时分片权重与激活,TSP以增加通信量为代价换取更低内存开销。我们提供了通信与内存的理论分析,描述了TSP注意力与门控MLP模块的实现,并与TP、SP及TP+SP进行了基准测试。结果表明,TSP是长序列、内存受限场景下的一种硬件感知替代方案,亦可与流水线并行、专家并行等现有策略协同用于密集模型与混合专家模型。
原文摘要 · Abstract (English)
We present tensor and sequence parallelism (TSP), a parallel execution strategy that folds tensor parallelism and sequence parallelism onto a single device axis. In conventional multi-dimensional parallelism layouts, tensor parallelism (TP) shards model weights while sequence parallelism (SP) shards tokens, reducing per-device parameter or activation memory, respectively. Traditionally, each scheme is assigned its own mesh dimension. TSP instead assigns each rank both a weight shard and a sequence shard, reducing both parameter and activation memory along the same device axis. We implement this design with two runtime schedules. For attention, ranks iterate over broadcast parameter shards and reconstruct context through a sequence-wise key/value exchange. For gated MLPs, weight shards circulate in a ring while partial outputs accumulate locally. By sharding both weights and activations across the same devices, TSP trades additional communication volume for reduced memory overhead. We provide a theoretical communication and memory analysis, describe our implementation of TSP attention and gated MLP blocks, and benchmark TSP against TP, SP, and TP+SP. These results position TSP as a hardware-aware alternative for long-context and memory-constrained model training, and as a viable axis of parallelism in concert with existing parallelism schemes such as pipeline and expert parallelism for dense and mixture-of-expert models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。