엔비디아, 분자동역학 계산 최대 1.7배 빨라져

엔비디아(NVIDIA, $NVDA)가 원자 1만7280개로 구성한 물 시스템에서 B300 GPU의 BF16x9 기법으로 PET 분자동역학 계산을 FP32 방식보다 최대 1.7배 빠르게 처리했다. 이 수치는 행렬곱만이 아니라 분자동역학 한 단계 전체를 기준으로 했다.

엔비디아는 29일 기술 블로그에서 기계학습 원자간 포텐셜인 PET(Point Edge Transformer)에 텐서 코어 기반 행렬곱 기법 두 가지를 적용한 결과를 공개했다. BF16x9는 B300(블랙웰) GPU에서, TF32x3는 A100 GPU에서 시험했다. 성능 수치는 특정 모델과 장비로 측정한 엔비디아 자체 평가 결과다.

기계학습 원자간 포텐셜은 원자 구조를 바탕으로 에너지와 원자 간 힘을 추정해 분자와 재료의 움직임을 계산하는 데 쓰인다. 이번 시험은 계산량이 큰 FP32 행렬곱을 텐서 코어로 처리할 때 성능과 결과가 어떻게 달라지는지를 살폈다.

PET는 스위스 로잔연방공과대(EPFL) COSMO 연구실이 개발했다. 엔비디아는 매개변수 2600만개인 PET-MAD 소형 모델을 원자 192개와 1만7280개 물 시스템에서 평가했다. 매개변수 2억5500만개인 대형 모델은 원자 192개 시스템에서 시험했다.

BF16x9는 FP32 값을 BF16 성분 세 개로 나눠 아홉 개의 부분 곱을 계산한 뒤 결과를 FP32로 합산한다. 이 방식은 큰 행렬곱을 텐서 코어에서 처리하도록 설계됐다.

원자 1만7280개 시스템에서 BF16x9의 주요 행렬곱은 FP32보다 2.2~2.5배 빨랐다. 행렬곱 이외의 처리까지 포함한 분자동역학 한 단계는 운용 방식에 따라 1.4~1.7배 빨라졌다.

원자 192개 시스템에서는 BF16x9의 주요 행렬곱이 1.2~1.4배, 분자동역학 단계가 1.1~1.2배 빨랐다. 엔비디아는 이 방식으로 계산한 단일 시점의 에너지와 힘 오차가 PET 모델 자체의 비결정성보다 작았다고 설명했다.

A100의 TF32x3는 FP32 값을 주성분과 잔차 성분으로 나눠 세 차례 TF32 곱을 수행하고 결과를 FP32로 합산한다. BF16x9와 다른 정밀도 형식을 활용하는 방식이다.

원자 1만7280개 시스템에서 TF32x3의 주요 행렬곱은 1.7~1.9배 빨랐고, 분자동역학 단계는 최대 1.5배 빨라졌다. 원자 192개 시스템의 분자동역학 단계는 소형 모델에서 1.2배, 대형 모델에서 1.4배 빨랐다.

두 기법 모두 학습된 모델의 가중치나 구조를 바꾸지 않고 행렬곱 실행 방식을 바꿨다. 모델을 다시 학습할 필요는 없지만 적용 방법은 다르다. BF16x9는 cuBLAS 지원 기능을 이용했고, TF32x3는 PET의 선형층에 별도 커널을 적용했다.

장시간 에너지 보존과 구조 최적화 시험은 TF32x3 결과로 제시됐다. 원자 192개 물 시스템의 에너지 변동은 FP32와 비슷한 범위에 머물렀다. 백금 표면에서 일산화탄소 분자를 이동시키며 측정한 에너지 곡선도 FP32와 비슷하게 매끄러웠다.

원자 539개인 칸토르 합금 표면의 구조 최적화에서는 FP32가 218회, TF32x3가 218~222회 만에 힘 허용치 1.0×10⁻³ eV/Å에 도달했다. 단일 패스 TF32는 500회 안에 수렴하지 않았고, 1000회까지 계산해도 같은 기준을 충족하지 못했다.

BF16x9에는 이와 같은 장시간 에너지 보존·구조 최적화 시험 결과가 제시되지 않았다. 따라서 두 기법이 같은 범위의 물리적 일관성을 검증받았다고 볼 수는 없다.

성능 차이는 시스템 크기와 전체 계산에서 행렬곱이 차지하는 비율에 따라 달라진다. 엔비디아는 이번 수치가 제시된 PET 모델과 GPU, 커널 설정에서 나온 결과라고 밝혔다.

NVIDIA cuBLAS 문서는 BF16x9 FP32 에뮬레이션의 지원 대상으로 컴퓨트 능력 10.0·10.3인 GPU와 CUDA 12.9 이상을 명시한다.