NVIDIA PET: 재학습 없는 텐서코어 가속, 분자역학 스텝 단축

엔비디아가 원자의 에너지와 힘을 예측하는 신경망 모델 PET을 재학습 없이 가속한 결과를 공개했다. 엔비디아의 평가에서는 행렬 연산뿐 아니라 원자 위치와 속도를 갱신하는 분자역학의 한 스텝 전체도 빨라져, 기존 모델로 시뮬레이션 시간을 줄일 가능성을 보였다.

이 계산은 작은 힘 오차가 스텝마다 쌓여 원자의 움직임을 흔들 수 있다. 행렬 연산 전용 GPU 연산부인 텐서코어를 쓰더라도 속도와 수치 안정성을 함께 확보해야 하는 이유다(원문).

엔비디아 측정에서 B300에서는 큰 행렬의 BF16x9 가속이 두드러졌고(Figure 3), A100에서도 TF32x3가 모든 제시 모양에서 자체 FP32보다 빨랐다(Figure 5). 장치가 달라 두 기법의 우열을 뜻하지는 않는다.

BF16x9의 performant 전략은 cuBLAS가 더 빠르다고 예상할 때만 쓴다. 엔비디아 측정에서 전체 스텝은 컴파일된 FP32 대비 192원자에서 1.2배, 17,280원자에서 1.7배 빨랐다. 엔비디아가 측정한 TF32x3의 가속률은 큰 시스템에서 자체 FP32 기준 1.5배였다(Figure 4·6과 설명).

작은 시스템에서는 GPU 작업 시작 시간과 이웃 원자 처리 같은 비행렬 작업의 비중이 커 이득이 줄어든다. TF32x3는 행렬 모양·GPU에 맞춘 커널 튜닝과 모델의 정확도·물리 안정성 요구에 따른 검증이 필요하다(한계와 선택 기준).

엔비디아의 정확도 평가에서는 BF16x9가 만든 오차가 같은 모델을 실행해도 생기는 자체 변동보다 작았다. 엔비디아 측정에서 TF32x3의 힘 오차는 FP32를 재컴파일해 실행했을 때의 차이보다 약 3배였다(정밀도 분석).

539원자 합금 표면의 위치를 안정한 배치로 조정하는 구조 완화에는 PET·FIRE2·NVIDIA ALCHEMI Toolkit을 썼다. 엔비디아의 3회 반복 평가에서 FP32는 218스텝, TF32x3는 218~222스텝에 힘 허용오차 1.0×10⁻³ eV/Å에 도달했다. 엔비디아 측정에서 TF32 변환을 한 번만 하는 경로는 스텝당 더 빨랐지만 1,000스텝까지도 수렴하지 못했다(Figure 8과 설명).

엔비디아가 백금 표면에서 일산화탄소 분자를 떼어내며 평가한 시험에서도 TF32x3는 FP32의 매끄러운 에너지 곡선을 따랐다. 엔비디아가 별도로 입자 수·부피·에너지를 고정한 192원자 물의 30 ps 시뮬레이션을 평가한 결과에서는 FP32와 비슷한 에너지 변동을 보였다. 엔비디아 평가에서 단일 변환 TF32는 각각 곡선의 흔들림과 에너지가 한쪽으로 치우치는 현상을 보였다(Figure 7과 설명).