yujunjun/LLMs-from-scratch
1
0
Fork 0
mirror of https://github.com/rasbt/LLMs-from-scratch.git synced 2025-09-26 16:52:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #55 from rayedbw/patch-4

Browse Source 
Update mha-implementations.ipynb
This commit is contained in:
Sebastian Raschka 2024-03-07 06:31:01 -06:00 committed by GitHub
commit 083d11fbd0
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
ch03/02_bonus_efficient-multihead-attention/mha-implementations.ipynb
View File

@ -168,7 +168,7 @@
" # (3, b, num_heads, num_tokens, head_dim) -> 3 times (b, num_head, num_tokens, head_dim)\n", " # (3, b, num_heads, num_tokens, head_dim) -> 3 times (b, num_head, num_tokens, head_dim)\n",
" queries, keys, values = qkv.unbind(0)\n", " queries, keys, values = qkv.unbind(0)\n",
"\n", "\n",
" # (b, num_head, num_tokens, head_dim) --> (b, num_heads, num_tokens, num_tokens)\n", " # (b, num_heads, num_tokens, head_dim) --> (b, num_heads, num_tokens, num_tokens)\n",
" attn_scores = queries @ keys.transpose(-2, -1)\n", " attn_scores = queries @ keys.transpose(-2, -1)\n",
" attn_scores = attn_scores.masked_fill(\n", " attn_scores = attn_scores.masked_fill(\n",
" self.mask.bool()[:num_tokens, :num_tokens], -torch.inf\n", " self.mask.bool()[:num_tokens, :num_tokens], -torch.inf\n",
@ -258,12 +258,12 @@
" # (b, num_tokens, 3, num_heads, head_dim) --> (3, b, num_heads, num_tokens, head_dim)\n", " # (b, num_tokens, 3, num_heads, head_dim) --> (3, b, num_heads, num_tokens, head_dim)\n",
" qkv = qkv.permute(2, 0, 3, 1, 4)\n", " qkv = qkv.permute(2, 0, 3, 1, 4)\n",
"\n", "\n",
" # (3, b, num_heads, num_tokens, head_dim) -> 3 times (b, num_head, num_tokens, head_dim)\n", " # (3, b, num_heads, num_tokens, head_dim) -> 3 times (b, num_heads, num_tokens, head_dim)\n",
" q, k, v = qkv.unbind(0)\n", " queries, keys, values = qkv.unbind(0)\n",
"\n", "\n",
" use_dropout = 0. if not self.training else self.dropout\n", " use_dropout = 0. if not self.training else self.dropout\n",
" context_vec = torch.nn.functional.scaled_dot_product_attention(q, k, v, \n", " context_vec = nn.functional.scaled_dot_product_attention(\n",
" attn_mask=None, dropout_p=use_dropout, is_causal=True)\n", " queries, keys, values, attn_mask=None, dropout_p=use_dropout, is_causal=True)\n",
"\n", "\n",
" # Combine heads, where self.d_out = self.num_heads * self.head_dim\n", " # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
" context_vec = context_vec.transpose(1, 2).contiguous().view(batch_size, num_tokens, self.d_out)\n", " context_vec = context_vec.transpose(1, 2).contiguous().view(batch_size, num_tokens, self.d_out)\n",
@ -396,7 +396,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.6" "version": "3.10.12"
} }
}, },
"nbformat": 4, "nbformat": 4,