repro-expressivity-efficiency-hybrid-sequence / source /official-code /micro /models /transformer.py
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # Causal mask: Prevent attending to future tokens | |
| def generate_mask(sz, window=None): | |
| if not window: | |
| mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) | |
| else: | |
| mask = (torch.triu(torch.ones(sz, sz)) - torch.triu(torch.ones(sz, sz), diagonal=window) == 1).transpose(0, 1) | |
| mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) | |
| return mask | |
| #################################################################################### | |
| def scaled_dot_product_attention(q, k, v, mask=None): | |
| d_k = q.size(-1) | |
| scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32)) | |
| if mask is not None: | |
| scores = scores.masked_fill(mask == float('-inf'), float('-inf')) | |
| # scores = scores.masked_fill(mask == 0, float('-inf')) | |
| attn = F.softmax(scores, dim=-1) | |
| output = torch.matmul(attn, v) | |
| return output, attn | |
| class MultiHeadAttention(nn.Module): | |
| def __init__(self, embed_dim, num_heads, dropout=0.1): | |
| super().__init__() | |
| assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads" | |
| self.embed_dim = embed_dim | |
| self.num_heads = num_heads | |
| # self.head_dim = embed_dim // num_heads | |
| self.head_dim = embed_dim | |
| self.q_proj = nn.Linear(embed_dim, self.head_dim*self.num_heads) | |
| self.k_proj = nn.Linear(embed_dim, self.head_dim*self.num_heads) | |
| self.v_proj = nn.Linear(embed_dim, self.head_dim*self.num_heads) | |
| self.out_proj = nn.Linear(embed_dim, embed_dim) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, query, key, value, mask=None): | |
| B, T, _ = query.size() | |
| # Linear projections | |
| q = self.q_proj(query).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) | |
| k = self.k_proj(key).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) | |
| v = self.v_proj(value).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) | |
| # Apply attention on all the projected vectors in batch | |
| attn_output, _ = scaled_dot_product_attention(q, k, v, mask) | |
| # Concatenate heads and run through final linear layer | |
| # attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.embed_dim) | |
| attn_output = attn_output.transpose(1, 2).sum(dim=2).view(B, T, self.embed_dim) | |
| output = self.out_proj(attn_output) | |
| return output | |
| class TransformerHead(nn.Module): | |
| def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1, do_norm=True): | |
| super().__init__() | |
| self.mha = MultiHeadAttention(embed_dim, num_heads, dropout) | |
| if do_norm: | |
| self.norm1 = nn.LayerNorm(embed_dim) | |
| self.norm2 = nn.LayerNorm(embed_dim) | |
| else: | |
| self.norm1 = nn.Identity(embed_dim) | |
| self.norm2 = nn.Identity(embed_dim) | |
| self.ffn = nn.Sequential( | |
| nn.Linear(embed_dim, ff_dim), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(ff_dim, embed_dim), | |
| ) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, mask=None): | |
| # Multi-head attention + residual + norm | |
| x_norm = self.norm1(x) | |
| attn_out = self.mha(x_norm, x_norm, x_norm, mask) | |
| # x = x + self.norm1(self.dropout(attn_out)) | |
| x = x + self.dropout(attn_out) | |
| # Feedforward + residual + norm | |
| x_norm = self.norm2(x) | |
| ff_out = self.ffn(x_norm) | |
| # x = x + self.norm2(self.dropout(ff_out)) | |
| x = x + self.dropout(ff_out) | |
| return x | |
| class SimpleHandmadeTFLayer(nn.Module): | |
| def __init__(self, d_model, nhead, causal=True, do_norm=True): | |
| super().__init__() | |
| self.transformer_encoder = TransformerHead(d_model, nhead, d_model, dropout=0.2, do_norm=do_norm) | |
| self.d_model = d_model | |
| self.nhead = nhead | |
| self.causal = causal | |
| def forward(self, x, mask): | |
| # xx = x.permute(1, 0, 2) # (seq_len, batch_size, d_model) | |
| xx = x | |
| if self.causal: | |
| xx = self.transformer_encoder(xx, mask=mask) | |
| else: | |
| xx = self.transformer_encoder(xx) | |
| # return xx.permute(1, 0, 2) # Already includes the skip connection | |
| return xx # Already includes the skip connection | |
| # # Finds the average required memory for select data | |
| # def required_activation_memory(self, x, mask, thres=0.01): | |
| # B, T, _ = x.size() | |
| # # Linear projections | |
| # mha = self.transformer_encoder.mha | |
| # q = mha.q_proj(x).view(B, T, self.nhead, self.d_model // self.nhead).transpose(1, 2) | |
| # k = mha.k_proj(x).view(B, T, self.nhead, self.d_model // self.nhead).transpose(1, 2) | |
| # v = mha.v_proj(x).view(B, T, self.nhead, self.d_model // self.nhead).transpose(1, 2) | |
| # attn_output, attn_weights = scaled_dot_product_attention(q, k, v, mask) | |
| # # This calculation is saying 'we only need the tokens where the last token attends highly with it' | |
| # # It's not exactly what we are going for, so we need to think about this more | |
| # return (torch.sum(attn_weights[:, 0, -1, :] > thres) / x.shape[0] * self.d_model).item() | |
| #################################################################################### | |
| # A simple wrapper for the PyTorch transformer encoder layers | |
| class SimplePyTorchTFLayer(nn.Module): | |
| def __init__(self, d_model, nhead, causal=True): | |
| super().__init__() | |
| encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=d_model, dropout=0.2) | |
| self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=1) | |
| self.d_model = d_model | |
| self.nhead = nhead | |
| self.causal = causal | |
| def forward(self, x, mask): | |
| xx = x.permute(1, 0, 2) # (seq_len, batch_size, d_model) | |
| # xx = x | |
| if self.causal: | |
| xx = self.transformer_encoder(xx, mask=mask) | |
| else: | |
| xx = self.transformer_encoder(xx) | |
| # return x + xx | |
| return x + xx.permute(1, 0, 2) # (batch_size, seq_len, vocab_size) |