🧬 biology

MegaFold: Efficient Training of Next-Generation 3D Attention Protein Models on Cross-Platform GPUs

MegaFold는 메모리 효율적인 커널, 최적화된 샤딩, 융합 연산자(fused operators), 그리고 결정론적 인지 파이프라인(determinism-aware pipeline)을 결합하여 NVIDIA 및 AMD GPU 모두에서 훨씬 더 긴 서열 길이와 더 빠른 학습 시간을 달ace하기 위해 차세대 3D 어텐션 단백질 모델의 메모리 및 계산 병목 현상을 극복하는 새로운 크로스 플랫폼 시스템입니다.

원저자: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

게시일 2026-06-16

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Hoa La, Ahan Gupta, Alex Morehead, Jianlin Cheng, Minjia Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ⚕️ 이것은 동료 심사를 거치지 않은 프리프린트의 AI 생성 설명입니다. 의학적 조언이 아닙니다. 이 내용을 바탕으로 건강 관련 결정을 내리지 마세요. 전체 면책 조항 읽기

당신이 거대하고 정교한 단백질 3D 퍼즐을 만들려고 노력하고 있다고 상상해 보세요. 오랫동안 과학자들은 잘 작동하는 특정 규칙 세트(예전의 "AlphaFold2"와 같은)를 사용해 왔습니다. 하지만 최신 버전인 AlphaFold3는, 평면적인 2D 퍼즐에서 모든 조각이 3D 공간 내에서 서로 상호작용하는 복잡한 다층 구조의 3D 퍼즐로 업그레이드된 것과 같습니다.

이 새로운 모델은 과학적으로는 놀랍지만, 한 가지 큰 문제가 있습니다. 바로 컴퓨터 메모리를 블랙홀처럼 집어삼키며 믿기지 않을 정도로 느리게 돌아간다는 점입니다.

이 논문은 이러한 차세대 3D 단백질 모델을 현대적인 컴퓨터에서 충돌 없이, 혹은 너무 오래 걸리지 않게 학습시키는 것을 가능하게 하는 새로운 "시스템"(특화된 도구 및 규칙 세트)인 MegaFold를 소개합니다.

다음은 MegaFold가 네 가지 주요 문제를 어떻게 해결하는지에 대한 설명입니다. 쉬운 비유를 통해 설명하겠습니다.

문제점: 왜 기존 방식은 실패하는가

새로운 AlphaFold3 모델은 단백질을 바라보는 방식이 데이터의 "입방체(cubic)" 폭발을 일으킵니다.

비유: 당신이 파티를 준비한다고 상상해 보세요. 예전 모델에서는 누가 누구와 친구인지(단순한 명단)만 알면 되었습니다. 하지만 새로운 모델에서는 모든 손님이 동시에 다른 모든 손님과 어떻게 상호작용하는지에 대한 3D 지도를 추적해야 합니다. 손님이 100명이라면, 예전 방식은 10,000개의 기록이 필요하지만, 새로운 방식은 1,000,000개의 기록이 필요합니다. 만약 이 모든 기록을 단 하나의 메모장(컴퓨터 메모리)에 다 적으려고 한다면, 그 메모장은 터져버릴 것입니다.

해결책: MegaFold의 네 가지 도구

MegaFold는 다음과 같은 네 가지 특정 업그레이드를 도입하여 이를 해결합니다.

1. "스마트 메모장" (EvoFlash-3D)

문제: 컴퓨터가 전체 3D 상호작용 지도를 한꺼번에 적으려고 시도하다가 순식간에 메모리를 가득 채워버립니다.
MegaFold의 해결책: 전체 지도를 메모장에 한 번에 적는 대신, MegaFold는 "연습장"(빠른 임시 메모리)을 사용하여 상호작용을 작은 덩어리로 하나씩 계산한 다음, 그 기록을 즉시 버립니다.
결과: 전체 거대한 지도를 한꺼번에 메모리에 담아둘 필요가 없습니다. 덕분에 컴퓨터는 메모리 부족 현상 없이 훨씬 더 긴 단백질 사슬을 처리할 수 있습니다.

2. "팀 릴레이" (EvoSP-3D)

문제: 여러 대의 컴퓨터(GPU)를 사용하여 협업할 때, 보통 각 컴퓨터에 서로 다른 손님 명단을 주는 방식으로 작업을 나눕니다. 하지만 AlphaFold3의 데이터는 2D 상호작용 그리드이기 때문에, 단순히 명단을 나누는 방식은 제대로 작동하지 않습니다. 컴퓨터들이 서로 누가 누구와 대화하고 있는지 혼란을 겪기 때문입니다.
MegaFold의 해결책: MegaFold는 컴퓨터들이 데이터를 주고받는 새로운 방식을 발명했습니다. 단순히 명단을 나누는 대신, 상호작용의 그리드를 나눕니다. 그들은 특정하고 효율적인 패턴(마치 릴레이 경주처럼)으로 퍼즐 조각을 주고받으며, 모든 컴퓨터가 다른 컴퓨터를 기다리며 시간을 낭비하지 않고도 자신이 책임져야 할 3D 지도의 어느 부분을 담당하고 있는지 정확히 알 수 있게 합니다.

3. "조립 라인" (EvoFusion)

문제: 컴퓨터가 매우 작고 비효율적인 단계로 작업을 수행하고 있었습니다. 숫자를 하나 계산하고, 멈추고, 저장하고, 다시 도구를 불러오고, 다시 계산하고, 또 멈추고, 저장하는 식입니다. 이는 마치 요리사가 양파를 썰고 나서, 칼을 가지러 냉장고로 갔다가, 다시 돌아와서 양파를 써는 것과 같습니다.
MegaFold의 해결책: MegaFold는 이러한 작은 단계들을 하나의 크고 매끄러운 동작으로 결합합니다. "썰기", "섞기", "요리하기"를 하나의 연속적인 동작으로 융합합니다.
결과: 컴퓨터는 작업을 시작하고 멈추는 데 시간을 낭비하지 않습니다. 중단 없이 "조립 라인"을 계속 움직이게 함으로써 훨씬 더 빠르게 작동합니다.

4. "미리 준비된 재료" (EvoPipe)

문제: 컴퓨터가 단백질을 만들기 시작하기도 전에, 인간(또는 느린 CPU)이 적절한 데이터(진화적 역사 등)를 모으기 위해 많은 번거로운 조사를 수행해야 합니다. 이 과정은 시간이 오래 걸리며, 강력한 컴퓨터가 데이터를 기다리느라 유휴 상태로 머물게 만듭니다.
MegaFold의 해결책: MegaFold는 동일한 단백질에 대해서는 이러한 연구 작업의 상당 부분이 항상 동일하다는 점을 깨달았습니다. 이 어려운 작업을 미리 한 번 수행하여 "식료품 저장실"(캐시)에 저장해 둡니다. 컴퓨터가 데이터가 필요할 때, 준비된 재료를 즉시 가져다 쓰기만 하면 됩니다.
결 결과: 강력한 컴퓨터가 기다리는 일이 없습니다. 컴퓨터는 100%의 시간 동안 계속 바쁘게 작동합니다.

결과

MegaFold는 두 종류의 강력한 컴퓨터 칩(NVIDIA 및 AMD)에서 테스트되었습니다. 결과는 다음과 같습니다.

더 긴 단백질: MegaFold를 통해 이전보다 3.36배 더 긴 단백질 서열을 학습할 수 있었습니다.
더 빠른 속도: NVIDIA 칩에서는 학습 과정이 1.73배 더 빨라졌고, AMD 칩에서는 1.62배 더 빨라졌습니다.
적은 메모리 사용량: 메모리를 현저히 적게 사용하여, 이러한 모델에서 흔히 발생하는 "메모리 부족(out of memory)" 충돌을 방und했습니다.

요약

AlphaFold3를 컴퓨터 시스템(자동차 섀시)에 비해 너무 무거운 초강력 엔진이라고 생각해보세요. MegaFold는 단순히 더 큰 자동차를 만든 것이 아니라, 그 강력한 엔진이 실제로 효율적으로 달릴 수 있도록 엔진, 변속기, 그리고 연료 시스템을 재설계했습니다. 이를 통해 과학자들은 더 다양한 컴퓨터에서 차세대 모델을 학습시킬 수 있게 되었으며, 단백질 구조 결정 과학을 더 빠르고 접근 가능하게 만들었습니다.

기술 요약: MegaFold

문제 정의

생물 분자 모델링의 최근 발전, 특히 AlphaFold3(AF3)는 표준 트랜스포머 아키텍처에서 2D 쌍별(pairwise) 표현에 대한 3D 어텐션을 활용하는 모델로의 패러다임 전환을 도입했습니다. AF3는 단백질-단백질, 단백질-DNA/RNA, 그리고 단백질-리간드 상호작용을 정확하게 예측할 수 있게 해주지만, 이러한 아키텍처의 변화는 기존의 LLM 최적화 기법으로는 해결할 수 없는 심각한 시스템 수준의 병목 현상을 초د합니다.

본 논문은 3D 어텐션의 세제곱 스케일링( $O(N^3)$ )에 기인하는 네 가지 근본적인 과제를 식별합니다:

세제곱 활성화 메모리(Cubic Activation Memory): 트랜스포머의 어텐션이 제곱 스케일링( $O(N^2)$ )을 따르는 것과 달리, AF3의 2D 쌍별 표현에 대한 3D 어텐션은 세제곱으로 스케일링되는( $O(N^3)$ ) 텐서를 생성합니다. 이는 활성화 메모리가 GPU 사용량의 대부분(최대 96%)을 차지하게 만들어, 작은 배치 크기에서도 중간 정도의 시퀀스 길이(예: 590개 잔기 초과)에서 메모리 부족(OOM) 오류를 유발합니다.
시퀀스 병렬성(Sequence Parallelism, SP)과의 불일치: 트랜스포머를 위한 표준 SP 전략은 1D 시퀀스 샤딩에 의존합니다. 그러나 AF3의 2D 쌍별 표현과 레이어 전반에 걸친 교차 축 리덕션(reduction) 패턴은 과도한 리샤딩(resharding)이나 동기화 없이 기존 SP가 지원할 수 없는 복잡한 다중 축 통신 패턴을 생성합니다.
커널 파편화(Kernel Fragmentation): 3D 어텐션, 트랜지션 레이어, 그리고 혼합 차원 텐서에 대한 프로젝션 연산이 인터리빙(interleaving)되면서 반복당 수천 개의 미세한 커널이 생성됩니다. 이러한 파편화는 효과적인 커널 퓨전을 방해하여 워프 점유율(warp occupancy) 저하, 런칭 오버헤드 증가, 메모리 대역폭 활용 비효율성을 초래합니다.
호스트 측 전처리 스톨(Host-Side Preprocessing Stalls): 3D 어텐션에 필요한 진화적 컨텍스트(예: 다중 서열 정렬, MSA)를 구축하는 과정은 의존성이 높은 무거운 CPU 전처리를 포함합니다. 이는 기존의 프리페칭(prefetching) 파이프라인으로는 숨길 수 없는 상당한 GPU 유휴 시간을 생성합니다.

방법론: MegaFold

MegaFold는 메모리 효율적인 커널, 통신 전략, 연산자 퓨전 및 호스트-디바이스 파이프라이닝의 공동 설계(co-design)를 통해 이러한 과제를 해결하도록 설계된 크로스 플랫폼 시스템입니다. 이는 오픈 소스 AF3 훈련 프레임워크를 기반으로 구현되었으며 NVIDIA와 AMD GPU를 모두 지원합니다.

1. EvoFlash-3D: 메모리 효율적인 3D 어텐션

세제곱 메모리 증가 문제를 해결하기 위해, MegaFold는 3D 어텐션을 위한 Triton 기반 커널인 EvoFlash-3D를 도입합니다.

메커니즘: 글로벌 메모리(HBM)에 전체 $O(N^3)$ 어텐션 맵을 실체화(materialize)하는 대신, 이 커널은 온라인 방식의 수치적으로 안정적인 소프트맥스 전략을 사용합니다. 순전파(forward)와 역전파(backward) 과정 모두에서 빠른 스크래치패드 메모리(공유 메모리) 내에서 어텐션 로짓을 점진적으로 계산합니다.
최적화: 이를 통해 수치적 정확성을 유지하면서 피크 활성화 메모리 복잡도를 $O(N^3)$ 에서 $O(N^2)$ 로 줄입니다. 이 커널은 표준 FlashAttention에는 없는 다차원 데이터 의존성을 유발하는 AF3 특유의 가산적 브로드캐스트 쌍 바이어스(additive broadcasted pair-bias) 항을 처리합니다.

2. EvoSP-3D: 통신 효율적인 샤딩

더 긴 시퀀스 학습을 가능하게 하기 위해, MegaFold는 2D 쌍별 표현에 맞춤화된 샤딩 전략인 EvoSP-3D를 제안합니다.

메커니즘: 하이브리드 데이터 병렬(DP) 및 시퀀스 병렬(SP) 접근 방식을 활용합니다. DP 그룹이 서로 다른 배치 입력을 처리하는 동안, SP 그룹은 2D 쌍별 텐서( $[N_{token}, N_{token}, h_{pair}]$ )를 분할합니다.
통신: AF3 레이어가 어텐션 축을 교차(예: 행 방향에서 열 방향으로)함에 따라, EvoSP-3D는 단계 간에 쌍별 텐서의 샤딩 레이아웃을 회전시키기 위해 all_to_all 컬렉티브를 사용합니다. 이를 통해 연속적인 어텐션 레이어가 전체 텐서를 실체화하지 않고도 올바른 2D 슬라이스 상에서 동작할 수 있게 합니다.
디퓨전 단계: 디퓨전 디코더에서는 증강된 원자 좌표가 통신 없이 SP 그룹에 분산되는 반면, 쌍별 표현은 디퓨전이 전체 컨텍스트 상에서 동작할 수 있도록 all_gather를 통해 수집됩니다.

3. EvoFusion: 퓨즈드 연산자 스택

커널 파편화를 완화하기 위해, MegaFold는 단순히 경계면을 퓨전하는 것을 넘어 핵심 연산자의 내부 실행 구조를 재구성하는 EvoFusion을 도입합니다.

범위: 레이어 정규화(Layer Normalization), 선형 프로젝션(Linear Projections), SwiGLU 활성화 함수를 대상으로 하며, 이들은 결합적으로 약 20%의 활성화 메모리와 커널 호출을 차지합니다.
메커니즘: 이러한 연산자들의 내부 루프를 재구조화함으로써, 정규화와 선형 프로젝션 사이에 전체 $N \times N$ 텐서를 HBM에 쓰는 중복된 실체화 과정을 제거합니다. 이는 공유 메모리 압박과 커널 런칭 오버헤드를 줄이는 동시에 효율적인 리덕션 및 GeMM 스케줄을 유지합니다.

4. EvoPipe: 결정론적 호스트-디바이스 파이프라인

CPU 전처리로 인한 GPU 유휴 시간을 제거하기 위해, MegaFold는 EvoPipe를 구현합니다.

메커니즘: 결정론적인 전처리 단계(예: RDKit 기반 원자 구성, MSA 특징 추출)를 확률적인 훈련 루프로부터 분리합니다. 이러한 결정론적 특징들은 오프라인에서 미리 계산되어 캐싱됩니다.
실행: 훈련 중 시스템은 캐싱된 특징을 검색하기 위한 경량 룩업(lookup)을 수행하며, 비결정론적인 증강(예: 랜덤 크로핑)만을 실시간으로 적용합니다. 이를 통해 가변적인 호스트 측 스톨을 제거하여 일관된 GPU 활용도를 보장합니다합니다.

주요 결과

시스템은 AlphaFold3-PyTorch 프레임워크를 사용하여 NVIDIA H200 (141GB) 및 AMD MI250 (64GB) GPU에서 평가되었습니다.

시퀀스 길이 확장성: MegaFold는 32개의 GPU를 사용할 때 베이스라인 PyTorch 구현 대비 최대 3.36배 더 긴 시퀀스 길이의 학습을 가능하게 합니다. 32개의 AMD GPU에서 베이스라인은 약 590개 잔기에서 실패하는 반면, MegaFold는 약 1450개 잔기까지 지원합니다.
실행 시간: MegaFold는 엔드 투 엔드 반복당 실행 시간을 NVIDIA에서는 최대 1.73배, AMD에서는 최대 1.62배 단축합니다.
메모리 효율성: 피크 GPU 메모리 사용량을 최대 1.23배(또는 23% 감소) 줄여, 더 긴 시퀀스에 메모리 용량을 재할당할 수 있게 합니다.
모델 FLOPs 활용도 (MFU): MegaFold는 베이스라인이 대규모 스케일에 도달하기 전에 OOM 한계에 부딪히는 것과 달리, 훨씬 높은 MFU(1.81B 파라미터에서 최대 44.9%)를 달s성합니다.
크로스 플랫폼 성능: 시스템은 NVIDIA 및 AMD 하드웨어 전반에서 일관된 이득을 보여주며, Triton 기반 커널 설계의 이식성을 입증했습니다.

의의 및 주장

본 논문은 MegaFold가 차세대 3D 어텐션 단백질 모델의 훈련을 통합적으로 지원하는 최초의 엔드 투 엔드 시스템이라고 주장합니다. 그 의의는 다음과 같습니다:

LLM 최적화의 불충분함: 표준 트랜스포머 최적화(FlashAttention, 표준 SP, 컴파일러 퓨전 등)는 AF3 스타일 모델의 고유한 세제곱 스케일링과 혼합 차원 의존성을 해결하는 데 실패합니다.
시스템 수준의 공동 설계 필수성: 이러한 과학적 모델의 효율적인 훈련을 위해서는 커널 실행, 병렬화, 연산자 퓨전 및 호스트-디바이스 파이프라이닝을 결합하여 다루는 전용 시스템 지원이 필요합니다.
확장성 장벽의 이동: 3D 어텐션의 특정 병목 현상을 해결함으로써, MegaFold는 OOM 장벽을 옮겨 놓았으며, 이를 통해 이전에는 불가능했던 더 큰 생물 분자 복합체의 훈련을 가능하게 하여 현대적 GPU 플랫폼에서 신흥 과학적 워크로드를 확장할 수 있는 길을 열었습니다.