Активация FlashAttention в comfyui для Nvidia Tesla V100 в linux

AI-вычисления требуют наличия большого объема быстродействующей памяти, а также использования новейшего программных модулей, в полной мере раскрывающих возможности имеющегося оборудования. Для видеокарты Nvidia Tesla V100 в comfui для этого можно установить модули Triton 3.2.0 и SageAttention 1.0.6:

  pip install sageattention==1.0.6 triton=3.2.0

Для еще большего увеличения быстродействия вычислений на видеокарте Nvidia Tesla V100 (и на других подходящих ускорителях) в comfyui желательно использовать программный модуль FlashAttention. В нем задействован алгоритм эффективного вычисления внимания (attention) в AI-трансформерах, который значительно ускоряет инференс и уменьшает загрузку памяти на GPU.

Технология FlashAttention (Fast and Memory-Efficient Exact Attention with IO-Awareness) использует оптимизированный алгоритм работы с памятью GPU, включающий следующие элементы:

• tiling (блочное разбиение) — разбиение Q, K, V-данных на маленькие блоки, которые помещаются в кэш GPU;
  
• вычисление attention по блокам, без хранения полной объемной матрицы внимания в VRAM;
  
• использование online softmax (механизма пересчёта на лету) и переупорядочивание операций;
  
• минимизация количества операций чтения/записи между медленной VRAM-памятью и быстрым кешем видеокарты.
  

Благодаря FlashAttention, сохраняя тот же математический результат, можно увеличить скорость проведения AI-вычислений в 2–4 раза, а также снизить потребление видеопамяти.

Проверить поддержку технологии flash_attn в использующейся конфигурации comfyui можно командой:

  python -c 'import flash_attn; print(f"Version: {flash_attn.__doc__}")'

На компьютере с видеокартой Nvidia Tesla V100 в штатной конфигурации comfyui отображается информация об отсутствии поддержки FlashAttention:

Traceback (most recent call last):
File "<string>", line 1, in <module> import flash_attn; print(f"Version: {flash_attn.__doc__}") 
ModuleNotFoundError: No module named 'flash_attn' 

Из-за отсутствия модуля flash_attn при AI-вычислениях не в полной мере задействуются тензорные ядра этого графического ускорителя. В частности, не используется инструкция mma.sync.aligned.m8n8k4, необходимая для эффективного выполнения матричных операций. Она напрямую задействует Tensor Core GPU Nvidia Tesla V100, что позволяет достичь максимальной производительности и уменьшить потребление видеопамяти на 30-40%.

Для запуска comfyui с flash-attention можно использовать команду:

  python main.py --enable-manager --use-flash-attention

К сожалению, полноценная поддержка flash-attn доступна только при использовании видеокарт Nvidia с архитектурой Ampere и новее. При проведении определенніх манипуляций, можно запустить модуль и на видеокартах Nvidia Turing, а также Volta (Nvidia Tesla V100).

Как добавить поддержку flash-attn в comfyui для Nvidia Tesla V100?

При запуске comfyui на компьютере с видеокартой Nvidia Tesla V100 в терминале появляется сообщение об опциональной возможности установить flash-attn командой:

  pip install flash-attn

Однако, при попытке установить модуль FlashAttention этой командой возникает ошибка Failed to build ‘flash-attn’ when getting requirements to build wheel:

Для видеокарт Nvidia Tesla V100 в comfyui можно установить неофициальный модуль flash-attn c репозитория peisuke на github:

Чтобы добавить поддержку flash_attn для Nvidia Tesla V100, выполняются команды:

  python3 -m venv venv && source venv/bin/activate && python -m pip install --upgrade pip wheel setuptools

  cd custom_nodes && git clone https://github.com/peisuke/flash-attention.git && cd flash-attention

  MAX_JOBS=1 pip install flash-attn --no-build-isolation

Аргумент MAX_JOBS=1 нужен для ограничения потребления памяти. Даже с таким ограничением при компиляции модуля в gentoo потребление памяти периодически превышало 30GB! Время сборки flash-attn на компьютере с 40GB VRAM с восьмиядерным процессором Intel QTJ0 в ОС gentoo заняло больше двух часов:

Получить информацию об установленой версии flash-attn можно командой:

  pip show flash_attn

При загрузке comfyui теперь виден активный модуль flash-attn:

AI-вычисления с flash-attn на NvidiaTesla V100 можно производить только с FP16-инструкциями, так как BF16 не поддерживается аппаратно, а FP32 — модулем flash-attn. Поэтому при работе в comfyui необходимо использовать AI-модели в формате FP16. Несмотря на ​это, прирост быстродействия при генерации в comfyui ощутим. При использовании неподходящей модели возникает ошибка:

  Flash Attention failed, using default SDPA: FlashAttention only supports Ampere GPUs or newer. 

В этом случае инференс все равно работает, используя pytorch-attn.

P.S. В списке custom-нод comfyui есть модули Flash Attention с поддержкой Nvidia Tesla V100, например, ComfyUI-ComfyUI-Flash-Attention_v100 от fearl0rd.