[논문 리뷰] FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours

Shenggan Cheng, Zhao, Xuanlei외 8명|arXiv (Cornell University)|2022. 03. 02.

Protein Structure and DynamicsBiochemistry, Genetics and Molecular Biology인용 수 17

한 줄 요약

FastFold는 동적 축 병렬 처리, 이중 비동기 연산, 오토체인징을 도입하여 메모리 사용량을 80% 이상 감소시키고 GPU 스케일링 효율을 햖스르며, AlphaFold 학습을 11일에서 67시간으로 단축하고 장거리 서열 추론 속도를 7.5–9.5배 빠르게 한다. 이 시스템은 512개의 A100 GPU에서 6.02 PetaFLOP/s의 처리량을 달성하며, 병렬 효율은 90.1%이다.

ABSTRACT

Protein structure prediction helps to understand gene translation and protein function, which is of growing interest and importance in structural biology. The AlphaFold model, which used transformer architecture to achieve atomic-level accuracy in protein structure prediction, was a significant breakthrough. However, training and inference of the AlphaFold model are challenging due to its high computation and memory cost. In this work, we present FastFold, an efficient implementation of AlphaFold for both training and inference. We propose Dynamic Axial Parallelism and Duality Async Operations to improve the scaling efficiency of model parallelism. Besides, AutoChunk is proposed to reduce memory cost by over 80% during inference by automatically determining the chunk strategy. Experimental results show that FastFold reduces overall training time from 11 days to 67 hours and achieves 7.5X - 9.5X speedup for long-sequence inference. Furthermore, we scale FastFold to 512 GPUs and achieve an aggregate throughput of 6.02 PetaFLOP/s with 90.1% parallel efficiency.

연구 동기 및 목표

• 대규모에서 AlphaFold 모델을 학습하고 추론할 때 발생하는 높은 계산 및 메모리 비용을 해결한다.
• 다중 GPU 학습에서 모델 병렬 처리의 확장성과 통신 오버헤드의 한계를 극복한다.
• 초장거리 단백질을 처리할 수 있도록 장거리 서열 추론 중 메모리 소비를 줄인다.
• 대규모 GPU 클러스터(예: 512개의 A100 GPU)에서 높은 처리량과 병렬 효율을 달성한다.
• 구조 생물학 및 그 외 분야의 다른 어텐션 기반 모델에 적용 가능한 일반적인 최적화 프레임워크를 개발한다.

제안 방법

• 통신 오버헤드를 줄인 텐서 분할 전략을 통해 GPU 간 최적화된 텐서 분할을 수행하는 동적 축 병렬 처리를 제안한다.
• 계산과 통신의 겹침을 숨기기 위해 계산 지연을 감추는 PyTorch 확장인 이중 비동기 연산을 도입한다.
• 추론 중 메모리 사용량을 최소화하기 위해 최적의 서열 청크를 자동으로 결정하는 오토체인징 기법을 개발한다.
• Evoformer 블록에 저수준 커널 최적화를 적용하여 계산 속도를 향상시키고 메모리 프로파일을 줄인다.
• 하이브리드 병렬 처리와 통신 최적화된 학습 파이프라인을 사용하여 512개의 A100 GPU에서 학습을 확장한다.
• 추론 중 속도와 메모리 사용량의 균형을 유지하기 위해 메모리 인식 기반 청크 조정을 동적으로 적용한다.

실험 결과

연구 질문

• RQ1 대규모 단백질 구조 예측에서 모델 병렬 처리를 최적화하여 통신 오버헤드를 줄이고 확장성 효율을 향상시킬 수 있는가?
• RQ2 어텐션 기반 모델인 AlphaFold에서 계산과 통신의 겹침을 효과적으로 달성할 수 있는가?
• RQ3 자동화된 청크 전략이 장거리 서열 추론 중 메모리 사용량을 줄이는 데 전문가가 수작업으로 설정한 구성보다 뛰어난 성능을 낼 수 있는가?
• RQ4 저수준 커널 및 메모리 최적화가 어텐션 기반 모델의 학습 및 추론 비용을 어느 정도 줄일 수 있는가?
• RQ5 제안된 최적화 기법이 AlphaFold를 초월한 다른 대규모 어텐션 모델에 일반화될 수 있는가?

주요 결과

• FastFold는 512개의 A100 GPU에서 AlphaFold 학습 시간을 11일에서 67시간으로 단축하여 10.5배의 속도 향상을 달성한다.
• 512개 GPU에서 6.02 PetaFLOP/s의 집합 처리량을 기록하며, 병렬 효율은 90.1%이다.
• AutoChunk는 추론 중 메모리 사용량을 80% 이상 감소시켜 전문가가 수작업으로 설계한 청크 전략보다 메모리 효율에서 30.6%–34.4% 높은 성능을 달성한다.
• FastFold는 기준 구현 대비 장거리 서열 추론 속도를 7.5–9.5배 빠르게 한다.
• 오토체인징은 전문가가 수작업으로 설정한 청크(64)를 사용한 OpenFold 대비 추론 속도를 12% 향상시키며, 청크 없이 추론하는 것과 비교해 메모리 사용량은 4%만 증가한다.
• 이 최적화 기법은 RoseTTAFold, MSA Transformer, 영상 트랜스포머 등 다른 어텐션 기반 모델에도 일반화되어 적용 가능하다.