用torch.compile优化大模型分布式训练，内存降28.5%，吞吐升68.7%。
SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile
- 基于torch.compile实现可编译的全分片数据并行，支持自动通信重叠。
- 在Llama 3 405B上实现28.54%内存减少和68.67%吞吐提升。
- 适合需要高效训练超大规模模型的研究者与工程团队。
大规模模型分布式训练消耗巨大算力，且需复杂工程整合多种训练技术。本文提出SimpleFSDP，一个原生PyTorch、基于torch.compile的全分片数据并行框架，具备简洁可维护的实现，支持完整计算-通信图追踪，并通过编译器后端优化提升性能。其创新点在于利用现有PyTorch原语（参数化、选择性激活检查点、DTensor）实现兼容torch.compile的集合通信，首次在TorchInductor后端引入中间表示节点分桶与重排序，实现有效计算-通信重叠。用户可通过自动或手动包裹模型组件，最小化通信暴露。在Llama 3模型（含405B超大规模版本）上使用TorchTitan进行的评估显示，相比最广泛使用的FSDP2 eager框架，内存最多降低28.54%，吞吐最高提升68.67%，且可与其他分布式训练技术组合使用。
原文摘要 · Abstract (English)
Distributed training of large models consumes enormous computation resources and requires substantial engineering efforts to compose various training techniques. This paper presents SimpleFSDP, a PyTorch-native compiler-based Fully Sharded Data Parallel (FSDP) framework, which has a simple implementation for maintenance and composability, allows full computation-communication graph tracing, and brings performance enhancement via compiler backend optimizations. SimpleFSDP's novelty lies in its unique $torch.compile$-friendly implementation of collective communications using existing PyTorch primitives, namely parametrizations, selective activation checkpointing, and DTensor. It also features the first-of-its-kind intermediate representation (IR) nodes bucketing and reordering in the TorchInductor backend for effective computation-communication overlapping. As a result, users can employ the aforementioned optimizations to automatically or manually wrap model components for minimal communication exposure. Extensive evaluations of SimpleFSDP on Llama 3 models (including the ultra-large 405B) using TorchTitan demonstrate up to 28.54% memory reduction and 68.67% throughput improvement compared to the most widely adopted FSDP2 eager framework, when composed with other distributed training techniques.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。