SandAI-org/MagiAttention

해결하는 문제

MagiAttention은 분산 학습 환경에서 초장문맥 및 비균일한 마스크 패턴에 대한 어텐션 메커니즘의 확장성 문제를 해결합니다. 일반적으로 컨텍스트 병렬(CP) 학습 설정에서 발생하는 계산 및 통신 오버헤드를 줄이기 위해 선형 확장성과 높은 성능을 제공하는 것을 목표로 합니다.

작동 방식

다음과 같은 주요 혁신을 통해 분산 어텐션 메커니즘을 구현합니다:

• Flex-Flash-Attention (FFA): 일반화된 어텐션 마스크 정의(AttnSlice)와 다양한 마스크 유형을 컴팩트하게 표현하고 효율적인 분산 파티셔닝을 가능하게 하는 맞춤형 커널.
• 계산 부하 균형: 세밀한 청크 단위 샤딩 전략과 디스패치 솔버를 사용하여 CP 랭크 간 작업 부하를 균일하게 분배합니다.
• 제로 중복 통신: 전통적인 릴링형 P2P 통신을 대체하여 전방 및 역방향 전파 중 중복 통신량을 제거하는 새로운 GroupCast 및 GroupReduce 원시 기능을 도입합니다.
• 적응형 다단계 오버랩: 계산과 통신을 스케줄링하여 지연 시간을 숨기고 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와의 간편한 통합을 지원합니다.
• 고급 마스킹: 다양한 및 겹치는 어텐션 마스크 패턴을 네이티브로 지원합니다.
• 어텐션 싱크: 분산 학습 가능한 어텐션 싱크 메커니즘 확장을 포함합니다.