Проблема точности в атомистическом моделировании
Машинные интератомные потенциалы (MLIP), такие как Point Edge Transformer (PET) от EPFL, позволяют изучать поведение материалов с точностью, близкой к квантовой механике, но при значительно меньших вычислительных затратах. Однако ключевая проблема MLIP — чувствительность к ошибкам округления. Силы, вычисляемые через дифференцирование энергии, должны быть сверхточными, иначе траектория молекулярной динамики (МД) становится нестабильной. Из-за этого большинство моделей до сих пор работают в FP32, игнорируя огромный прирост производительности Tensor Cores, оптимизированных под BF16/TF32.
Решение: эмуляция FP32 через Tensor Cores
Исследователи NVIDIA разработали два метода, которые разлагают FP32-операнды на компоненты более низкой точности, вычисляют их произведения на Tensor Cores и суммируют результат в FP32. Это позволяет использовать аппаратное ускорение без переобучения моделей или изменения их архитектуры.
- BF16x9 (NVIDIA Blackwell): FP32-операнд разбивается на три BF16-компонента. Вычисляется 9 попарных произведений, которые аккумулируются в FP32. Активируется через переменную окружения
CUBLAS_EMULATE_SINGLE_PRECISION=1. - TF32x3 (NVIDIA Ampere): FP32-операнд разбивается на старший (TF32) и остаточный (TF32) компоненты. Вычисляются три перекрестных произведения, остаточный член опускается. Реализовано через CUTLASS
OpMultiplyAddFastF32.
Результаты бенчмарков на модели PET
Тестирование проводилось на моделях PET-MAD (26M и 255M параметров) на системах из 192 и 17 280 атомов. Методы показали значительный ускорение GEMM-операций при сохранении физической точности.
| Метод / Платформа | Ускорение GEMM | Ускорение шага МД | Точность (Force Error) |
|---|---|---|---|
| BF16x9 (Blackwell) / 17 280 атомов | 2.2–2.5x | До 1.5x (оценка) | Ниже уровня недетерминизма модели |
| BF16x9 (Blackwell) / 192 атома | 1.2–1.4x | — | Соответствует FP32 |
| TF32x3 (Ampere) / Общий случай | 1.7–1.9x | До 1.5x | Близка к базовому FP32 |
| TF32x3 (Ampere) / Релаксация (539 атомов) | 1.5x за шаг | — | Сходится за 218–222 шага (как FP32) |
Важность для науки: сохранение энергии
Критическим тестом стала проверка сохранения энергии в микроканоническом ансамбле (NVE). Обычное использование TF32 (одна прохода) приводило к дрейфу энергии и «шероховатости» потенциальной поверхности. Метод TF32x3 сохранил гладкость потенциальной энергии и корректное сохранение энергии, в то время как упрощенные методы TF32 не смогли достичь требуемого допуска по силам даже после 1000 шагов релаксации.
Как внедрить
Для использования BF16x9 на Blackwell достаточно установить переменную окружения CUBLAS_EMULATE_SINGLE_PRECISION=1. Для Ampere рекомендуется использовать NVIDIA ALCHEMI Toolkit. Исходный код и чекпоинты модели PET доступны в репозитории UPET.
Источник: NVIDIA dev blog ↗