FlashAttention

SMP v2 支援FlashAttention核心，並可輕鬆將它們套用至 Hugging Face 變壓器模型的各種場景。請注意，如果您使用 v2.0 或更新版本的 FlashAttention 套件，SMP 會使用 FlashAttention v2；然而，Triton 快閃記憶體注意預設為 FlashAttention v1.x 中的快閃記憶體注意核心，因此僅在 v1 中受到支援。 FlashAttention

模塊（nn.Module）是定義模型注意層的低級 API。它應該在模型創建之後立即應用，例如從 AutoModelForCausalLM.from_config() API 中，以及在模型被轉換或使用 FSDP 包裝之前。

使用 FlashAttention 內核進行自我關注

下面的代碼片段演示了如何使用 SMP v2 提供的 torch.sagemaker.nn.attn.FlashSelfAttention API。

def new_attn(self, q, k, v, attention_mask=None, head_mask=None):
    return (
        self.flashmod((q, k, v), causal=True, cast_dtype=torch.bfloat16, layout="b h s d"),
        None,
    )

for layer in model.gpt_neox.layers:
    layer.attention.flash_mod = torch.sagemaker.nn.attn.FlashSelfAttention()
    layer.attention._attn = functools.partial(new_attn, layer.attention)

使用 FlashAttention 內核進行大量查詢注意

SMP v2 也支援針對群組查詢注意 (GQA) 的FlashAttention核心，並可輕鬆將它們套用至 Hugging Face 部變壓器模型的各種場景。與原始的注意力架構不同，GQA 同樣將查詢頭分成群組，同一組中的查詢頭共用相同的鍵和價值頭。因此，q 和 kv 頭分別傳遞給正向呼叫。注意：q 頭的數量需要被 kv 頭的數量整除。

使用示例 FlashGroupedQueryAttention

下面的代碼片段演示了如何使用 SMP v2 提供的 torch.sagemaker.nn.attn.FlashGroupedQueryAttention API。

from transformers.models.llama.modeling_llama import LlamaAttention
from torch.sagemaker.nn.attn import FlashGroupedQueryAttention

class LlamaFlashAttention(LlamaAttention):
    def __init__(self, config: LlamaConfig):
        super().__init__(config)

        self.flash_attn = FlashGroupedQueryAttention(
            attention_dropout_prob=0.0,
        )
        
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        ...
    ):
        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)
        ...
        kv = (key_states, value_states)
        attn_output = self.flash_attn(
            query_states,
            kv,
            attn_mask=attention_mask,
            causal=True,
            layout="b h s d",
        )
        ...
        attn_output = self.o_proj(attn_output)
        ...
        return attn_output

SMP 庫還提供了torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention使用低級別的 torch.sagemaker.nn.attn.FlashGroupedQueryAttention API。Hugging Face 變壓器有一個類似的實現，稱為 v LlamaFlashAttention24.36.0。下面的代碼片段演示了如何使用 SMP v2 LlamaFlashAttention API 或變形金剛 LlamaFlashAttention2 API 來替換現有駱駝模型的注意層。

from torch.sagemaker.nn.huggingface.llama_flashattn import LlamaFlashAttention
from transformers.models.llama.modeling_llama import LlamaFlashAttention2

flash_attn_class = LlamaFlashAttention # or flash_attn_class = LlamaFlashAttention2

attn_name = "self_attn"
for layer in model.model.layers:
    prev_layer = getattr(layer, attn_name)
    setattr(layer, attn_name, flash_attn_class(model.config))