TorchAcc：基于 TorchXLA 的分布式训练框架
内容提要
阿里云 TorchAcc是一个基于PyTorch/XLA的大模型分布式训练框架，提供多样化的并行策略和显存优化功能，通过图形优化和通信优化提高了分布式训练的效率和性能。该框架在多个模型的分布式训练场景中表现出显著的性能优势。
延伸解读
分布式训练复杂性的根源
文章指出，大模型规模增长远超单GPU算力和显存提升速度，迫使采用分布式训练。数据并行逻辑简单，仅需AllReduce同步梯度；而模型并行需分割模型、选择通信原语并调度算子以最大化计算通信重叠，开发难度显著更高。TorchAcc旨在通过自动探索并行策略降低这一复杂性，让开发者聚焦模型设计。
IR Graph：优化策略的统一基础
TorchAcc将PyTorch或TensorFlow前端模型转化为统一的中间表示层（IR Graph），并基于此实施计算、存储、通信和分布式策略优化。不同并行策略在不同抽象层级操作：数据并行、流水并行和FSDP在FX Graph层实现，而张量并行和序列并行则在更底层的HLO Graph层通过mark sharding接口和SPMD优化完成。这种分层设计平衡了优化效果与实现成本。
显存优化：ROAM方法的分治策略
显存优化是NP-hard问题，传统启发式或约束求解方法难以兼顾时效与高效。TorchAcc提出ROAM方法，采用分治策略：引入显存感知的权重更新调度器，根据梯度产生时机和优化器类型选择更新点；将大图分割为内存无关子图，子图内用线性规划求解最优执行顺序。实验显示，ROAM相比原生PyTorch、启发式算法和Facebook的整数线性规划方法分别节省约16%、13%和27%显存，且求解时间显著缩短。
性能提升的实际验证
文章通过Llama 2-7B模型测试展示了优化效果：计算优化中，XLA的kernel fusion合并访存密集型算子，叠加FlashAttention后性能进一步提升；通信优化中，合并collective算子、移至独立CUDA Stream异步执行，并利用Latency Hiding Scheduler，在128张GPU上加速比从88提升至116。这些结果验证了TorchAcc在真实训练场景中的有效性。
Q&A
TorchAcc 是什么？
TorchAcc 是一个基于 PyTorch/XLA 的大模型分布式训练框架，旨在提高分布式训练的效率和性能。
TorchAcc 如何优化显存使用？
TorchAcc 通过显存智能分配器优化显存资源，精细调度和地址分配，降低显存开销。
TorchAcc 支持哪些并行策略？
TorchAcc 支持数据并行、模型并行、FSDP（Fully Sharded Data Parallel）等多种并行策略。
TorchAcc 如何提升训练效率？
TorchAcc 通过自动探寻并整合并行策略、优化计算密集度和减少访存开销来提升训练效率。
TorchAcc 在分布式训练中有哪些性能优势？
TorchAcc 在多个模型的分布式训练中表现出显著的性能优势，部分模型训练速度提升高达 3 倍。
TorchAcc 如何处理复杂的通信模式？
TorchAcc 通过合并算子、异步执行和优化调度来有效处理复杂的通信模式，提升训练效率。
