Отзыв об установке FlashAttention 2 для вывода модели на DGX Spark

Предыстория установки

После установки MOSS-TTS я не испытывал особых неудобств со скоростью вывода. Однако, поскольку было сказано, что установка FlashAttention 2 может ускорить вывод и сократить использование памяти GPU, мне стало любопытно, насколько сильно улучшится производительность. Поэтому я приступил к установке FlashAttention 2.

Что такое FlashAttention 2 и как он повышает эффективность вывода?

Насколько я понимаю, FlashAttention 2 — это реализация, предназначенная для более эффективной обработки операций Attention в архитектурах Transformer. Я предполагал, что улучшение скорости и эффективности памяти достигается за счет сокращения обращений к памяти и создания промежуточных тензоров в процессе вычисления Attention, либо за счет оптимизации потока операций. Однако, я также полагал, что этот эффект может варьироваться в зависимости от структуры модели, длины входных данных, типа данных (torch.float16 / torch.bfloat16) и архитектуры GPU.

Установка FlashAttention 2 на DGX Spark

В README MOSS-TTS есть следующее примечание:

FlashAttention 2 is only available on supported GPUs and is typically used with torch.float16 or torch.bfloat16

Исходя из того, что его можно использовать с torch.float16, я решил, что он должен работать и на DGX Spark, и приступил к установке.

flash-attn2-on-spark

1. Предварительные проверки

  • DGX Spark использует CUDA 13.0, поэтому я искал и устанавливал зависимые пакеты из --extra-index-url https://download.pytorch.org/whl/cu130. Для дополнительной проверки я сверил версию CUDA с помощью nvidia-smi.

  • Во время установки, при поиске PyTorch, мне пришлось использовать уже установленный PyTorch из существующей среды venv. Чтобы предотвратить создание временной изолированной среды при сборке в venv, я добавил опцию --no-build-isolation.

  • В среде Spark установка wheel-пакета завершилась неудачей, и в логе установки было выведено следующее сообщение. Это связано с архитектурой aarch64. Для меня это стало уже привычным явлением при работе со Spark, так что я даже не удивился. И не расстроился. Знакомое сообщение.

Precompiled wheel not found. Building from source...

У меня не было другого выбора, кроме как продолжить сборку из исходников. Поскольку ninja необходим для процесса сборки из исходников, я установил его в venv.

bash pip install ninja

  • На хост-системе требуются библиотеки для разработки Python 3.12, поэтому, если у вас их нет, установите их.
sudo apt update
sudo apt install python3.12-dev -y

flash-attn компилирует код C++ и CUDA и связывает его с Python, для чего требуется файл Python.h, который определяет внутреннюю структуру Python. В обычной среде выполнения Python этот файл не включен, поэтому необходимо отдельно установить пакет для разработчиков.

2. Команда установки

Это ключевой момент данного поста. Учитывая все вышеперечисленные моменты, я установил FlashAttention 2 с помощью следующей команды:

TORCH_CUDA_ARCH_LIST="12.0" MAX_JOBS=1 pip install --no-build-isolation --extra-index-url https://download.pytorch.org/whl/cu130 -e ".[flash-attn]"

3. Обоснование выбора команды (по результатам проб и ошибок)

Изначально на машине также работал gpt-oss-120b. В этом состоянии я запустил pip install ... -e ".[flash-attn]", что привело к резкому скачку использования CPU, система зависла, а терминал перестал отвечать. Пришлось принудительно перезагружать систему с помощью физического переключателя. После этого я остановил все ресурсоемкие задачи и сосредоточился на установке.

После нескольких проб и ошибок мне удалось завершить установку с помощью вышеуказанной команды. Общее время установки составило от 1 до 2 часов. Я не измерял точно, но установка длилась более часа, я отлучился, и когда вернулся после еды, она уже была завершена.

Во время установки память постоянно занимала около 24 ГБ. Проблема была в CPU, и для стабильной установки было лучше остановить другие задачи.

Опции были добавлены по следующим причинам:

  • TORCH_CUDA_ARCH_LIST="12.0": Целью было явно указать архитектуру Blackwell для сокращения времени установки.
  • MAX_JOBS=1: Из-за недавнего зависания системы я установил это значение консервативно на 1. В результате установка заняла более 60 минут.

Эффект повышения производительности вывода после установки

1. Скорость

Честно говоря, значительного прироста скорости я не почувствовал. Даже без FlashAttention я уже получал достаточно высокую скорость, поэтому, даже если и было сокращение на несколько секунд, ощущение «стало заметно быстрее» не было сильным.

  • Для генерации результата длительностью около 7 секунд потребовалось 8-9 секунд.
  • Для генерации результата длительностью 25 секунд потребовалось около 32 секунд.
  • Результат длительностью 16 секунд был получен за 21 секунду.

То есть, время вывода составляло примерно в 1.3 раза больше длины сгенерированного результата.

2. Память

Изменений в использовании памяти также было трудно заметить. Во время вывода я постоянно наблюдал за показаниями nvidia-smi, но использование памяти заметно не увеличивалось и не уменьшалось. Во время вывода потребляемая мощность составляла около 36 Вт, а температура поднималась с 46 до 53 градусов.

Заключение

  • В среде DGX Spark установка FlashAttention 2 через wheel-пакет завершилась неудачей, поэтому пришлось устанавливать из исходников.
  • Сама установка прошла успешно, но сборка заняла много времени и создала значительную нагрузку на CPU.
  • После установки ожидаемого заметного улучшения скорости и использования памяти не произошло.
  • В следующий раз, если придется снова собирать, можно будет увеличить MAX_JOBS до 4. Теоретически, это должно сократить время сборки в 4 раза.

Похожие статьи