[박찬 기자]

<이미지를 클릭하시면 크게 보실 수 있습니다>

엔비디아는 딥마인드가 개발한 단백질 생성 인공지능(AI) 모델 '알파폴드(AlphaFold)'의 훈련 시간을 단축하는 기술을 선보였다. 이를 통해 당초 7일 걸리던 모델의 훈련을 10시간 만에 마칠 수 있게 됐다.

마크테크포스트는 22일(현지시간) 엔비디아 연구진이 딥마인드의 단백질 생성 AI 모델인 알파폴드의 훈련 시간을 단축하는 '스케일폴드(ScaleFold)' 기술에 관한 논문을 온라인 아카이브에 게재했다고 전했다.

이에 따르면 스케일폴드는 벤치마크 기준보다 6배 이상 빠른 7.51분 만에 알파폴드에서 파생한 '오픈폴드(OpenFold)' 모델의 부분 훈련 작업을 완료, 궁극적으로 알파폴드의 7일 걸리던 초기 훈련 시간을 10시간으로 단축한다.

연구진은 알파폴드의 훈련 속도를 지연시키는 비효율적인 통신 방식과 로컬 CPU 오버헤드, 병렬화 불가능한 워크로드, 최적이 아닌 커널 확장성 등 컴퓨팅 리소스 저활용 요소를 발견했다.

이런 문제를 극복하기 위해 새로운 비차단 데이터 파이프라인을 도입하고 쿠다(CUDA) 그래프를 활용해 훈련을 추적하고 오버헤드를 최소화하는 등 세분화된 최적화를 구현, 통신 효율성을 크게 향상했다고 설명했다.

<이미지를 클릭하시면 크게 보실 수 있습니다>

또 알파폴드 훈련 내에서 필수 계산 패턴을 식별하고 각 패턴에 대한 전용 트리톤(Triton) 커널을 개발, 모델 전반에 걸쳐 단편화된 계산을 통합하고 다양한 작업 부하 크기 및 하드웨어 아키텍처에 맞게 커널 구성을 세심하게 조정했다. 트리톤은 신경망의 계산을 빠르고 효율적으로 실행하도록 GPU 하드웨어의 성능을 최대로 유지하도록 하는 프로그래밍 기법이다. 이런 최적화의 융합을 통해 스케일폴드가 만들어졌다.


ML퍼프 HPC v3.0 벤치마크에서 스케일폴드는 2080개의 'H100' GPU에서 엔비디아의 오픈폴드 모델의 부분 훈련 작업을 단 7.51분에 완료했는데, 이는 벤치마크 기준보다 6배 빠른 것이다. 특히 알파폴드 모델에 대한 총 사전 훈련 시간을 7일에서 단 10시간으로 단축하며 새로운 기록을 세웠다.

결론적으로 이 연구는 비효율적인 통신과 오버헤드 중심 계산을 완화하기 위해 맞춤화된 체계적인 접근 방식인 스케일폴드를 도입, 알파폴드 훈련의 확장성 문제를 해결했다는 평이다.

연구진은 "이 연구가 복잡한 단백질 접힘 문제를 해결하기 위해 딥러닝 기반 계산 방법을 확장하기 위한 효과적인 프레임워크를 제공, 고성능 컴퓨팅과 생물정보학 연구에 도움이 될 것"이라고 말했다.

박찬 기자 cpark@aitimes.com

<저작권자 Copyright ⓒ AI타임스 무단전재 및 재배포 금지>