| 132 | |
| 133 | |
| 134 | class Attention(nn.Module): |
| 135 | def __init__(self, dim: int, n_heads: int, n_kv_heads: Optional[int], qk_norm: bool): |
| 136 | """ |
| 137 | Initialize the Attention module. |
| 138 | |
| 139 | Args: |
| 140 | dim (int): Number of input dimensions. |
| 141 | n_heads (int): Number of heads. |
| 142 | n_kv_heads (Optional[int]): Number of kv heads, if using GQA. |
| 143 | |
| 144 | Attributes: |
| 145 | n_kv_heads (int): Number of key and value heads. |
| 146 | n_local_heads (int): Number of local query heads. |
| 147 | n_local_kv_heads (int): Number of local key and value heads. |
| 148 | n_rep (int): Number of repetitions for local heads. |
| 149 | head_dim (int): Dimension size of each attention head. |
| 150 | wq (ColumnParallelLinear): Linear transformation for queries. |
| 151 | wk (ColumnParallelLinear): Linear transformation for keys. |
| 152 | wv (ColumnParallelLinear): Linear transformation for values. |
| 153 | wo (RowParallelLinear): Linear transformation for output. |
| 154 | cache_k (torch.Tensor): Cached keys for attention. |
| 155 | cache_v (torch.Tensor): Cached values for attention. |
| 156 | |
| 157 | """ |
| 158 | super().__init__() |
| 159 | self.n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads |
| 160 | model_parallel_size = fs_init.get_model_parallel_world_size() |
| 161 | self.n_local_heads = n_heads // model_parallel_size |
| 162 | self.n_local_kv_heads = self.n_kv_heads // model_parallel_size |
| 163 | self.n_rep = self.n_local_heads // self.n_local_kv_heads |
| 164 | self.head_dim = dim // n_heads |
| 165 | |
| 166 | self.wq = ColumnParallelLinear( |
| 167 | dim, n_heads * self.head_dim, bias=False, gather_output=False, |
| 168 | init_method=nn.init.xavier_uniform_, |
| 169 | ) |
| 170 | self.wk = ColumnParallelLinear( |
| 171 | dim, self.n_kv_heads * self.head_dim, bias=False, |
| 172 | gather_output=False, init_method=nn.init.xavier_uniform_, |
| 173 | ) |
| 174 | self.wv = ColumnParallelLinear( |
| 175 | dim, self.n_kv_heads * self.head_dim, bias=False, |
| 176 | gather_output=False, init_method=nn.init.xavier_uniform_, |
| 177 | ) |
| 178 | self.wo = RowParallelLinear( |
| 179 | n_heads * self.head_dim, dim, bias=False, |
| 180 | input_is_parallel=True, init_method=nn.init.xavier_uniform_, |
| 181 | ) |
| 182 | |
| 183 | if qk_norm: |
| 184 | self.q_norm = nn.LayerNorm(self.n_local_heads * self.head_dim) |
| 185 | self.k_norm = nn.LayerNorm(self.n_local_kv_heads * self.head_dim) |
| 186 | else: |
| 187 | self.q_norm = self.k_norm = nn.Identity() |
| 188 | |
| 189 | @staticmethod |
| 190 | def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): |
| 191 | """ |