NVIDIA, PET 원자 시뮬레이션에 Tensor Core FP32 시험… 커널과 전체 단계 속도 구분 필요
NVIDIA는 10월 29일 기계학습 원자간 퍼텐셜의 정밀도와 속도를 비교한 실험을 발표했다. 이 사이트는 원 발표일을 유지한 채 11월 1일 기술호에 싣는다. 특정 Point Edge Transformer, 즉 PET 모델과 GPU에서 Tensor Core로 FP32에 가까운 행렬곱을 수행한 결과다.
NVIDIA는 10월 29일 기계학습 원자간 퍼텐셜의 정밀도와 속도를 비교한 실험을 발표했다. 이 사이트는 원 발표일을 유지한 채 11월 1일 기술호에 싣는다. 특정 Point Edge Transformer, 즉 PET 모델과 GPU에서 Tensor Core로 FP32에 가까운 행렬곱을 수행한 결과다. 모든 과학 모델이 같은 배수로 빨라진다는 증거는 아니다. 행렬 커널의 속도와 분자 동역학 한 단계 전체의 속도를 분리해서 읽어야 한다.
원자간 퍼텐셜은 원자 배치의 에너지를 예측하고 그 기울기에서 얻은 힘이 여러 단계의 위치와 속도를 결정한다. 작은 계산 오차도 긴 궤적에서 쌓일 수 있어 많은 모델이 FP32 추론을 유지한다. NVIDIA의 방법은 PET를 재학습하거나 구조를 바꾸지 않는다. FP32 피연산자를 낮은 정밀도 성분으로 나눠 Tensor Core에서 일부 곱을 계산하고 FP32로 누적한다. Blackwell의 BF16x9는 각 값을 세 성분으로 나눠 아홉 곱을 사용한다. Ampere의 TF32x3는 세 곱만 쓰고 가장 작은 잔차끼리의 곱을 생략한다. 둘의 오차와 속도를 하나의 수치로 합칠 수 없다.
실험 조건도 구체적이다. 2,600만 매개변수 PET-MAD 소형 체크포인트를 192개 원자 물 상자와 1만 7,280개 원자 초격자에서 시험했다. 2억 5,500만 매개변수 대형 체크포인트는 192개 원자에서 시험했다. 비교 기준은 각 조건의 기본 FP32 경로다. Blackwell B300에서 BF16x9를 쓸 때 큰 시스템의 주요 GEMM 커널은 2.2~2.5배 빨랐다. 분자 동역학 한 단계 전체는 같은 큰 시스템에서 performant 방식으로 1.7배, eager 방식으로 1.4배 빨랐다. 192개 원자에서는 전체 단계가 각각 1.2배와 1.1배였다. 커널의 2.5배를 전체 시뮬레이션의 배수로 쓰면 안 된다.
Ampere A100의 TF32x3 시험은 CUTLASS `OpMultiplyAddFastF32` 커널과 맞는 선형층을 대체하는 `TF32Linear`를 사용하며 입력과 출력은 FP32다. 큰 시스템의 주요 GEMM은 FP32 기준 약 1.7~1.9배, 전체 단계는 약 1.5배 빨랐다. 192개 원자의 대형 모델은 전체 단계가 1.4배, 같은 크기의 소형 모델은 1.2배였다. NVIDIA는 단일 TF32 방식이 더 빠를 수 있지만 정확성과 궤적 안정성에서 더 큰 손실이 있다고 설명한다. 가장 빠른 산술 방식이 반드시 유용한 물리 계산은 아니다.
물리적 일관성도 검사했다. CO 분자를 Pt(111) 표면에서 떼어 내는 실험에서 TF32x3의 에너지 곡선은 FP32처럼 매끄러웠지만 단일 TF32에는 눈에 띄는 흔들림이 있었다. 192개 원자 물 상자의 30피코초 NVE 시뮬레이션에서도 TF32x3는 FP32와 비슷한 제한된 에너지 변동을 보였고 단일 TF32는 더 큰 추세적 드리프트를 보였다. 이는 선택된 재료와 기간에서의 결과이지 모든 물질의 장기 보존을 보장하지 않는다.
539개 원자 합금 판의 구조 이완 시험 세 번에서는 FP32가 218단계에 수렴했고 TF32x3는 218~222단계에 1.0×10⁻³ eV/Å의 힘 기준을 충족했다. 단일 TF32 세 번은 500단계 제한에서 모두 실패했고 한도를 1,000으로 늘려도 기준에 도달하지 못했다고 원문은 전한다. TF32x3는 단계별 속도를 1.5배 높이면서 시험된 수렴을 유지했다. 단일 TF32는 단계별 2배 빠를 수 있어도 목표를 놓쳤다.
실제 적용 전에는 자기 모델에서 큰 GEMM이 차지하는 시간과 에너지, 힘, 긴 궤적, 수렴을 함께 검사해야 한다. 작은 시스템은 이웃 처리와 그래프 작업, 실행 오버헤드 때문에 전체 이득이 작을 수 있다. 여기에 적은 배수는 NVIDIA가 해당 PET와 GPU에서 보고한 것이며 이 사이트가 재측정한 결과가 아니다. 모델 크기, 원자 수, 측정 단위와 FP32 기준을 붙여 비교해야 한다.
성능 수치를 비교하려면 분모와 측정 범위를 고정해야 한다. GEMM 커널은 전체 분자 동역학 단계의 일부이므로 커널의 최고 배속을 전체 시뮬레이션 속도로 바꿔 말할 수 없다. 에너지 오차가 작더라도 힘의 미분과 긴 궤적이 안정적이라는 보장은 없다. 자체 모델에서는 원자 수, 체크포인트 크기, GPU, FP32 기준 구현, 수렴 임계값을 기록하고 에너지와 힘을 함께 비교해야 한다. 이 기사에 적힌 배속은 NVIDIA가 시험한 PET 작업의 보고값이며 다른 재료에 대한 독립 검증 결과가 아니다.
Sources
FAQ
BF16x9와 TF32x3의 계산법과 GPU 범위는 어떻게 다른가요?
Blackwell의 BF16x9는 FP32 값마다 BF16 성분 세 개로 나눠 아홉 곱을 계산한다. Ampere의 TF32x3는 주요 TF32 곱 세 개를 쓰고 가장 작은 잔차 곱을 생략한다. 둘 다 FP32로 누적하고 PET를 재학습하지 않지만 하드웨어와 정확도 특성이 다르다.
2.5배와 1.7배를 모두 전체 시뮬레이션 가속으로 불러도 되나요?
안 된다. 1만 7,280개 원자 PET 부하에서 B300 BF16x9의 큰 GEMM 커널은 최고 2.5배였지만 전체 MD 단계는 performant 방식으로 1.7배였다. A100 TF32x3에는 커널 약 1.7~1.9배와 전체 단계 약 1.5배의 별도 결과가 있다. GPU, 규모와 FP32 기준을 함께 밝혀야 한다.
단일 예측 외에 궤적과 구조 수렴을 시험한 이유는 무엇인가요?
힘은 예측 에너지에서 나오므로 작은 오차도 여러 단계에 쌓일 수 있다. NVIDIA는 CO/Pt(111) 곡선, 30피코초 물 궤적과 539개 원자 합금 이완을 살폈다. TF32x3는 218~222단계에 수렴했지만 빠른 단일 TF32는 500단계 안에 실패했다. 다른 모델은 자체 검증이 필요하다.