КПД

Метод

В работе Flash Attention 2 по существу еще слегка подкрутили процедуру вычисления и повысили степень параллелизма самой операций.

Алгоритм вычисления
Автор заметил, что операции, не являющиеся матричным умножением, выполняются куда медленнее (в 16 раз), чем матричные умножения, потому переписал алгоритм так, чтобы уменьшить их количество. Казалось бы, их количество невелико, но тем не менее, они занимают существенную часть общего времени работы. Кроме того, при авторегрессионной генерации нужна лишь верхнетреугольная часть матрицы Attention, и вместо того, чтобы считать ее, а затем занулять, ее просто не считают. Вот так вот!

Благодаря перечисленным выше нововведениям удается добиться ускорения 2-3x.

Параллелизм
Flash Attention-1 параллелизует вычисления по размеру батча и числу голов в трансформере, но если батч не слишком большой или трансформер не очень огромный, то многие streaming multiprocessors (SM) простаивают. И чтобы не оставлять их без дела, предлагается паралеллизовывать вычисления и по длине последовательности. На прямом проходе ряды матрицы Attention можно считать независимо, а на обратном проходе - колонки. И каждый поток обрабатывает свой токен. Кроме того, для уменьшения коммуникации между варпами (группами потоков), оказывается целесообразным держать куски матриц ключей (Key) и значений (Values) общими для групп поток, а Query свою на варп (в Flash Attention-1 было наоборот). Уменьшение количество операций чтения/записи приводит к дополнительному ускорению.

Результаты

Flash-Attention-2 сравнивается с Flash-Attention из оригинального репозитория, реализации на triton и xformers. Для замеров рассматривают последовательности длиной от 512 до 16k токенов, и слой attention со скрытой размерностью 2048 (64 или 128 голов).
FlashAttention-2 в 1.3-1.5x быстрее на прямом проходе, и до 2x быстрее на обратном проходе по сравнению с Flash-Attention - 1 (особенно велик выигрыш при использовании causal mask). Flash-Attention - 2 использует до 72% теоретической производительности A100. На H100 разница еще заметнее.

Выводы

Данная история поучительна тем, что одна и та же математическая операция в зависимости от реализации, может выполняться принципиально разное время. Замечательный пример того, что насколько учет особенностей железа, время работы различных компонент, сильных и слабых сторон ускорителя вычислений важен при проектировании алгоритмов.

👍2

www.tgoop.com/quant_prune_distill/55

128 viewsedited  Jul 19, 2023 at 04:29

Метод

В работе Flash Attention 2 по существу еще слегка подкрутили процедуру вычисления и повысили степень параллелизма самой операций.

Алгоритм вычисления
Автор заметил, что операции, не являющиеся матричным умножением, выполняются куда медленнее (в 16 раз), чем матричные умножения, потому переписал алгоритм так, чтобы уменьшить их количество. Казалось бы, их количество невелико, но тем не менее, они занимают существенную часть общего времени работы. Кроме того, при авторегрессионной генерации нужна лишь верхнетреугольная часть матрицы Attention, и вместо того, чтобы считать ее, а затем занулять, ее просто не считают. Вот так вот!

Благодаря перечисленным выше нововведениям удается добиться ускорения 2-3x.

Параллелизм
Flash Attention-1 параллелизует вычисления по размеру батча и числу голов в трансформере, но если батч не слишком большой или трансформер не очень огромный, то многие streaming multiprocessors (SM) простаивают. И чтобы не оставлять их без дела, предлагается паралеллизовывать вычисления и по длине последовательности. На прямом проходе ряды матрицы Attention можно считать независимо, а на обратном проходе - колонки. И каждый поток обрабатывает свой токен. Кроме того, для уменьшения коммуникации между варпами (группами потоков), оказывается целесообразным держать куски матриц ключей (Key) и значений (Values) общими для групп поток, а Query свою на варп (в Flash Attention-1 было наоборот). Уменьшение количество операций чтения/записи приводит к дополнительному ускорению.

Результаты

Flash-Attention-2 сравнивается с Flash-Attention из оригинального репозитория, реализации на triton и xformers. Для замеров рассматривают последовательности длиной от 512 до 16k токенов, и слой attention со скрытой размерностью 2048 (64 или 128 голов).
FlashAttention-2 в 1.3-1.5x быстрее на прямом проходе, и до 2x быстрее на обратном проходе по сравнению с Flash-Attention - 1 (особенно велик выигрыш при использовании causal mask). Flash-Attention - 2 использует до 72% теоретической производительности A100. На H100 разница еще заметнее.

Выводы

Данная история поучительна тем, что одна и та же математическая операция в зависимости от реализации, может выполняться принципиально разное время. Замечательный пример того, что насколько учет особенностей железа, время работы различных компонент, сильных и слабых сторон ускорителя вычислений важен при проектировании алгоритмов.

BY КПД