Глубокое изучение алгоритма FlashAttention – часть 3
Исчерпывающее изучение алгоритма FlashAttention - Часть 3
Добро пожаловать на третью часть нашей серии о Flash Attention! В этом сегменте мы рассмотрим работу алгоритма FlashAttention V1, разобрав его основные концепции и принципы. Если вы новичок в этой теме или хотите узнать больше о графических процессорах и о том, как работает FlashAttention на высоком уровне, не забудьте посмотреть раздел Understanding GPU & advancement in GPU acceleration в этой серии.
Чтобы начать, давайте уточним, что оптимизации и увеличение скорости FlashAttention в первую очередь направлены на графические процессоры. Хотя в статье упоминается кэш L1 и L2, эти оптимизации фундаментально связаны с производительностью графического процессора, а не с оперативной памятью или другими компонентами памяти.
Быстрый обзор
В типичной графической архитектуре данные хранятся на жестком диске, но для выполнения каких-либо значимых вычислений данные должны быть перемещены в оперативную память. Затем они проходят путь через различные иерархии памяти, пока не попадут в графический процессор. Алгоритм FlashAttention тонко настроен для использования возможностей тензорных ядер в современных графических процессорах. Это особенно важно, так как во время тренировки моделей, таких как GPT-3, тензорные ядра были бездействующими примерно 50% времени.
FlashAttention – заметный алгоритм по двум основным причинам: деление на блоки и повторное вычисление.
- 7 алгоритмов машинного обучения, которые нельзя пропустить
- Продвижение искусственного интеллекта инновационная система памяти университета Сунгкюнкван под названием «Мемория» повышает производительность трансформатора при выполнении сложных задач с длинными последовательностями.
- Sсhneider Electric применяет Retrieval Augmented LLMs на SageMaker, чтобы обеспечить обновления в режиме реального времени в их системах управления предприятием (ERP).
Деление на блоки – это техника, которая разделяет матрицы Q, K и V на более мелкие блоки. Это позволяет алгоритму считывать и обрабатывать эти матрицы блок за блоком, а не загружать все сразу в память графического процессора.
Повторное вычисление, с другой стороны, связано с обратным распространением ошибки, важным аспектом обучения моделей. Вместо сохранения значений в памяти повышенной пропускной способности (HBM) и многократного доступа к этой памяти, FlashAttention перевычисляет значения при необходимости. Хотя повторное вычисление увеличивает количество операций с плавающей точкой, оно значительно сокращает время, затраченное на доступ к памяти.
Теперь давайте перейдем к деталям алгоритма, упомянутым в статье:
Хотите визуальное объяснение? Посмотрите мое видео об алгоритме FlashAttention V1
Алгоритм FlashAttention и вычисление онлайн-нормализатора для Softmax
Алгоритм FlashAttention
Концепция деления на блоки
Первое ключевое понятие в FlashAttention – это деление на блоки. Каждой единице в модели трансформера соответствуют матрицы Q, K и V. Процесс деления на блоки разделяет эти матрицы на управляемые блоки для обработки. Обычно размер блока устанавливается на уровне 128, как упоминают авторы.
Для начала нам нужно определить размеры блоков для матриц Q, K и V. Кроме того, мы инициализируем промежуточные переменные, такие как “l” и “m”, чтобы сохранять результаты. Конечный результат является произведением всех промежуточных переменных. Это деление на блоки и сохранение промежуточных результатов важно для эффективной комбинации этих результатов позже в процессе.
Safe Softmax
Суть FlashAttention заключается в его реализации функции Softmax. Стандартная функция Softmax часто сталкивается с проблемами, связанными с переполнением и недостатком вычислительной точности, поскольку экспоненциальные значения могут стать слишком большими или слишком малыми. FlashAttention использует “safe Softmax” для решения этих проблем.
Safe Softmax работает путем поиска максимального значения во входном массиве и вычитания этого максимального значения из каждого элемента массива перед возведением в экспоненту. Это позволяет избежать возможных проблем с переполнением или недостатком вычислительной точности, обеспечивая численную стабильность вычислений.
Формула для Safe Softmax:
SafeSoftmax с онлайн-вычислением нормализатора
Safe Softmax от FlashAttention черпает вдохновение из концепции, известной как «онлайн-вычисление нормализатора». Этот метод, описанный в статье от NVIDIA, Онлайн вычисление нормализатора для Softmax, предоставляет способ вычисления Softmax без избыточного обращения к памяти. Подход разработан с целью уменьшения количества операций с памятью, делая алгоритм более эффективным с точки зрения использования памяти.
В ходе этого процесса алгоритм поддерживает текущую сумму и изменяет вычисление Softmax с использованием промежуточных значений, рассчитанных на лету. Он достигает этого путем отмены изменений, произведенных в предыдущих итерациях, и применения необходимых коррекций при обработке новых элементов. Этот подход позволяет производить вычисления Softmax без повторного обращения ко всем элементам входного массива, значительно сокращая доступ к памяти.
Комбинирование тайлинга и Safe Softmax
FlashAttention искусно сочетает в себе концепции тайлинга и Safe Softmax, чтобы максимизировать эффективность своего механизма внимания. Путем разделения входных данных на управляемые блоки и применения Safe Softmax на уровне блока, FlashAttention минимизирует доступ к памяти, сохраняя при этом численную стабильность. Этот подход обеспечивает беспрепятственную работу алгоритма с тензорными ядрами графического ускорителя, которые часто используются мало при обучении глубоких моделей обучения языка.
Вывод
В заключение, FlashAttention V1 – это высокоэффективный и численно стабильный механизм внимания, использующий возможности графических процессоров и тензорных ядер. Его инновационное использование тайлинга и Safe Softmax гарантирует минимизацию узких мест доступа к памяти, что позволяет LLM-моделям, таким как GPT-3, обучаться более эффективно.
Интеграция этих техник демонстрирует прекрасное взаимодействие математических концепций, алгоритмического гения и оптимизации оборудования в области глубокого обучения.
