Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

FlashAttention

SMP v2 supporta i FlashAttention kernel e ne semplifica l'applicazione a vari scenari per i modelli Hugging Face Transformer. Nota che se usate il FlashAttention pacchetto v2.0 o successivo, SMP usa FlashAttention v2; tuttavia, l'attenzione flash di Triton per impostazione predefinita è il flash attention kernel della versione 1.x, quindi è supportato esclusivamente nella v1. FlashAttention FlashAttention

Il modulo (nn.Module) è un’API di basso livello che definisce i livelli di attenzione di un modello. Dovrebbe essere applicato subito dopo la creazione del modello, ad esempio dall’API AutoModelForCausalLM.from_config(), e prima che si esegua la trasformazione o il wrapping del modello con FSDP.

FlashAttention Usate i kernel per concentrarvi su voi stessi

Il frammento di codice riportato di seguito mostra come utilizzare l’API torch.sagemaker.nn.attn. FlashSelfAttention fornita da SMP v2.

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)

Usa i FlashAttention kernel per attirare l'attenzione sulle interrogazioni raggruppate

SMP v2 supporta anche i FlashAttention kernel per l'attenzione alle query raggruppate (GQA) e ne semplifica l'applicazione a vari scenari per i modelli Hugging Face Transformer. A differenza dell’architettura di attenzione originale, GQA suddivide equamente le teste di query in gruppi e le teste di query dello stesso gruppo condividono le stesse teste di chiavi e valori. Pertanto, le teste q e kv vengono passate alla chiamata in avanti separatamente. Nota: il numero di teste q deve essere divisibile per il numero di teste kv.

Esempio di utilizzo FlashGroupedQueryAttention

Il frammento di codice riportato di seguito mostra come utilizzare l’API torch.sagemaker.nn.attn. FlashGroupedQueryAttention fornita da SMP v2.

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

La libreria SMP fornisce anche torch.sagemaker.nn.huggingface.llama_flashattn. LlamaFlashAttention, che utilizza l’API torch.sagemaker.nn.attn. FlashGroupedQueryAttention a basso livello. A partire dalla v4.36.0., Hugging Face Transformers dispone di un’implementazione simile chiamata LlamaFlashAttention2. Il frammento di codice riportato di seguito mostra come utilizzare l’API LlamaFlashAttention SMP v2 o l’API LlamaFlashAttention2 Transformers per sostituire i livelli di attenzione di un modello Llama esistente.

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))