В области глубокого обучения механизм внимания является ключевым компонентом повышения производительности модели. Однако традиционные механизмы внимания потребляют много памяти и вычислительных ресурсов при обработке длинных последовательностей. Чтобы решить эту проблему, Три Дао и др. предложили FlashAttention — быстрый и эффективный механизм внимания. В этой статье будут представлены основные концепции и методы установки FlashAttention и его улучшенной версии FlashAttention-2и Пример. использования。

Введение в бумагу

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

• автор: Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
• Бумажная ссылка: arxiv.org/abs/2205.14135

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

• автор: Tri Dao
• Бумажная ссылка: flash2.pdf

Установка и возможности

Экологические требования

• CUDA: 11,6 и выше
• PyTorch: 1.12 и выше
• Операционная система: Linux (начиная с версии 2.3.2 были некоторые положительные отзывы о Windows, но компиляция для Windows все еще требует дополнительного тестирования)

Мы рекомендуем использовать контейнер PyTorch от Nvidia, который содержит все инструменты, необходимые для установки FlashAttention.

Этапы установки

1. Убедитесь, что PyTorch установлен.
2. Установитьpackaging：pip install packaging
3. Установитьninjaи убедитесь, что он работает правильно：ninja --version && echo $?должен вернуть код выхода0。если Не возвращено0,重新Установитьninja：pip uninstall -y ninja && pip install ninja

Установить с помощью pip

pip install flash-attn --no-build-isolation

Скомпилировать из исходников

python setup.py install

Контролируйте количество задач параллельной компиляции (для машин с объемом оперативной памяти менее 96 ГБ и несколькими ядрами ЦП)

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

Пример использования

FlashAttention в основном реализует масштабирование внимания скалярного произведения (softmax(Q @ K^T * softmax_scale) @ V). Ниже приведены основные функции использования FlashAttention:

from flash_attn import flash_attn_qkvpacked_func, flash_attn_func

# Когда К, K, Когда V помещен в тензор, используйте flash_attn_qkvpacked_func.
out = flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,
                                window_size=(-1, -1), alibi_slopes=None, deterministic=False)

# Используйте Q напрямую, K, Ви, используй flash_attn_func
out = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,
                      window_size=(-1, -1), alibi_slopes=None, deterministic=False)

Описание параметра

• qkv: (batch_size, seqlen, 3, nheads, headdim)тензор формата,Содержит Q, K, V
• dropout_p: с плавающей запятой, вероятность выпадения
• softmax_scale: float, коэффициент масштабирования QK^T перед softmax, значение по умолчанию — 1 / sqrt(headdim)
• causal: bool, применять ли маску причинного внимания (например, для авторегрессионного моделирования)
• window_size: (left, верно), если нет (-1, -1), то достигается локальное внимание скользящего окна
• alibi_slopes: (nheads,) или (batch_size, нхеды), fp32. Добавьте смещение (-alibi_slope) к оценке внимания запроса iи ключа j. * |i - j|)
• deterministic: bool, использовать ли детерминированную реализацию Обратного распространение ошибки (немного медленнее и используйте больше Память)

Производительность

эффект ускорения

FlashAttention на А100 80GB SXM5 эффект при использовании формата FP16/BF16 на графическом процессоре ускоренияследующее：

• Head Dimension: 64 или 128
• Hidden Dimension: 2048 (т.е. 32 или 16 голов)
• Sequence Length: 512, 1k, 2k, 4k, 8k, 16k
• Batch Size: 16k / seqlen

экономия памяти

FlashAttention позволяет существенно экономить память при обработке более длинных последовательностей. В отличие от использования памяти стандартным механизмом внимания, которое увеличивается квадратично с длиной последовательности, использование памяти FlashAttention увеличивается линейно. Когда длина последовательности составляет 2 КБ, память можно сохранять 10 раз, а когда длина последовательности составляет 4 КБ, память можно сохранять 20 раз.

Полный код модели и сценарий обучения

Была выпущена полная реализация модели GPT, а также предоставлены оптимизированные реализации других уровней (таких как MLP, LayerNorm, перекрестная энтропия, встраивание вращения). В целом скорость обучения увеличена в 3-5 раз по сравнению с базовыми реализациями (такими как реализация Huggingface), достигнув 225 терафлопс/сек на A100, что эквивалентно использованию 72% модельных флопс.

Журнал обновлений FlashAttention

2.0: Полностью переписан, в 2 раза быстрее.

FlashAttention был полностью переписан в версии 2.0 и теперь работает в два раза быстрее. В этом обновлении представлено несколько изменений и улучшений, включая изменения в именах некоторых функций и упрощенное использование, когда входные данные имеют одинаковую длину последовательности. FlashAttention-2 — это серия улучшений исходного алгоритма FlashAttention, направленных на повышение производительности вычислений на графических процессорах. В этой статье подробно обсуждаются алгоритм, параллелизм и рабочий режим FlashAttention-2. раздел Стратегия。

алгоритм

Ключевым моментом оптимизации FlashAttention-2 является сокращение операций с плавающей запятой при нематричном умножении (matmul) для полного использования выделенных вычислительных блоков графического процессора (таких как тензорные ядра на графических процессорах Nvidia), которые используются при обработке операций matmul. (особенно в формате FP16/BF16) производительность значительно оптимизирована. Цель этой оптимизации — максимизировать пропускную способность графического процессора за счет выполнения как можно большего количества операций matmul.

прямое распространение

1. существовать ПроволокаSoftmaxНавык：FlashAttention-2верносуществовать ПроволокаSoftmaxРасчеты были изменены,минимизировать не-matmulоперации с плавающей запятой：
   • избегать прохождения diag(ℓ(2))^-1 Измените масштаб двух элементов, выводящих обновления.
   • Поддерживайте «немасштабированную» версию O(2) и сохраняйте статистику ℓ(2).
   • Только когда цикл закончится, пройдите diag(ℓ(last))^-1 Масштабируйте последний O(last), чтобы получить правильный результат.
2. Максимизируйте количество флопов Matmul：чтобы максимизироватьGPUпроизводительность,FlashAttention-2 фокусируется на операциях matmul.,Потому что специализированные модули на современных графических процессорах (такие как Tensor Ядра) существуют, хорошо справляются с этими операциями. Возьмите Нвидиа A100 Графический процессор, например, FP16/BF16. Теоретическая пропускная способность matmul может достигать 312 Тфлопс/с, не матмул Пропускная способность FP32 всего 19,5 TFLOPs/s。поэтому,FlashAttention-2проходитьоптимизацияалгоритм,Максимально сведите к минимуму операции, не связанные с matmul.,Таким образом, сохраняется высокая пропускная способность и эффективность выполнения.
3. алгоритм Подробнее：FlashAttention-2изпрямое распространениепроходить以下步骤实现：
   • 𝐵𝑟 × 𝑑 блоков 𝑇𝑟, выходная матрица Oиlogsumexp𝐿 также соответственно разбивается на блоки.
   • существуют. Работа распределяется внутри каждого блока, чтобы максимально эффективно использовать ресурсы графического процессора.
   • Представлена ​​линия существования Softmax трюк,Масштабирование промежуточных результатов за счет эффективного управления и,Уменьшите ненужные вычислительные затраты.

Обратное распространение ошибки

Обратное распространение ошибок FlashAttention-2 похоже на FlashAttention, но с некоторыми изменениями:

• Используйте только построчный logsumexp 𝐿, а максимальное значение и индекс в netsoftmaxi.
• Используйте аналогичную стратегию разбиения на блоки для оптимизации вычислений и Память доступа для улучшения Обратного распространение ошибки эффективности и производительности.

FlashAttention-2существовать Параллелизмирабочий Аспектом раздела была глубокая оптимизация для достижения более высокой эффективности вычислений и производительности на существующем графическом процессоре. В этом разделе подробно обсуждается стратегия распараллеливания FlashAttention-2 рабочий. разделметод。

Параллелизм

прямое распространение

существоватьFlashAttention-2середина,прямое распространениеиз并行化Стратегияследующее：

1. Планирование блоков потоков：Используйте по одному на каждую голову внимания.нитьблоки для обработки,Всего есть партии size × number of головы нить блоки. Каждый блок запланирован для выполнения на потоковом мультипроцессоре (SM). Например, Нвидиа A100 Таких СМ на ГПУ имеется 108. Этот тип планирования предполагает большое количество блоков (например, ≥ 80) очень эффективен, поскольку может полностью использовать вычислительные ресурсы графического процессора.
2. Оптимизация для длинных последовательностей：длядлинная последовательность（通常意味着较小изbatch size或较少из头数）,Чтобы лучше использовать мультипроцессоры на графическом процессоре,FlashAttention-2 дополнительно распараллеливает измерение длины последовательности. В данном случае это существование значительно повышает производительность и эффективность.

Обратное распространение ошибки

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

• Планирование блоков потоков：Используйте по одному на блок столбцовнитьблоки для обработки。проходитьиспользовать原子加操作来существоватьдругойнить Общайтесь между блоками,обновить dQ,Это позволяет избежать конфликтов чтения и записи в общей Память.

рабочий раздел

прямое распространение

существоватьпрямое распространениесередина,FlashAttention-2 улучшает стратегию рабочих разделов,Избегает схемы «split-K» в FlashAttention.,В частности, включите:

• Сегментация Киева：FlashAttention-2ВоляQразделить на4индивидуальныйнитьпучок（warp）середина,в то же время сделатьKиVверно所有нитьпучок可访问。每индивидуальныйнитьпучок执行矩阵乘法以获取QK>из一部分,и умножим его на часть V,Получить соответствующий выходной фрагмент. Это улучшение уменьшает связь между пакетами.,Уменьшено количество операций чтения и записи для общей Память.,Это улучшает производительность.

Обратное распространение ошибки

существовать Обратное распространение ошибкисередина,Во избежание проблем синхронизации, вызванных схемой «сплит-К»,FlashAttention-2 выбирает подходящую стратегию разделения пакета.,Используйте оптимизацию для расчета эффективности доступа и Память.

переименование функции

Названия следующих функций были обновлены, чтобы отразить их обновленную функциональность:

• flash_attn_unpadded_func -> flash_attn_varlen_func
• flash_attn_unpadded_qkvpacked_func -> flash_attn_varlen_qkvpacked_func
• flash_attn_unpadded_kvpacked_func -> flash_attn_varlen_kvpacked_func

Если входные данные имеют одинаковую длину последовательности в одном пакете, будет проще и быстрее использовать следующую функцию:

• flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)
• flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)

2.1: Изменение поведения флагов причин

если seqlen_q != seqlen_k и causal=True,Тогда маска причины будет выровнена по правому нижнему углу матрицы внимания.,И нет верхнего левого угла.

Например, если seqlen_q = 2 и seqlen_k = 5,затем причинная маска (1 = Зарезервировано, 0 = маска) следующим образом:

версия v2.0:

1 0 0 0 0
1 1 0 0 0

версия v2.1:

1 1 1 1 0
1 1 1 1 1

если seqlen_q = 5 и seqlen_k = 2,Тогда маска причины будет следующей:

версия v2.0:

1 0
1 1
1 1
1 1
1 1

версия v2.1:

0 0
0 0
0 0
1 0
1 1

Если все замаскированные строки содержат нули, то выходные данные также будут нулями.

2.2: Оптимизация вывода

Оптимизируйте вывод (итеративное декодирование), когда длина последовательности запроса очень коротка (например, длина последовательности запроса = 1). Узким местом здесь является максимально быстрая загрузка KV-кэша, мы разделяем нагрузку по разным блокам потоков и используем отдельное ядро ​​для слияния результатов.

См. flash_attn_with_kvcache Функция (выполнить встраивание ротации, обновить кэш KV на месте).

Спасибо команде xformers, особенно Даниэлю Хазизе за сотрудничество.

2.3: Локальное (т.е. скользящее окно) внимание

Внедрите скользящее окно внимания (т. е. локальное внимание). Спасибо команде Mistral AI, особенно Тимоти Лакруа за его вклад. В модели Мистраль 7Б используются раздвижные окна.

2.4: ALiBi (внимание к линейному смещению), детерминированное обратное распространение ошибки

Внедрение ALiBi (Press et al., 2021). Спасибо Сангуну Чо из Kakao Brain за вклад.

Достичь определенности. Обратное распространение ошибок. Спасибо инженерам Meituan.

2.5: Пейджинговый KV-кеш

Поддерживает страничный кэш KV (т. е. PagedAttention). Спасибо @beginlner за его вклад.

Каталог кода flash_attn/modules/block.py

Интерпретация кода: Блок

В этом блоге мы объясним это шаг за шагом. Block код класса. Этот класс реализует общую блочную структуру, которая широко используется в таких моделях, как Transformer середина. Мы подробно представим процесс. реализации кодаи所用приезжатьиз Теоретическая основа。

Теоретическая основа

В архитектуре Transformer самыми основными компонентами являются уровни (блоки) кодера и декодера. Каждый слой обычно включает в себя следующие части:

1. Многоголовый механизм самообслуживания：используется длявычислить每индивидуальный词верно其他词из注意力权重。
2. нейронная сеть прямого распространения：верно每индивидуальный词из表示进行非линейное преобразование。
3. Остаточная связь и нормализация слоев：Для стабильных тренировок,Добавлены остаточные соединения и нормализация слоев.

Существует две общие структуры слоев:

• Преднормальная структура：нормализация слоясуществовать Основные операции（внимание илинейронная сеть прямого распространения）применялся раньше。
• Постнормальная структура：нормализация слоясуществовать Основные операции之后应用。

Процесс реализации кода

class Block(nn.Module):
    def __init__(
        self,
        dim,
        mixer_cls=None,
        mlp_cls=None,
        norm_cls=nn.LayerNorm,
        dropout_cls=nn.Dropout,
        prenorm=True,
        resid_dropout1=0.0,
        resid_dropout2=0.0,
        drop_path1=0.0,
        drop_path2=0.0,
        fused_dropout_add_ln=False,
        return_residual=False,
        residual_in_fp32=False,
        sequence_parallel=False,
        mark_shared_params=False,
    ):

Этот код определяет Block Конструктор класса. Вот объяснение параметров:

• dim：входитьи输出из维度。
• mixer_cls：используется длявычислить注意力из类。
• mlp_cls：используется длянейронная сеть прямого распространенияиз类。
• norm_cls：используется длянормализация слояиз类。
• dropout_cls：используется дляDropoutиз类。
• prenorm：Стоит ли использовать Преднормальная структура。
• resid_dropout1 и resid_dropout2：残差连接изDropoutСтавка。
• drop_path1 и drop_path2：используется дляStochastic Параметр глубины.
• fused_dropout_add_ln：Стоит ли интегрироватьDropout、Операция AddиLayerNorm.
• return_residual：лисуществовать每индивидуальный子слой返回残差。
• residual_in_fp32：Стоит ли использоватьFP32Точность экономит остатки。
• sequence_parallel：ли并行处理序列。
• mark_shared_params：ли标记共享параметр。

        super().__init__()
        self.prenorm = prenorm
        self.fused_dropout_add_ln = fused_dropout_add_ln
        self.return_residual = return_residual
        self.residual_in_fp32 = residual_in_fp32
        if self.residual_in_fp32:
            assert self.prenorm, "residual_in_fp32 is only compatible with prenorm=True"
        if mixer_cls is None:
            mixer_cls = partial(MHA, num_heads=dim // 64)
        if mlp_cls is None:
            mlp_cls = partial(Mlp, hidden_features=4 * dim)
        self.mixer = mixer_cls(dim)
        self.dropout1 = dropout_cls(resid_dropout1)
        self.drop_path1 = StochasticDepth(drop_path1, mode="row")
        self.norm1 = norm_cls(dim)
        self.mlp = mlp_cls(dim)
        if not isinstance(self.mlp, nn.Identity):
            self.dropout2 = dropout_cls(resid_dropout2)
            self.drop_path2 = StochasticDepth(drop_path2, mode="row")
            self.norm2 = norm_cls(dim)

В конструкторе каждый параметр инициализируется первым. в соответствии с prenorm、fused_dropout_add_ln и т. д. флаги устанавливают некоторые утверждения и значения по умолчанию. если не предусмотрено mixer_cls и mlp_cls,По умолчанию используется механизм многоголового внимания (MHA).

Далее инициализируйте mixer、dropout、StochasticDepth и norm слой.

        if self.fused_dropout_add_ln:
            assert layer_norm_fn is not None, "Triton is not installed"
            assert isinstance(self.norm1, (nn.LayerNorm, RMSNorm)) and isinstance(
                self.dropout1, nn.Dropout
            )

        if sequence_parallel:
            for p in self.norm1.parameters():
                p._sequence_parallel = True
            if hasattr(self, "norm2"):
                for p in self.norm2.parameters():
                    p._sequence_parallel = True

        if mark_shared_params:
            for p in self.norm1.parameters():
                p._shared_params = True
            if hasattr(self, "norm2"):
                for p in self.norm2.parameters():
                    p._shared_params = True

Этот код обрабатывает fused_dropout_add_ln、sequence_parallel и mark_shared_params ситуация. если включено fused_dropout_add_ln,затем убедитесь, что Triton установлен,и norm1 и dropout1 является допустимым типом. если включено sequence_parallel,но Воля norm1 и norm2 Параметры помечены как требующие параллелизма последовательностей. если включено mark_shared_params,но Воля这些параметр Отметить как общийпараметр。

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)

    def forward(
        self,
        hidden_states: Tensor,
        residual: Optional[Tensor] = None,
        mixer_subset=None,
        mixer_kwargs=None,
    ):

Определены два метода:

• allocate_inference_cache：Выделить кеш для фазы вывода。
• forward：прямое передачафункция, обработка ввода hidden_states и residual。

        if self.prenorm:
            if not self.fused_dropout_add_ln:
                dropped = self.drop_path1(self.dropout1(hidden_states))
                residual = (dropped + residual) if residual is not None else dropped
                hidden_states = self.norm1(residual.to(dtype=self.norm1.weight.dtype))
                if self.residual_in_fp32:
                    residual = residual.to(torch.float32)
            else:
                if self.drop_path1.p == 0 or not self.training:
                    rowscale1 = None
                else:
                    rowscale1 = self.drop_path1(
                        torch.ones(
                            hidden_states.shape[:-1],
                            device=hidden_states.device,
                            dtype=hidden_states.dtype,
                        )
                    )
                hidden_states, residual = layer_norm_fn(
                    hidden_states,
                    self.norm1.weight,
                    self.norm1.bias,
                    residual=residual,
                    eps=self.norm1.eps,
                    dropout_p=self.dropout1.p if self.training else 0.0,
                    rowscale=rowscale1,
                    prenorm=True,
                    residual_in_fp32=self.residual_in_fp32,
                    is_rms_norm=isinstance(self.norm1, RMSNorm)
                )
            if mixer_kwargs is None:
                mixer_kwargs = {}
            if mixer_subset is not None:
                mixer_kwargs["mixer_subset"] = mixer_subset
            hidden_states = self.mixer(hidden_states, **mixer_kwargs)
            if mixer_subset is not None:
                residual = residual[:, mixer_subset]
            if not isinstance(self.mlp, nn.Identity):
                if not self.fused_dropout_add_ln:
                    dropped = self.drop_path2(self.dropout2(hidden_states))
                    residual = (dropped + residual) if residual is not None else dropped
                    hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
                    if self.residual_in_fp32:
                        residual = residual.to(torch.float32)
                else:
                    if self.drop_path2.p == 0 or not self.training:
                        rowscale2 = None
                    else:
                        rowscale2 = self.drop_path2(
                            torch.ones(
                                hidden_states.shape[:-1],
                                device=hidden_states.device,
                                dtype=hidden_states.dtype,
                            )
                        )
                    hidden_states, residual = layer_norm_fn(
                        hidden_states,
                        self.norm2.weight,
                        self.norm2.bias,
                        residual=residual,
                        eps=self.norm2.eps,
                        dropout_p=self.dropout2.p if self.training else 0.0,
                        rowscale=rowscale2,
                        prenorm=True,
                        residual_in_fp32=self.residual_in_fp32,
                        is_rms_norm=isinstance(self.norm2, RMSNorm)
                    )
                hidden_states = self.mlp(hidden_states)
            return hidden_states, residual

еслииспользовать prenorm,но首先处理 dropout и residual,затем подайте заявку norm1。в соответствии с fused_dropout_add_ln настройки, чтобы выбрать, следует ли объединять эти операции. позвони после mixer слой (обычно многоголовочный механизм внимания). Наконец, если mlp нет Identity,но进行类似из操作处理 mlp слой.

        else:
            assert residual is None
            mixer_out = self.mixer(
                hidden_states, **(mixer_kwargs if mixer_kwargs is not None else {})
            )
            if self.return_residual:  # mixer out is actually a pair here
                mixer_out, hidden_states = mixer_out
            if not self

.fused_dropout_add_ln:
                hidden_states = self.norm1(
                    (self.drop_path1(self.dropout1(mixer_out)) + hidden_states).to(
                        dtype=self.norm1.weight.dtype
                    )
                )
            else:
                if self.drop_path1.p == 0 or not self.training:
                    rowscale1 = None
                else:
                    rowscale1 = self.drop_path1(
                        torch.ones(
                            mixer_out.shape[:-1], device=mixer_out.device, dtype=mixer_out.dtype
                        )
                    )
                hidden_states = layer_norm_fn(
                    mixer_out,
                    self.norm1.weight,
                    self.norm1.bias,
                    residual=hidden_states,
                    eps=self.norm1.eps,
                    dropout_p=self.dropout1.p if self.training else 0.0,
                    rowscale=rowscale1,
                    prenorm=False,
                    is_rms_norm=isinstance(self.norm1, RMSNorm)
                )
            if not isinstance(self.mlp, nn.Identity):
                mlp_out = self.mlp(hidden_states)
                if self.return_residual:  # mlp out is actually a pair here
                    mlp_out, hidden_states = mlp_out
                if not self.fused_dropout_add_ln:
                    hidden_states = self.norm2(
                        (self.drop_path2(self.dropout2(mlp_out)) + hidden_states).to(
                            dtype=self.norm2.weight.dtype
                        )
                    )
                else:
                    if self.drop_path2.p == 0 or not self.training:
                        rowscale2 = None
                    else:
                        rowscale2 = self.drop_path2(
                            torch.ones(
                                mlp_out.shape[:-1], device=mlp_out.device, dtype=mlp_out.dtype
                            )
                        )
                    hidden_states = layer_norm_fn(
                        mlp_out,
                        self.norm2.weight,
                        self.norm2.bias,
                        residual=hidden_states,
                        eps=self.norm2.eps,
                        dropout_p=self.dropout2.p if self.training else 0.0,
                        rowscale=rowscale2,
                        prenorm=False,
                        is_rms_norm=isinstance(self.norm2, RMSNorm)
                    )
            return hidden_states

для postnorm Структура и процесс обработки аналогичны, но layer norm Применительно к основным операциям (внимание инейронная сеть прямого распространение) после. Вот residual изначально настроен на None。处理 mixer слой, обработанный позже mlp слой.

через этот код,Block Классы могут гибко поддерживать prenorm и postnorm структура,И различные Дропауты, Остаточная связь и нормализация слоевиз组合。Это делает этосуществовать实现другой类型изTransformerОчень эффективен в архитектуреи Универсальный。

Блог по интерпретации кода: ParallelBlock

В этом блоге мы объясним это шаг за шагом. ParallelBlock код класса. Этот класс реализует блок параллельного внимания (микшера) и MLP, аналогичный структуре модели GPT-J, GPT-NeoX и PaLM.

Теоретическая основа

ParallelBlock Класс имеет немного другую структуру, чем обычный блок Transformer. Традиционный блок Transformer обычно имеет следующую структуру: Слой. Norm (LN) -> Multi-Head Attention (MHA) / MLP -> Dropout -> Добавлять. и ParallelBlock Структура: Dropout -> Add -> LN -> MHA / МЛП. Преимущество этой структуры в том, что она может интегрировать операции dropout и addLayerNorm для повышения производительности.

Процесс реализации кода

class ParallelBlock(nn.Module):
    """The attention (mixer) and MLP blocks are done in parallel, similar to GPT-J, GPT-NeoX,
    and PaLM.
    """

    def __init__(
        self,
        dim,
        mixer_cls=None,
        mlp_cls=None,
        norm_cls=nn.LayerNorm,
        dropout_cls=nn.Dropout,
        resid_dropout1=0.0,
        resid_dropout2=0.0,
        tied_norm=False,
        fused_dropout_add_ln=False,
        residual_in_fp32=False,
        sequence_parallel=False,
        mark_shared_params=False,
    ):
        super().__init__()
        self.tied_norm = tied_norm
        self.fused_dropout_add_ln = fused_dropout_add_ln
        self.residual_in_fp32 = residual_in_fp32
        if mixer_cls is None:
            mixer_cls = partial(MHA, num_heads=dim // 64)
        if mlp_cls is None:
            mlp_cls = partial(Mlp, hidden_features=4 * dim)
        self.mixer = mixer_cls(dim)
        self.dropout1 = dropout_cls(resid_dropout1)
        self.norm1 = norm_cls(dim)
        self.mlp = mlp_cls(dim)
        self.dropout2 = dropout_cls(resid_dropout2)
        if not self.tied_norm:
            self.norm2 = norm_cls(dim)

        if self.fused_dropout_add_ln:
            assert layer_norm_fn is not None, "Triton is not installed"
            assert isinstance(self.norm1, (nn.LayerNorm, RMSNorm)) and isinstance(
                self.dropout1, nn.Dropout
            )

        if sequence_parallel:
            for p in self.norm1.parameters():
                p._sequence_parallel = True
            if hasattr(self, "norm2"):
                for p in self.norm2.parameters():
                    p._sequence_parallel = True

        if mark_shared_params:
            for p in self.norm1.parameters():
                p._shared_params = True
            if hasattr(self, "norm2"):
                for p in self.norm2.parameters():
                    p._shared_params = True

В конструкторе каждый параметр инициализируется первым. в соответствии с tied_norm、fused_dropout_add_ln и т. д. флаги устанавливают некоторые утверждения и значения по умолчанию. если не предусмотрено mixer_cls и mlp_cls,По умолчанию используется механизм многоголового внимания (MHA).Инициализировано mixer、dropout、norm и mlp слой,并в соответствии с条件设置了нормализация слояпараметриз并行序列иподелиться флагом。

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)

    def forward(
        self,
        hidden_states1: Tensor,
        hidden_states2: Optional[Tensor] = None,
        residual: Optional[Tensor] = None,
        mixer_kwargs=None,
    ):
        r"""Pass the input through the encoder layer.

        Args:
            hidden_states1: the output of the previous attention (mixer) or embedding layer.
            hidden_states2: the output of the previous MLP layer (if None, will use hidden_states1).
            residual.
        """
        if not self.fused_dropout_add_ln:
            dropped1 = self.dropout1(hidden_states1)
            if hidden_states2 is not None:
                dropped2 = self.dropout2(hidden_states2)
                residual = (
                    (residual + dropped1 + dropped2)
                    if residual is not None
                    else dropped1 + dropped2
                )
            else:
                residual = (residual + dropped1) if residual is not None else dropped1
            hidden_states1 = self.norm1(residual.to(dtype=self.norm1.weight.dtype))
            hidden_states2 = (
                self.norm2(residual.to(dtype=self.norm2.weight.dtype))
                if not self.tied_norm
                else hidden_states1
            )
            if self.residual_in_fp32:
                residual = residual.to(torch.float32)
        else:
            weight2, bias2 = (
                (self.norm2.weight, self.norm2.bias) if not self.tied_norm else (None, None)
            )
            hidden_states1, *rest, residual = layer_norm_fn(
                hidden_states1,
                self.norm1.weight,
                self.norm1.bias,
                residual=residual,
                x1=hidden_states2,
                weight1=weight2,
                bias1=bias2,
                eps=self.norm1.eps,
                dropout_p=self.dropout1.p if self.training else 0.0,
                prenorm=True,
                residual_in_fp32=self.residual_in_fp32,
                is_rms_norm=isinstance(self.norm1, RMSNorm)
            )
            if self.tied_norm:
                hidden_states2 = hidden_states1
            else:
                hidden_states2, = rest

        if mixer_kwargs is None:
            mixer_kwargs = {}
        hidden_states1 = self.mixer(hidden_states1, **mixer_kwargs)
        hidden_states2 = self.mlp(hidden_states2)
        return hidden_states1, hidden_states2, residual

существовать forward метод, согласно fused_dropout_add_ln Настройки, выберите, интегрировать ли операции dropout и addиLayerNorm. в соответствии с hidden_states2 Это None,决定ли添加dropoutприезжать residual середина. затем подайте заявку norm1 и norm2,и разобраться с остатками. вызов mixer и mlp слой и, наконец, вернитесь hidden_states1、hidden_states2 и residual。

через этот код,ParallelBlock Класс реализует блочную структуру параллельного внимания и MLP, предоставляя эффективный метод оптимизации производительности модели.

Каталог кода flash_attn/modules/mha.py

Интерпретация кода: FlashSelfAttention

В этом блоге мы объясним это шаг за шагом. FlashSelfAttention код класса. Этот класс реализует Softmax из Многоголовый механизм самообслуживания. Мы подробно представим процесс реализации кодаи所用приезжатьиз Теоретическая основа。

Теоретическая основа

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

1. линейное преобразование：Волявходить向量проходить Проволока性слой分别映射приезжать查询、Вектор ключа и значения.
2. Рассчитать вес внимания：использовать查询и键向量из点积Рассчитать вес внимание, обычно масштабируемое и нормализованное с помощью функции Softmax.
3. Взвешенный поиск：Воля归一化后из注意力权重与值向量相乘,Получите выходной вектор.

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

Процесс реализации кода

class FlashSelfAttention(nn.Module):
    """Реализация масштабируемого точечного произведения с помощью Softmax.
    параметр
    ---------
        softmax_scale: используется дляSoftmax注意力из温度параметр。
                      (По умолчанию: 1/sqrt(d_keys), где d_keysсуществовать вычисляется во время выполнения)
        attention_dropout: Коэффициент отсева применительно к вниманию
                           (По умолчанию: 0,0)
    """

    def __init__(
        self,
        causal=False,
        softmax_scale=None,
        attention_dropout=0.0,
        window_size=(-1, -1),
        alibi_slopes=None,
        deterministic=False,
    ):
        super().__init__()
        assert flash_attn_varlen_qkvpacked_func is not None, «FlashAttention не установлен»
        assert flash_attn_qkvpacked_func is not None, «FlashAttention не установлен»
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.drop = nn.Dropout(attention_dropout)
        self.register_buffer("alibi_slopes", alibi_slopes, persistent=False)
        self.window_size = window_size
        self.deterministic = deterministic

Этот код определяет FlashSelfAttention Конструктор класса. Вот объяснение параметров:

• causal：Этопричинное внимание,То есть, следует ли учитывать временной порядок последовательности.
• softmax_scale：Softmaxиз温度параметр,Используется для масштабирования результатов скалярного произведения.
• attention_dropout：注意力机制серединаизDropoutСтавка。
• window_size：локальный размер окна。
• alibi_slopes：используется для调整注意力偏置из斜Ставка。
• deterministic：Стоит ли использовать确定性操作。

В конструкторе сначала убедитесь, что требуемая функция FlashAttention установлена ​​с помощью утверждений, затем инициализируйте каждый параметр и передайте register_buffer Зарегистрируйте параметры, которые не требуют градиентов.

    def forward(self, qkv, causal=None, cu_seqlens=None, max_seqlen=None):
        """Достижение многоголового внимания Softmax.
        параметр
        ---------
            qkv: Тензор, содержащий запрос, ключи и значения.
                Если cu_seqlens имеет значение Noneиmax_seqlen имеет значение None, то форма qkv равна (B, S, 3, H, D)。
                если cu_seqlens не Noneиmax_seqlen не None, то форма qkv равна (всего, 3, H, D),
                где total — общая длина последовательности в пакете.
            causal: если пройдет, переопределит self.causal
            cu_seqlens: (batch_size + 1,) Тензор формы типа torch.int32. Совокупная длина последовательностей в пакете, используемая для индексации в qkvсередина.
            max_seqlen: инт. Максимальная длина последовательности в пакете.
        возвращаться:
        --------
            out: еслиcu_seqlens не Noneиmax_seqlen не None, то форма (всего, H, D),
                В противном случае (В, S, H, D)。
        """
        assert qkv.dtype in [torch.float16, torch.bfloat16]
        assert qkv.is_cuda
        causal = self.causal if causal is None else causal
        unpadded = cu_seqlens is not None
        if self.alibi_slopes is not None:
            self.alibi_slopes = self.alibi_slopes.to(torch.float32)
        if unpadded:
            assert cu_seqlens.dtype == torch.int32
            assert max_seqlen is not None
            assert isinstance(max_seqlen, int)
            return flash_attn_varlen_qkvpacked_func(
                qkv,
                cu_seqlens,
                max_seqlen,
                self.drop.p if self.training else 0.0,
                softmax_scale=self.softmax_scale,
                causal=causal,
                alibi_slopes=self.alibi_slopes,
                window_size=self.window_size,
                deterministic=self.deterministic,
            )
        else:
            return flash_attn_qkvpacked_func(
                qkv,
                self.drop.p if self.training else 0.0,
                softmax_scale=self.softmax_scale,
                causal=causal,
                alibi_slopes=self.alibi_slopes,
                window_size=self.window_size,
                deterministic=self.deterministic,
            )

Этот код реализует forward метод,Прямо сейчаспрямое распространениепроцесс。Ниже приводитсяпараметриз解释：

• qkv：Тензор, содержащий запрос, ключи и значения.
• causal：еслипередача,Воля覆盖 self.causal。
• cu_seqlens：批次середина序列из累计长度,используется для索引приезжать qkv середина.
• max_seqlen：批次середина最大序列长度。

существоватьпрямое процесс отправки, первая проверка qkv Тип данных и тип устройства. Тогда в зависимости от того, есть ли cu_seqlens чтобы определить, следует ли использовать незаполненные последовательности. если использовать незаполненную последовательность, то передать flash_attn_varlen_qkvpacked_func Функция вычисляет внимание, иначе проходит; flash_attn_qkvpacked_func Функция вычисляет внимание.

Краткое описание кода

FlashSelfAttention Класс реализует механизм самообслуживания с несколькими головками с помощью Softmax. Его основные этапы включают в себя:

1. Параметры инициализации.
2. существоватьпрямое распространениепроцесссередина,в соответствии свходить张量из形状ипараметр选择适当из Функция вычисляет внимание.

Через приведенный выше код,Мы можем эффективно вычислять механизмы самообслуживания,И существовать, когда это необходимо, применять причинное внимание Dropout.

Блог по интерпретации кода: FlashCrossAttention

В этом блоге мы объясним это шаг за шагом. FlashCrossAttention код класса. Этот класс реализует Softmax Механизм перекрестного внимания масштабированного скалярного произведения. Мы подробно представим процесс реализации кодаи所用приезжатьиз Теоретическая основа。

Теоретическая основа

Механизм перекрестного внимания аналогичен механизму самообслуживания, но использует разные запросы, ключи и значения из разных последовательностей. Его основные этапы включают в себя:

1. линейное преобразование：Воля查询、Векторы ключей и значений отображаются через линейные слои.
2. Рассчитать вес внимания：использовать查询и键向量из点积Рассчитать вес внимание, обычно масштабируется и проходит Softmax Нормализация функций.
3. Взвешенный поиск：Воля归一化后из注意力权重与值向量相乘,Получите выходной вектор.

Существование перекрестного внимания имеет широкое применение во многих задачах, таких как архитектура кодера-декодера в машинном переводе.

Процесс реализации кода

class FlashCrossAttention(nn.Module):
    """Реализация масштабируемого точечного произведения с помощью Softmax.
    параметр
    ---------
        softmax_scale: используется дляSoftmax注意力из温度параметр。
                      (По умолчанию: 1/sqrt(d_keys), где d_keysсуществовать вычисляется во время выполнения)
        attention_dropout: Коэффициент отсева применительно к вниманию
                           (По умолчанию: 0,0)
    """

    def __init__(
        self,
        causal=False,
        softmax_scale=None,
        attention_dropout=0.0,
        alibi_slopes=None,
        window_size=(-1, -1),
        deterministic=False,
    ):
        super().__init__()
        assert flash_attn_varlen_kvpacked_func is not None, «FlashAttention не установлен»
        assert flash_attn_kvpacked_func is not None, «FlashAttention не установлен»
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.drop = nn.Dropout(attention_dropout)
        self.register_buffer("alibi_slopes", alibi_slopes, persistent=False)
        self.window_size = window_size
        self.deterministic = deterministic

Этот код определяет FlashCrossAttention Конструктор класса. Вот объяснение параметров:

• causal：Этопричинное внимание,То есть, следует ли учитывать временной порядок последовательности.
• softmax_scale：Softmaxиз温度параметр,Используется для масштабирования результатов скалярного произведения.
• attention_dropout：注意力机制серединаизDropoutСтавка。
• window_size：локальный размер окна。
• alibi_slopes：используется для调整注意力偏置из斜Ставка。
• deterministic：Стоит ли использовать确定性操作。

В конструкторе сначала убедитесь, что требуемая функция FlashAttention установлена ​​с помощью утверждений, затем инициализируйте каждый параметр и передайте register_buffer Зарегистрируйте параметры, которые не требуют градиентов.

    def forward(
        self,
        q,
        kv,
        causal=None,
        cu_seqlens=None,
        max_seqlen=None,
        cu_seqlens_k=None,
        max_seqlen_k=None,
    ):
        """Достижение многоголового внимания Softmax.
        параметр
        ---------
            q: Тензор, содержащий запрос. Форма (B, Sq, H, D)
            kv: Тензор, содержащий ключи и значения. Форма (B, Sk, 2, H_k, D)
            causal: если пройдет, переопределит self.causal
            cu_seqlens: (batch_size + 1,) Тензор формы типа torch.int32. Совокупная длина последовательностей в пакете, используемая для индексации в q середина.
            max_seqlen: инт. В пакетном режиме q максимальная длина последовательности.
            cu_seqlens_k: (batch_size + 1,) Тензор формы типа torch.int32. Совокупная длина последовательностей в пакете, используемая для индексации в kv середина.
            max_seqlen_k: инт. В пакетном режиме k и v максимальная длина последовательности.
        """
        assert q.dtype in [torch.float16, torch.bfloat16]
        assert q.is_cuda and kv.is_cuda
        causal = self.causal if causal is None else causal
        unpadded = cu_seqlens is not None
        if self.alibi_slopes is not None:
            self.alibi_slopes = self.alibi_slopes.to(torch.float32)
        if unpadded:
            assert cu_seqlens.dtype == torch.int32
            assert max_seqlen is not None
            assert isinstance(max_seqlen, int)
            assert cu_seqlens_k is not None
            assert cu_seqlens_k.dtype == torch.int32
            assert max_seqlen_k is not None
            assert isinstance(max_seqlen, int)
            return flash_attn_varlen_kvpacked_func(
                q,
                kv,
                cu_seqlens,
                cu_seqlens_k,
                max_seqlen,
                max_seqlen_k,
                self.drop.p if self.training else 0.0,
                softmax_scale=self.softmax_scale,
                causal=causal,
                alibi_slopes=self.alibi_slopes,
                window_size=self.window_size,
                deterministic=self.deterministic,
            )
        else:
            batch_size, seqlen_q = q.shape[0], q.shape[1]
            seqlen_k = kv.shape[1]
            assert kv.shape[0] == batch_size and kv.shape[4] == q.shape[3]
            return flash_attn_kvpacked_func(
                q,
                kv,
                self.drop.p if self.training else 0.0,
                causal=causal,
                softmax_scale=self.softmax_scale,
                alibi_slopes=self.alibi_slopes,
                window_size=self.window_size,
                deterministic=self.deterministic,
            )

Этот код реализует forward метод,Прямо сейчаспрямое распространениепроцесс。Ниже приводитсяпараметриз解释：

• q：包含查询из张量,Форма (B, Sq, H, D)。
• kv：Содержит ключии值из张量,Форма (B, Sk, 2, H_k, D)。
• causal：еслипередача,Воля覆盖 self.causal。
• cu_seqlens：批次середина序列из累计长度,используется для索引приезжать q середина.
• max_seqlen：批次середина q максимальная длина последовательности.
• cu_seqlens_k：批次середина序列из累计长度,используется для索引приезжать kv середина.
• max_seqlen_k：批次середина kv максимальная длина последовательности.

существоватьпрямое процесс отправки, первая проверка q и kv Тип данных и тип устройства. Тогда в зависимости от того, есть ли cu_seqlens чтобы определить, следует ли использовать незаполненные последовательности. если использовать незаполненную последовательность, то передать flash_attn_varlen_kvpacked_func Функция вычисляет внимание, иначе проходит; flash_attn_kvpacked_func Функция вычисляет внимание.

Краткое описание кода

FlashCrossAttention Класс реализует механизм перекрестного внимания с несколькими головками с помощью Softmax. Его основные этапы включают в себя:

1. Параметры инициализации.
2. существоватьпрямое распространениепроцесссередина,в соответствии свходить张量из形状ипараметр选择适当из Функция вычисляет внимание.

Через приведенный выше код,Мы можем эффективно вычислять механизмы перекрестного внимания,И существовать, когда это необходимо, применять причинное внимание Dropout.

Блог по интерпретации кода: SelfAttention

В этом блоге мы объясним это шаг за шагом. SelfAttention код класса. Этот класс реализует Softmax Механизм самообслуживания скалярного произведения масштабирования. Мы подробно представим процесс реализации кодаи所用приезжатьиз Теоретическая основа。

Теоретическая основа

Механизм самообслуживания является ядром архитектуры Transformer, который позволяет модели учитывать все остальные слова в последовательности при вычислении представления каждого слова. Его основные этапы включают в себя:

1. линейное преобразование：Воля查询（query）、Векторы ключей и значений отображаются через линейные слои.
2. Рассчитать вес внимания：использовать查询и键向量из点积Рассчитать вес внимание, обычно масштабируется и проходитSoftmaxНормализация функций.
3. Взвешенный поиск：Воля归一化后из注意力权重与值向量相乘,Получите выходной вектор.

Процесс реализации кода

class SelfAttention(nn.Module):
    """Реализация масштабируемого точечного произведения с помощью Softmax.
    параметр
    ---------
        softmax_scale: используется дляSoftmax注意力из温度параметр。
                      (По умолчанию: 1/sqrt(d_keys), где d_keysсуществовать вычисляется во время выполнения)
        attention_dropout: Коэффициент отсева применительно к вниманию
                           (По умолчанию: 0,0)
    """

    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0):
        super().__init__()
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.drop = nn.Dropout(attention_dropout)

Этот код определяет SelfAttention Конструктор класса. Вот объяснение параметров:

• causal：Этопричинное внимание,То есть, следует ли учитывать временной порядок последовательности.
• softmax_scale：Softmaxиз温度параметр,Используется для масштабирования результатов скалярного произведения.
• attention_dropout：注意力机制серединаизDropoutСтавка。

Флаг причинного внимания инициализируется в конструкции Function、SoftmaxУвеличитьпараметриDropoutслой.

    def forward(self, qkv, causal=None, key_padding_mask=None):
        """Достижение многоголового внимания Softmax.
        параметр
        ---------
            qkv: Тензор, содержащий запрос, ключи и значения.Форма (B, S, 3, H, D)
            causal: если пройдет, переопределит self.causal
            key_padding_mask: Логическая маска, используемая для маскировки весов внимания. Правда означает сохранять, Ложь означает маскировать. Форма (B, S)
        """
        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
        causal = self.causal if causal is None else causal
        q, k, v = qkv.unbind(dim=2)
        softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])

существоватьпрямое В процессе распространения сначала определяются размер партии и длина последовательности. если пройдено causal параметры, переопределить self.causal。然后Воля qkv 张量沿第三индивидуальный维度进行拆分,Получите запрос, ключ и значение. наконец,вычислитьSoftmaxиз Увеличитьпараметр。

        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)

Эта строка кода использует соглашение Эйнштейна (einsum) для вычисления скалярного произведения ключа запроса.,并进行Увеличить。scores Форма тензора (B, H, T, S),Представляет сходство каждого запроса со всеми ключами.

        if key_padding_mask is not None:
            padding_mask = torch.full(
                (batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device
            )
            padding_mask.masked_fill_(key_padding_mask, 0.0)
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")

если предлагает key_padding_mask,Тогда оценка внимания маскируется. Сначала создайте маску заполнения,Устанавливает очень маленькое значение в указанной позиции (например, -10000,0). Затем,Примените эту маску к показателю внимания,Маскируйте значения, которые не нужно учитывать.

        if causal:
            # "triu_tril_cuda_template" not implemented for 'BFloat16'
            # So we have to construct the mask in float
            causal_mask = torch.triu(
                torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1
            )
            # TD [2022-09-30]: Adding is faster than masked_fill_ (idk why, just better kernel I guess)
            scores = scores + causal_mask.to(dtype=scores.dtype)

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

        attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
        attention_drop = self.drop(attention)
        output = torch.einsum("bhts,bshd->bthd", attention_drop, v)
        return output

Рассчитайте Softmax, чтобы получить вес внимания,Затем примените Dropout. наконец,Взвешенная оценка ценностей с использованием весов внимания,Получите окончательный результат.

Краткое описание кода

SelfAttention Класс реализует механизм самообслуживания с несколькими головками с помощью Softmax. Его основные этапы включают в себя:

1. Параметры инициализации.
2. В процессе существования прямого распространения оценка внимания рассчитывается на основе входного тензора.
3. Примените маску заполнения клавиш и причинную маску.
4. Рассчитайте Softmax, чтобы получить вес внимания, и примените Dropout.
5. Взвешенная оценка ценностей с использованием весов внимания,Получите окончательный результат.

в заключение

FlashAttention и его улучшенная версия FlashAttention-2 обеспечивают значительную скорость применения механизма внимания в глубоком обучении, делая обработку данных длинных последовательностей более эффективной. Я надеюсь, что эта статья поможет вам понять и использовать FlashAttention.

Если у вас есть какие-либо вопросы или предложения по FlashAttention, добро пожаловать на GitHub. Свяжитесь с нами.

Справочная ссылка:

• FlashAttention Бумага
• Бумага FlashAttention-2