#每日一篇学术文章# 【Memory-efficient tensor parallelism for long-sequence Transformer training】（面向长序列Transformer训练的内存高效张量并行方法）
http://t.cn/AX7ZnCJB
近年来, 基于 Transformer 架构的大语言模型 (LLM) 凭借卓越性能引发广泛关注。 工业级 LLM 需处理长序列输入以提供优质服务。 然而, 内存消耗随序列长度呈平方级增长, 制约长序列训练的扩展能力。 现有并行方法在执行过程中产生冗余张量, 存在内存优化空间; 同时, 张量并行 (TP) 无法实现计算与通信的有效重叠。 针对上述问题, 本文提出一种通用并行方法——内存高效张量并行 (METP), 专为 Transformer 训练核心计算单元 (即两个连续矩阵乘法及其间可能存在的函数运算 O=f(AB)C 设计)。 METP 将计算 O 的子任务分配到多设备, 采用点对点通信 (send/recv) 替代集合通信交换子矩阵完成计算, 避免生成冗余张量。 通过双缓冲技术实现计算与通信的深度重叠, 并提出完全重叠的理论条件以指导长序列 Transformer 训练。 理论分析表明: 当并行度为 p 时, METP 在未使用 FlashAttention计 算注意力时的内存开销为 O(1/p3); 在使用 FlashAttention 计算多头自注意力时, 相比 TP 至少可节省 41.7% 内存。 实验证明, 基于 8 块 A100 GPU 的配置, METP 可使序列长度较其他方法提升 2.38–2.99 倍。

来源：Frontiers of Information Technology & Electronic Engineering