flash cosine sim attention

_{Погружение в глубокое обучение, переработанное журналом Quanta}

Внимание: Flash-косинус-сходство

Реализация внимания на сходство слитого косинуса в том же стиле, что и Flash Attention. Наблюдение состоит в том, что, приняв нормализованные запросы и ключи l2, вам больше не нужно отслеживать максимумы строк для обеспечения числовой стабильности. Это значительно упрощает алгоритм мгновенного внимания, предполагая, что внимание к косинусному сходству не требует затрат на обобщение.

Другими словами, стабильное, быстрое, эффективное использование памяти и длительное внимание к контексту без каких-либо недостатков.

Обновление: К сожалению, эксперименты Робина показали гораздо худшие оценки FID, что не отразилось на потерях. Ожидаются новые эксперименты. Используйте эту библиотеку с осторожностью.

Обновление 2: Единственным спасением будет использование сгруппированного l2norm, которое потенциально может обеспечить большую выразительность. Если кто-нибудь сможет оценить эту технику в своей генеративной работе и получить некоторые оценки FID, я буду очень признателен.

Обновление 3. Подход, аналогичный косинусному симуляции внимания, был проверен в масштабе с помощью модели зрения с параметрами 22B от Brain.

Статус (wip)

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

Однако при неавторегрессии без маскировки архитектура на A100 по-прежнему медленнее для F16. Цель состоит в том, чтобы заставить его работать быстрее на A100 вперед и назад как для F32, так и для F16, поскольку общая память еще не полностью использована.

На старых графических картах без достаточного количества общей памяти придется оценивать соотношение эффективности и скорости памяти в зависимости от длины обучаемой последовательности.

Признательность

• Артуру Хеннекину за обучение моему первому ядру CUDA и за написание простой эталонной реализации, которая помогла мне загрузить первое ядро, производительность которого достигла разумного уровня, до базового уровня. Эта работа была бы невозможна без его опыта.

• Борису Дайме и Робину Ромбаху за проведение экспериментов с упрощенной косинусной симуляцией внимания с фиксированным масштабированием на некоторых важных моделях преобразования текста в изображение и проверку того, что она действительно работает так же хорошо, как обычное внимание.

• Маркусу Рабе за написание статьи, в которой показано, что внимание не требует памяти O(n²), и Три Дао за объединение всего этого в реализации ядра CUDA для регулярного внимания, демонстрацию превосходства в скорости с использованием мозаичного подхода, минимизирующего доступ к HBM (и за вычисление out dO * O == dP * P для прохода назад). Без их открытий я бы не смог завершить свое паломничество в поисках окончательной формулировки внимания.

• Stability.ai за щедрую спонсорскую поддержку передовых исследований в области искусственного интеллекта.

Установить

$ pip install flash-cosine-sim-attention

Использование

Внимание к себе

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v )  # (1, 8, 1024, 64)

Перекрестное внимание

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v ) # (1, 8, 1024, 64)

С маскированием ключа/значения

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

mask = torch . ones ( 1 , 2048 ). bool (). cuda ()

out = flash_cosine_sim_attention ( q , k , v , mask = mask ) # (1, 8, 1024, 64)

авторегрессионный

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True )  # (4, 8, 1024, 64) 

Разнообразный

Одноглавый ключ/значения (Shazeer и др. и используются в PaLM)

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True )  # (4, 8, 1024, 64)

Если вам нужно выполнять операции с запросами и ключами между l2norm и фактическим шагом внимания, просто установите l2norm_qk = False

бывший.

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention , l2norm_tensors

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 1024 , 64 ). cuda ()

q , k = l2norm_tensors ( q , k )

# do your rotation of queries and keys
# say with https://github.com/lucidrains/rotary-embedding-torch

out = flash_cosine_sim_attention ( q , k , v , l2norm_qk = False )  # (4, 8, 1024, 64)

Перекрестите внимание с причинно-следственными действиями, как и ожидалось - (кэширование ключей и значений в авторегрессии во время вывода или обучение, подобное Transformer-XL)

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True ) # (1, 8, 1024, 64)

Если у вас объединены размеры партии и головки, это нормально.

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 32 , 1024 , 64 ). cuda ()
k = torch . randn ( 32 , 2048 , 64 ). cuda ()
v = torch . randn ( 32 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True ) # (32, 1024, 64) 

Поддерживаемые размеры головки

• 16 - ф32

• 32

• 64

• 96

• 128

• 16 -f16

• 80 - в процессе

Тодо

• поддержка bfloat16, используйте sfinae, как рекомендует Артур

• поток из qk_mma в разделяемую память частями для расчета mma, посмотрите, можно ли использовать освобожденный smem для кэширования большего количества

• поддержка динамического позиционного смещения O(n) 1d

• выяснять, почему кэширование smem-фрагментов приведет к снижению производительности, это не имеет смысла

• подумайте об использовании logsumexp — работает, но дополнительный журнал приводит к снижению производительности

• подготовить механизм кэширования фрагментов smem, чтобы обеспечить максимально возможное кэширование на A100 (или f16)

• сделать обработку размера плитки внимания настраиваемой для обратного прохода

• переместить атомарное добавление в перегруженную функцию внутри mma

• гибкий, какой тип используется для накопления

• протестируйте тайлы 64х96 на f16

• ввести версию с эффективным использованием памяти ЦП (только для вывода, поскольку обучение не имеет смысла), используя простой код Pytorch

• выяснить, как распределять по-разному для архитектур (скажем, A100), в случае обратного изменения можно по-другому использовать увеличение общей памяти

• разделить размеры строк и столбцов для плиток внимания

• dk и dv теперь находятся в f16, когда это возможно (не одноголовый kv)

• поддержка большего количества стандартных размеров головки (wip)

• еще раз отладить и исправить смещение градиентов назад для размера головы 32

• исправить градиенты смещения внимания

• разрешить использование одноглавых ключей/значений, как в PaLM

• исправить атомное дополнение для f16

• Смещение внимания должно иметь возможность принимать измерения дополнительного измерения пакета, для Alphafold2, например, смещение внимания

• автоматизировать очистку кеша ядра, используя версию в качестве суффикса к имени пакета

• решить причинно-следственные числовые проблемы f16

• перенять все знания от прямого ядра к обратному ядру и убедиться, что оно превосходит по производительности хотя бы на A100

Описание

В настоящее время внимание к косинусному подобию не получило широкого применения в промышленности. Единственная крупная модель, обученная с его помощью, — это SwinV2. Если кто-то может аннулировать этот подход, пожалуйста, откройте проблему или отправьте мне электронное письмо. Проводить эксперименты с регулярным вниманием можно с помощью репозитория x-transformers.

Обновление: Борис Дайма любезно начал эксперимент (синий с красным в качестве базовой линии) для проверки внимания к косинусному сходству с фиксированной шкалой 10 в условиях реальной модели.

Обновление 2: Внимание к косинусному сходству было доказано в реальной сети внимания к изображению с использованием постоянной шкалы 10 . Не хуже, чем обычное внимание. Благодарность принадлежит Борису Дайме за то, что он потратил время на проведение эксперимента и развеял сомнения относительно техники.

Обновление 3: Робин Ромбах протестировал ядро ​​в этом репозитории с размером головы 64 и фиксированным масштабом 10 в модели преобразования текста в изображение, не наблюдая никаких отличий от обычного внимания. Ожидаются дополнительные оценки.

Обновление 4. Улучшение производительности, наблюдаемое в экспериментах Бориса, вероятно, связано с тем, что внимание косинусного моделирования позволяет переключаться с конфигурации до Layernorm на конфигурацию после Layernorm в преобразователях (поскольку l2norm фактически заменяет предварительную норму). норма слоя). Внимание, вызванное косинусным моделированием, скорее всего, даст такие же результаты, как и обычное внимание, без каких-либо других изменений в преобразователе.

Тестирование

Для тестирования выходные данные и градиенты равны для сценариев без авторегрессии и авторегрессии.

$ python setup.py test 

Бенчмаркинг

Обязательно сначала установите ядро ​​CUDA

$ python setup . py install

Затем

$ python benchmark . py

Для сравнения только вперед или назад добавьте к приведенному выше флаг --only-forwards или --only-backwards . Чтобы оценить авторегрессию, добавьте --causal

Тесты – уничтожение

ГТХ 2080 Ти

Вперед

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 1.05x   kernel: 0.24ms  baseline: 0.23ms
seq_len: 256    slower: 1.27x   kernel: 0.38ms  baseline: 0.30ms
seq_len: 512    slower: 1.28x   kernel: 0.87ms  baseline: 0.68ms
seq_len: 1024   slower: 1.15x   kernel: 2.63ms  baseline: 2.28ms
seq_len: 2048   slower: 0.99x   kernel: 7.99ms  baseline: 8.10ms
seq_len: 4096   slower: 0.88x   kernel: 30.82ms baseline: 34.84ms
seq_len: 8192   slower: 0.00x   kernel: 121.96ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.85x   kernel: 0.20ms  baseline: 0.24ms
seq_len: 256    slower: 0.97x   kernel: 0.24ms  baseline: 0.25ms
seq_len: 512    slower: 1.22x   kernel: 0.43ms  baseline: 0.35ms
seq_len: 1024   slower: 0.95x   kernel: 0.93ms  baseline: 0.98ms
seq_len: 2048   slower: 0.90x   kernel: 3.16ms  baseline: 3.50ms
seq_len: 4096   slower: 0.85x   kernel: 11.06ms baseline: 13.07ms
seq_len: 8192   slower: 0.00x   kernel: 42.61ms baseline: oom

Назад - еще нужно поработать

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 1.07x   kernel: 0.61ms  baseline: 0.57ms
seq_len: 256    slower: 1.40x   kernel: 0.91ms  baseline: 0.65ms
seq_len: 512    slower: 1.70x   kernel: 2.34ms  baseline: 1.38ms
seq_len: 1024   slower: 1.26x   kernel: 5.67ms  baseline: 4.50ms
seq_len: 2048   slower: 1.29x   kernel: 20.60ms baseline: 15.91ms
seq_len: 4096   slower: 1.30x   kernel: 78.93ms baseline: 60.81ms
seq_len: 8192   slower: 0.00x   kernel: 314.51ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.91x   kernel: 0.50ms  baseline: 0.55ms
seq_len: 256    slower: 1.06x   kernel: 0.58ms  baseline: 0.55ms
seq_len: 512    slower: 1.13x   kernel: 0.81ms  baseline: 0.72ms
seq_len: 1024   slower: 0.97x   kernel: 2.09ms  baseline: 2.16ms
seq_len: 2048   slower: 0.96x   kernel: 7.06ms  baseline: 7.35ms
seq_len: 4096   slower: 0.97x   kernel: 26.08ms baseline: 26.84ms
seq_len: 8192   slower: 0.00x   kernel: 101.02ms    baseline: oom

Вперед и назад — F32 определенно медленнее

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 1.05x   kernel: 0.83ms  baseline: 0.79ms
seq_len: 256    slower: 1.34x   kernel: 1.26ms  baseline: 0.95ms
seq_len: 512    slower: 1.44x   kernel: 3.14ms  baseline: 2.18ms
seq_len: 1024   slower: 1.15x   kernel: 7.83ms  baseline: 6.81ms
seq_len: 2048   slower: 1.20x   kernel: 28.83ms baseline: 24.03ms
seq_len: 4096   slower: 1.20x   kernel: 111.13ms    baseline: 92.51ms
seq_len: 8192   slower: 0.00x   kernel: 441.70ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.89x   kernel: 0.68ms  baseline: 0.77ms
seq_len: 256    slower: 1.03x   kernel: 0.80ms  baseline: 0.77ms
seq_len: 512    slower: 1.06x   kernel: 1.16ms  baseline: 1.10ms
seq_len: 1024   slower: 0.93x   kernel: 2.94ms  baseline: 3.16ms
seq_len: 2048   slower: 0.93x   kernel: 10.06ms baseline: 10.87ms
seq_len: 4096   slower: 0.93x   kernel: 37.09ms baseline: 39.96ms
seq_len: 8192   slower: 0.00x   kernel: 143.13ms    baseline: oom

Для авторегрессии явный выигрыш python benchmark.py --causal

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.97x   kernel: 0.81ms  baseline: 0.84ms
seq_len: 256    slower: 1.07x   kernel: 1.12ms  baseline: 1.05ms
seq_len: 512    slower: 0.83x   kernel: 2.23ms  baseline: 2.68ms
seq_len: 1024   slower: 0.55x   kernel: 4.83ms  baseline: 8.82ms
seq_len: 2048   slower: 0.49x   kernel: 15.89ms baseline: 32.68ms
seq_len: 4096   slower: 0.46x   kernel: 57.50ms baseline: 126.00ms
seq_len: 8192   slower: 0.00x   kernel: 224.76ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.69ms  baseline: 0.84ms
seq_len: 256    slower: 0.95x   kernel: 0.79ms  baseline: 0.83ms
seq_len: 512    slower: 0.78x   kernel: 1.06ms  baseline: 1.37ms
seq_len: 1024   slower: 0.50x   kernel: 2.10ms  baseline: 4.24ms
seq_len: 2048   slower: 0.37x   kernel: 5.85ms  baseline: 15.92ms
seq_len: 4096   slower: 0.31x   kernel: 19.80ms baseline: 64.42ms
seq_len: 8192   slower: 0.00x   kernel: 75.25ms baseline: oom

Для последовательностей переменной длины с маскированием тоже явный выигрыш. Предположим, что в среднем 25% токенов замаскированы. python benchmark.py --mask-prob 0.25

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.95x   kernel: 0.84ms  baseline: 0.89ms
seq_len: 256    slower: 1.19x   kernel: 1.28ms  baseline: 1.08ms
seq_len: 512    slower: 1.23x   kernel: 3.19ms  baseline: 2.59ms
seq_len: 1024   slower: 0.92x   kernel: 8.19ms  baseline: 8.88ms
seq_len: 2048   slower: 0.92x   kernel: 30.08ms baseline: 32.57ms
seq_len: 4096   slower: 0.94x   kernel: 123.20ms    baseline: 131.22ms
seq_len: 8192   slower: 0.00x   kernel: 461.77ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.85x   kernel: 0.77ms  baseline: 0.90ms
seq_len: 256    slower: 0.93x   kernel: 0.86ms  baseline: 0.93ms
seq_len: 512    slower: 0.93x   kernel: 1.31ms  baseline: 1.40ms
seq_len: 1024   slower: 0.76x   kernel: 3.31ms  baseline: 4.35ms
seq_len: 2048   slower: 0.71x   kernel: 11.19ms baseline: 15.65ms
seq_len: 4096   slower: 0.70x   kernel: 41.27ms baseline: 59.01ms
seq_len: 8192   slower: 0.00x   kernel: 158.60ms    baseline: oom

А100 40ГБ (wip)

Благодарим компанию Stability за предоставление доступа к A100 для тестирования. Спасибо Энрико за то, что он нашел время провести несколько тестов, когда у меня еще не было доступа.

A100 все еще находится в стадии разработки. Общая память еще не полностью использована. Как ни странно, у F32 дела обстоят лучше, чем у F16.

Нападающие

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.98x   kernel: 0.29ms  baseline: 0.30ms
seq_len: 256    slower: 1.19x   kernel: 0.35ms  baseline: 0.29ms
seq_len: 512    slower: 0.94x   kernel: 0.52ms  baseline: 0.55ms
seq_len: 1024   slower: 0.75x   kernel: 1.23ms  baseline: 1.65ms
seq_len: 2048   slower: 0.88x   kernel: 4.17ms  baseline: 4.73ms
seq_len: 4096   slower: 0.79x   kernel: 14.53ms baseline: 18.36ms
seq_len: 8192   slower: 0.64x   kernel: 55.01ms baseline: 85.93ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.84x   kernel: 0.24ms  baseline: 0.29ms
seq_len: 256    slower: 1.02x   kernel: 0.29ms  baseline: 0.29ms
seq_len: 512    slower: 1.24x   kernel: 0.36ms  baseline: 0.29ms
seq_len: 1024   slower: 1.48x   kernel: 0.79ms  baseline: 0.54ms
seq_len: 2048   slower: 1.31x   kernel: 2.08ms  baseline: 1.59ms
seq_len: 4096   slower: 1.21x   kernel: 6.89ms  baseline: 5.70ms
seq_len: 8192   slower: 1.07x   kernel: 24.80ms baseline: 23.15ms

Назад

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.94x   kernel: 0.57ms  baseline: 0.60ms
seq_len: 256    slower: 1.29x   kernel: 0.75ms  baseline: 0.58ms
seq_len: 512    slower: 1.16x   kernel: 1.30ms  baseline: 1.12ms
seq_len: 1024   slower: 0.98x   kernel: 3.14ms  baseline: 3.19ms
seq_len: 2048   slower: 1.05x   kernel: 11.13ms baseline: 10.63ms
seq_len: 4096   slower: 0.98x   kernel: 40.11ms baseline: 40.79ms
seq_len: 8192   slower: 0.97x   kernel: 154.96ms    baseline: 159.70ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.91x   kernel: 0.55ms  baseline: 0.60ms
seq_len: 256    slower: 1.03x   kernel: 0.62ms  baseline: 0.60ms
seq_len: 512    slower: 1.36x   kernel: 0.82ms  baseline: 0.60ms
seq_len: 1024   slower: 1.52x   kernel: 1.52ms  baseline: 1.01ms
seq_len: 2048   slower: 1.37x   kernel: 4.14ms  baseline: 3.03ms
seq_len: 4096   slower: 1.33x   kernel: 14.23ms baseline: 10.71ms
seq_len: 8192   slower: 1.34x   kernel: 53.90ms baseline: 40.28ms

Вперед и назад

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.92x   kernel: 0.80ms  baseline: 0.87ms
seq_len: 256    slower: 1.23x   kernel: 1.07ms  baseline: 0.87ms
seq_len: 512    slower: 1.08x   kernel: 1.80ms  baseline: 1.66ms
seq_len: 1024   slower: 0.94x   kernel: 4.33ms  baseline: 4.62ms
seq_len: 2048   slower: 0.99x   kernel: 15.26ms baseline: 15.44ms
seq_len: 4096   slower: 0.93x   kernel: 54.78ms baseline: 59.21ms
seq_len: 8192   slower: 0.91x   kernel: 210.38ms    baseline: 230.97ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.90x   kernel: 0.78ms  baseline: 0.86ms
seq_len: 256    slower: 1.00x   kernel: 0.87ms  baseline: 0.87ms
seq_len: 512    slower: 1.36x   kernel: 1.18ms  baseline: 0.86ms
seq_len: 1024   slower: 1.49x   kernel: 2.31ms  baseline: 1.55ms
seq_len: 2048   slower: 1.33x   kernel: 6.17ms  baseline: 4.63ms
seq_len: 4096   slower: 1.28x   kernel: 21.08ms baseline: 16.44ms
seq_len: 8192   slower: 1.24x   kernel: 78.75ms baseline: 63.45ms

авторегрессионный

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.82ms  baseline: 1.01ms
seq_len: 256    slower: 1.02x   kernel: 1.00ms  baseline: 0.98ms
seq_len: 512    slower: 0.82x   kernel: 1.55ms  baseline: 1.89ms
seq_len: 1024   slower: 0.51x   kernel: 2.79ms  baseline: 5.44ms
seq_len: 2048   slower: 0.45x   kernel: 8.37ms  baseline: 18.67ms
seq_len: 4096   slower: 0.40x   kernel: 29.16ms baseline: 72.97ms
seq_len: 8192   slower: 0.38x   kernel: 108.68ms    baseline: 285.47ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.81ms  baseline: 0.98ms
seq_len: 256    slower: 0.90x   kernel: 0.88ms  baseline: 0.98ms
seq_len: 512    slower: 1.16x   kernel: 1.13ms  baseline: 0.97ms
seq_len: 1024   slower: 0.80x   kernel: 1.68ms  baseline: 2.10ms
seq_len: 2048   slower: 0.54x   kernel: 3.66ms  baseline: 6.81ms
seq_len: 4096   slower: 0.45x   kernel: 11.43ms baseline: 25.32ms
seq_len: 8192   slower: 0.41x   kernel: 40.58ms baseline: 99.14ms

Последовательности переменной длины (замаскировано до 25% токенов)

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.80x   kernel: 0.85ms  baseline: 1.07ms
seq_len: 256    slower: 1.07x   kernel: 1.15ms  baseline: 1.08ms
seq_len: 512    slower: 1.00x   kernel: 1.94ms  baseline: 1.94ms
seq_len: 1024   slower: 0.84x   kernel: 4.64ms  baseline: 5.55ms
seq_len: 2048   slower: 0.84x   kernel: 15.86ms baseline: 18.86ms
seq_len: 4096   slower: 0.76x   kernel: 55.19ms baseline: 72.47ms
seq_len: 8192   slower: 0.75x   kernel: 212.48ms    baseline: 282.71ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.80x   kernel: 0.83ms  baseline: 1.04ms
seq_len: 256    slower: 0.90x   kernel: 0.93ms  baseline: 1.03ms
seq_len: 512    slower: 1.18x   kernel: 1.22ms  baseline: 1.04ms
seq_len: 1024   slower: 1.10x   kernel: 2.40ms  baseline: 2.17ms
seq_len: 2048   slower: 0.89x   kernel: 6.27ms  baseline: 7.06ms
seq_len: 4096   slower: 0.82x   kernel: 21.19ms baseline: 25.95ms
seq_len: 8192   slower: 0.78x   kernel: 79.45ms baseline: 101.83ms

Обучение небольшого GPT на Enwik8

$ make train

Попробуйте длину последовательности 8192. Это будет медленно, но будет работать (нормальное внимание будет нарушено при > 2048, вы увидите это, если удалите флаг --use-cuda-kernel )

$ python train . py - - seq - len 8192 - - use - cuda - kernel 

Цитаты

 @article { Dao2022FlashAttentionFA ,
    title   = { FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness } ,
    author  = { Tri Dao and Daniel Y. Fu and Stefano Ermon and Atri Rudra and Christopher R'e } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2205.14135 }
}

 @misc { rabe2021selfattention ,
    title   = { Self-attention Does Not Need $O(n^2)$ Memory } , 
    author  = { Markus N. Rabe and Charles Staats } ,
    year    = { 2021 } ,
    eprint  = { 2112.05682 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @inproceedings { Henry2020QueryKeyNF ,
    title   = { Query-Key Normalization for Transformers } ,
    author  = { Alex Henry and Prudhvi Raj Dachapally and Shubham Vivek Pawar and Yuxuan Chen } ,
    booktitle = { FINDINGS } ,
    year    = { 2020 }
}

 @article { Wang2022DeepNetST ,
    title   = { DeepNet: Scaling Transformers to 1, 000 Layers } ,
    author  = { Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Dongdong Zhang and Furu Wei } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2203.00555 }
}