Lambda, Технологический институт Стивенса и команда OpenFold создали ядра GPU с открытым исходным кодом, которые ускоряют обучение OpenFold3 до 1,73 раза, обеспечивая при этом более точные и полностью воспроизводимые градиенты.
Такие модели, как AlphaFold3 и OpenFold3, изменили подход исследователей к изучению белков, нуклеиновых кислот, малых молекул и их взаимодействий. Их обучение и дообучение по-прежнему обходятся дорого, и затраты быстро растут по мере увеличения молекулярных систем.
Значительная часть этих затрат приходится на одну операцию: треугольное внимание (triangle attention).
Поэтому исследователи из Lambda, Технологического института Стивенса и команды OpenFold создали TurboPairFormer — высокопроизводительную реализацию треугольного внимания и треугольного умножения для моделей прогнозирования биомолекулярных структур. Это решение с открытым исходным кодом доступно на PyPI.
Можно выделить три ключевых результата:
- Более быстрое обучение: до 1,73 раза быстрее, чем оригинальный бэкенд OpenFold3 Triton, с наименьшим временем вычислений на GPU на шаг оптимизатора среди восьми протестированных бэкендов на 16 GPU NVIDIA H100.
- Более точные градиенты: самые низкие средние ошибки градиентов запросов (query), ключей (key) и парных смещений (pair-bias) среди восьми реализаций внимания, а также снижение ошибки на 28–47% за счет компенсации выходных остатков в контролируемых экспериментах.
- Воспроизводимые запуски: побитово идентичные градиенты при повторных запусках во всех 600 тестовых случаях с фиксированными входными данными.
Почему треугольное внимание является «узким местом»
Треугольное внимание помогает модели анализировать пары молекулярных компонентов, направляя информацию через третий компонент. Именно так она улавливает геометрию, необходимую для прогнозирования 3D-структур.
Это также ресурсоемкая операция. Стандартное внимание масштабируется квадратично относительно количества токенов. Треугольное внимание масштабируется кубически, поэтому удвоение размера обучающей выборки делает его примерно в 8 раз более затратным.
Существующие реализации, включая NVIDIA® cuEquivariance™ как часть платформы NVIDIA BioNeMo, ByteDance Protenix и бэкенд OpenFold3 Triton, уже ускоряют этот процесс. Однако каждая из них требует компромисса между скоростью, численной точностью и воспроизводимостью. TurboPairFormer был создан для сохранения всех трех параметров на GPU NVIDIA Hopper™.
Как работает TurboPairFormer
TurboPairFormer сочетает аппаратные оптимизации GPU с переработанным алгоритмом вычисления внимания. Он использует три взаимосвязанных метода.
- Он вычисляет вероятности внимания один раз во время обратного прохода. Обратный проход, при котором рассчитываются градиенты для запросов, ключей, значений и парных смещений, является наиболее требовательной частью треугольного внимания. Существующие ядра часто пересчитывают вероятности внимания несколько раз, по одному разу для каждого градиента, что создает избыточные вычисления и накладные расходы. TurboPairFormer пересчитывает каждый тайл вероятностей внимания только один раз и использует этот результат для всех четырех градиентов.
Он вычисляет вероятности внимания один раз во время обратного прохода. Обратный проход, при котором рассчитываются градиенты для запросов, ключей, значений и парных смещений, является наиболее требовательной частью треугольного внимания. Существующие ядра часто пересчитывают вероятности внимания несколько раз, по одному разу для каждого градиента, что создает избыточные вычисления и накладные расходы. TurboPairFormer пересчитывает каждый тайл вероятностей внимания только один раз и использует этот результат для всех четырех градиентов.
- Он сохраняет то, что теряется при округлении BF16, для повышения точности. Существующие ядра часто смешивают точность FP32 и BF16 для баланса скорости и точности. Преобразование значений FP32 в BF16 вносит ошибки округления. TurboPairFormer сохраняет эту потерянную информацию. Его механизм компенсации выходных остатков фиксирует остаток округления при преобразовании выходных данных прямого прохода из FP32 в BF16, а затем использует его для повышения точности градиентов обратного прохода. В контролируемых экспериментах это снизило среднюю ошибку градиента примерно на 28–47%.
Он сохраняет то, что теряется при округлении BF16, для повышения точности. Существующие ядра часто смешивают точность FP32 и BF16 для баланса скорости и точности. Преобразование значений FP32 в BF16 вносит ошибки округления. TurboPairFormer сохраняет эту потерянную информацию. Его механизм компенсации выходных остатков фиксирует остаток округления при преобразовании выходных данных прямого прохода из FP32 в BF16, а затем использует его для повышения точности градиентов обратного прохода. В контролируемых экспериментах это снизило среднюю ошибку градиента примерно на 28–47%.
- Он делает градиенты детерминированными для стабилизации обучения. Ядра, которые накапливают градиенты с помощью атомарных операций с плавающей запятой, могут возвращать слегка различающиеся результаты при идентичных запусках, поскольку порядок операций меняется от запуска к запуску. TurboPairFormer использует фиксированный порядок накопления, поэтому градиенты каждый раз идентичны побитово, что способствует более стабильному обучению.
Он делает градиенты детерминированными для стабилизации обучения. Ядра, которые накапливают градиенты с помощью атомарных операций с плавающей запятой, могут возвращать слегка различающиеся результаты при идентичных запусках, поскольку порядок операций меняется от запуска к запуску. TurboPairFormer использует фиксированный порядок накопления, поэтому градиенты каждый раз идентичны побитово, что способствует более стабильному обучению.
Результаты тестирования
Мы интегрировали TurboPairFormer в OpenFold3 и сравнили его с семью другими конфигурациями бэкендов: ByteDance Protenix, NVIDIA cuEquivariance, TriFast, MegaFold, DeepSpeed-Evo, ScaleFold и оригинальным бэкендом OpenFold3 Triton. Тесты проводились на 16 GPU H100 на двух узлах при размерах выборки 384, 640 и 768 токенов.
TurboPairFormer показал наименьшее время вычислений на GPU на шаг оптимизатора при любом размере выборки.
Время обучения OpenFold3 на 16 GPU H100
Время вычислений на GPU в секундах на шаг оптимизатора. Меньшее значение лучше. Измерения охватывают шаги с 21 по 200 и исключают загрузку данных, передачу данных между хостом и устройством, а также обновления EMA.
При размере выборки 768 TurboPairFormer работает в 1,73 раза быстрее, чем OpenFold3 Triton, в 1,13 раза быстрее, чем NVIDIA cuEquivariance, и от 1,03 до 1,32 раза быстрее, чем остальные пять бэкендов. Разрыв увеличивается по мере роста размера выборки, где треугольное внимание становится наиболее затратным. На GPU NVIDIA H200 TurboPairFormer обучался в 1,54 раза быстрее, чем OpenFold3 Triton при размере выборки 768, сократив время на шаг с 22,53 до 14,67 секунд.
Точность
Мы предоставили каждому ядру идентичные входные данные BF16 и сравнили его градиенты с эталоном высокой точности FP64. TurboPairFormer показал наименьшие средние ошибки для градиентов запросов, ключей и парных смещений при всех протестированных размерах выборки и размерностях внимания. По сравнению с NVIDIA cuEquivariance, он снизил среднюю ошибку градиента на:
- 31,4–44,4% для градиентов запросов
- 28,7–42,4% для градиентов ключей
- 16,9–37,9% для градиентов парных смещений
Точность градиента значений была практически одинаковой. Это ожидаемо: компенсация выходных остатков улучшает градиенты запросов, ключей и парных смещений и не влияет напрямую на градиент значений.
Воспроизводимость и стабильность
Мы также проверили, возвращает ли каждое ядро идентичные градиенты для идентичных входных данных. На одном GPU NVIDIA H100 мы выполнили 600 случаев с фиксированными входными данными по пять раз каждый. Мы считали случай воспроизводимым только в том случае, если все четыре градиента (запросы, ключи, значения и парные смещения) совпадали побитово во всех пяти запусках. TurboPairFormer показал воспроизводимость во всех 600 случаях.
Воспроизводимые градиенты помогают поддерживать стабильность обучения. Чтобы проверить это, мы возобновили обучение с одной и той же контрольной точки модели для каждого ядра, используя идентичные случайные начальные числа и обучающие примеры на графическом процессоре NVIDIA H100. С TurboPairFormer функция потерь при обучении оставалась стабильной, без значительных колебаний и значений NaN.
Что это означает для разработки лекарств
Модели прогнозирования биомолекулярной структуры поддерживают такие задачи, как прогнозирование связывания белок-лиганд, разработка лекарств на основе структуры и изучение семейств белков, связанных с заболеваниями. Адаптация этих моделей обычно подразумевает дообучение на новых экспериментальных структурах, специализацию на семействе белков или многократную переоценку кандидатов. Каждый такой запуск требует вычислительных затрат.
Более быстрые базовые операции позволяют исследователям проводить больше экспериментов на тех же графических процессорах и проверять больше гипотез за меньшее время. А поскольку TurboPairFormer нацелен на оптимизацию базовых блоков, общих для моделей прогнозирования биомолекулярной структуры, полученные преимущества распространяются на широкий спектр приложений ИИ в биомолекулярной сфере.
Начало работы
TurboPairFormer в настоящее время поддерживает графические процессоры H100 и H200.
- Установите package: pip install turbopairformer
- Настройте его с помощью руководства по быстрому запуску
- Ознакомьтесь с paper для получения полной информации о методе и результатах
Авторы: Yide Ran, Chelsea Lowman, Jan Domański, David Hartmann, Jenke Scheen, Jennifer Wei, Chuan Li, Jianwen Xie, Zhaozhuo Xu. Благодарности: Stevens Institute of Technology, Lambda, OpenFold Consortium при Open Molecular Software Foundation
