Nvidia прискорила обчислення MLIP на Tensor Core за допомогою емуляції FP32

Машинне навчання міжатомних потенціалів (MLIP) дозволяє моделювати молекулярні системи значно швидше за методи квантової механіки на зразок теорії функціонала густини (DFT). Проте симуляції атомарного рівня чутливі до накопичення похибок при розрахунку сил, через що більшість таких моделей змушені виконувати інференс в одинарній точності (FP32), не використовуючи переваги низькоточних операцій тензорних ядер сучасних GPU Nvidia.

Команда Nvidia показала спосіб подолання цього бар'єра за допомогою емуляції FP32 через тензорні ядра без перенавчання нейромереж чи зміни їхньої архітектури. Підхід полягає в декомпозиції операндів FP32 на компоненти нижчої точності, обчисленні перехресних добутків на Tensor Core та їх подальшому накопиченні у FP32.

Реалізація для архітектур Blackwell та Ampere

Nvidia протестувала дві стратегії емуляції:

• BF16x9 на чипах Blackwell: кожен операнд FP32 розкладається на три компоненти BF16 зі шкалюванням, що формує дев'ять попарних добутків. У бібліотеці cuBLAS (для чипів з compute capability 10.0 або 10.3) функціонал доступний «з коробки» через змінну середовища CUBLAS_EMULATE_SINGLE_PRECISION=1 або відповідні прапорці математичного режиму, що не вимагає змін у коді моделей PyTorch.
• TF32x3 на чипах Ampere: операнд ділиться на провідний та залишковий компоненти TF32 без коефіцієнтів масштабування. Обчислюються три добутки (провідний на провідний та два перехресні із залишковими), тоді як найменший залишковий добуток відкидається. Реалізація спирається на оператор cutlass::arch::OpMultiplyAddFastF32 у бібліотеці CUTLASS C++, замінюючи стандартні шари nn.Linear на кастомні.

Результати тестів на моделі PET

Методики перевірили на трансформерній моделі Point Edge Transformer (PET), розробленій Lab COSMO в EPFL. Оцінювалися контрольні точки PET-MAD на 26 млн параметрів (для систем зі 192 та 17 280 атомів води) і на 255 млн параметрів.

• На графічних процесорах Blackwell (B300): метод BF16x9 прискорив великі матричні множення (GEMM) у 2,2–2,5 раза для навантаження на 17 280 атомів, а повний крок молекулярної динаміки (MD) — до 1,7 раза у режимі performant.
• На процесорах Ampere (A100): кастомний шар TF32x3 забезпечив прискорення GEMM у 1,7–1,9 раза, а повного MD-кроку — у 1,5 раза для великої системи. Водночас однопрохідний TF32 хоч і показував більший приріст, спричиняв деградацію точності та розбіжність сил.

Фізична стабільність та релаксація структури

Під час тестів гладкості потенціальної енергії молекули CO на поверхні платини Pt(111) та енергетичного дрейфу в симуляціях NVE TF32x3 зберіг стабільність і плавність кривих, близьку до базової FP32.

При тестуванні структурної релаксації зрізу сплаву Кантора (539 атомів) за допомогою оптимізатора FIRE2 в NVIDIA ALCHEMI Toolkit підхід TF32x3 зійшовся до заданої похибки ($1.0 \times 10^{-3}$ еВ/Å) за 218–222 кроки (проти 218 кроків у FP32), скоротивши час кожного кроку в 1,5 раза. Для порівняння, стандартний однопрохідний режим TF32 не зміг досягти збіжності навіть після подвоєння ліміту до 1000 кроків.