本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。
FlashAttention
SMP v2 支援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
使用示例 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 LlamaFlashAttention2
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))