1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100

    def forward(
        self,
        hidden_states: Optional[tuple[torch.FloatTensor]],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        **kwargs,
    ) -> tuple[Union[torch.Tensor, tuple[torch.Tensor]], ...]:
        # 判断是否是Cross Attention
        is_cross_attention = encoder_hidden_states is not None
        # Cross Attention使用cross_attention_cache
        # Self Attention使用self_attention_cache
        # 用is_updated表示当前层的KV是否已缓存 (用于Cross Attention)
        if past_key_value is not None:
            if isinstance(past_key_value, EncoderDecoderCache):
                is_updated = past_key_value.is_updated.get(self.layer_idx)
                if is_cross_attention:
                    curr_past_key_value = past_key_value.cross_attention_cache
                else:
                    curr_past_key_value = past_key_value.self_attention_cache
            else:
                curr_past_key_value = past_key_value


        if is_cross_attention:
            # Cross Attention
            query_states = self.q_attn(hidden_states)
            attention_mask = encoder_attention_mask
            # 尝试获取KV Cache
            if past_key_value is not None and is_updated:
                key_states = curr_past_key_value.layers[self.layer_idx].keys
                value_states = curr_past_key_value.layers[self.layer_idx].values
            else:
                key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
                # 变换成MHA的shape
                shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
                key_states = key_states.view(shape_kv).transpose(1, 2)
                value_states = value_states.view(shape_kv).transpose(1, 2)
        else:
            # Self Attention
            query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
            shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
            key_states = key_states.view(shape_kv).transpose(1, 2)
            value_states = value_states.view(shape_kv).transpose(1, 2)


        shape_q = (*query_states.shape[:-1], -1, self.head_dim)
        query_states = query_states.view(shape_q).transpose(1, 2)

        # 更新缓存: 启动KV Cache，且是Self Attention，或Cross Attention没有缓存过的情况
        if (past_key_value is not None and not is_cross_attention) or (
            past_key_value is not None and is_cross_attention and not is_updated
        ):
            cache_position = cache_position if not is_cross_attention else None
            key_states, value_states = curr_past_key_value.update(
                key_states, value_states, self.layer_idx, {"cache_position": cache_position}
            )
            if is_cross_attention:
                past_key_value.is_updated[self.layer_idx] = True

        # 判断是否是因果注意力 (Casual)
        is_causal = attention_mask is None and query_states.shape[-2] > 1 and not is_cross_attention

        # 选择注意力实现方式
        # [eager/flash_attention_2/sdpa/triton/xformers]
        using_eager = self.config._attn_implementation == "eager"
        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        # 选择精度提升(upcast)和重排(reorder)
        if using_eager and self.reorder_and_upcast_attn:
            attn_output, attn_weights = self._upcast_and_reordered_attn(
                query_states, key_states, value_states, attention_mask, head_mask
            )
        else:
            # 调用注意力计算函数
            attn_output, attn_weights = attention_interface(
                self,
                query_states,
                key_states,
                value_states,
                attention_mask,
                head_mask=head_mask,
                dropout=self.attn_dropout.p if self.training else 0.0,
                is_causal=is_causal,
                **kwargs,
            )

        # 将Attention结果用线性层c_proj投影回原始维度
        attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
        attn_output = self.c_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)


        return attn_output, attn_weights

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 多个Query头
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        # 共享的Key和Value
        self.k_proj = nn.Linear(embed_dim, self.head_dim)
        self.v_proj = nn.Linear(embed_dim, self.head_dim)

        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        B, T, C = x.size()

        # Q: (B, T, num_heads, head_dim)
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # (B, num_heads, T, head_dim)

        # K, V: shared (B, T, head_dim)->(B, 1, T, head_dim)
        k = self.k_proj(x).unsqueeze(1)
        v = self.v_proj(x).unsqueeze(1)

        # Attention Score: (B, num_heads, T, T)
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(attn_scores, dim=-1)

        # Attention output: (B, num_heads, T, head_dim)
        attn_output = torch.matmul(attn_weights, v)

        # 合并头->(B, T, embed_dim)
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)

        return self.out_proj(attn_output)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

using namespace std;

template<typename T>
void OnlineSoftmax(T* dst, const T* src, int n) {
    T m = -numeric_limits<T>::infinity();
    T d = 0.0f;

    for (int i = 0; i < n; +i) {
        T m_update = max(m, src[i]);
        d = d * exp(m - m_update) + exp(src[i] - m_update);
        m = m_update;
    }

    for (int i = 0; i < n; ++i) {
        dst[i] = exp(src[i] - m) / d;
    }
}

1
2

model = ...()
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import os
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '...'

dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
model = ...().to(rank)
model = DDP(model, device_ids=[rank])
...
dist.destroy_process_group()