SandAI-org/MagiAttention

解決的問題

MagiAttention 解決了在分散式訓練中，超長上下文與異質遮罩模式下注意力機制的可擴展性挑戰。旨在提供線性可擴展性與高效率，降低通常與上下文平行（CP）訓練設定相關的計算與通訊開銷。

工作原理

透過以下關鍵創新實現分散式注意力機制：

• Flex-Flash-Attention (FFA)：一種通用的注意力遮罩公式化（AttnSlice）與客製化核心，可緊湊表達多樣遮罩類型並實現高效的分散式分割。
• 計算負載平衡：採用細粒度的區塊級分片策略與調度求解器，確保 CP 排名之間的工作負載均勻分配。
• 零冗餘通訊：以新型的 GroupCast 與 GroupReduce 原語取代傳統的環狀 P2P 通訊，消除前向與反向傳播中的冗餘通訊量。
• 適應性多階段重疊：調度計算與通訊，隱藏延遲並最大化 GPU 使用率。

適用對象

專為訓練超長上下文大模型的研究人員與開發者設計，例如自回歸式影片生成模型（如 Magi-1），以及需要與 Megatron-LM、PyTorch FSDP 和 HuggingFace Transformers 等框架整合的場景。

核心亮點

• 線性可擴展性：在 H100 與 B200 GPU 上的分散式訓練環境中實現近乎線性的可擴展性。
• 廣泛硬體支援：在 Hopper GPU 上性能媲美 Flash-Attention 3，並透過 Flash-Attention 4 早期支援 Blackwell GPU。
• 框架整合：可輕鬆整合至 Megatron-LM、PyTorch FSDP 與 HuggingFace Transformers。
• 高階遮罩支援：原生支援多樣且重疊的注意力遮罩模式。
• 注意力池：包含對分散式可學習注意力池機制的擴展。