| 17 | |
| 18 | class Attention(nn.Module): |
| 19 | def __init__(self, args: ModelArgs): |
| 20 | super().__init__() |
| 21 | assert args.dim % args.n_heads == 0 |
| 22 | self.head_dim = args.dim // args.n_heads |
| 23 | self.n_heads = args.n_heads |
| 24 | self.dropout_p = args.dropout_p |
| 25 | self.resid_dropout = nn.Dropout(args.dropout_p) |
| 26 | |
| 27 | self.wq = nn.Linear(args.dim, args.dim, bias=False) |
| 28 | self.wk = nn.Linear(args.dim, args.dim, bias=False) |
| 29 | self.wv = nn.Linear(args.dim, args.dim, bias=False) |
| 30 | self.wo = nn.Linear(args.dim, args.dim, bias=False) |
| 31 | |
| 32 | def forward(self, x): |
| 33 | bsz, seq_len, _ = x.size() |