feat!(attention): enable attention logic

from nanochat
attention_config to be passed to model_config
init weights etc
also update _get_sentence_embedding

Author: meilame-tayebjee

torchTextClassifiers/model/components/__init__.py

@@ -1,3 +1,6 @@
+from .attention import (
+    AttentionConfig as AttentionConfig,
+)
 from .categorical_var_net import (
     CategoricalForwardType as CategoricalForwardType,
 )
@@ -6,3 +9,4 @@
 )
 from .classification_head import ClassificationHead as ClassificationHead
 from .text_embedder import TextEmbedder as TextEmbedder
+from .text_embedder import TextEmbedderConfig as TextEmbedderConfig

torchTextClassifiers/model/components/attention.py

@@ -0,0 +1,126 @@
+"""Largely inspired from Andrej Karpathy's nanochat, see here https://github.com/karpathy/nanochat/blob/master/nanochat/gpt.py"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+### Some utils used in text_embedder.py for the attention blocks ###
+
+
+def apply_rotary_emb(x, cos, sin):
+    assert x.ndim == 4  # multihead attention
+
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:]  # split up last time into two halves
+    y1 = x1 * cos + x2 * sin  # rotate pairs of dims
+    y2 = x1 * (-sin) + x2 * cos
+    out = torch.cat([y1, y2], 3)  # re-assemble
+    out = out.to(x.dtype)  # ensure input/output dtypes match
+    return out
+
+
+def norm(x):
+    # Purely functional rmsnorm with no learnable params
+    return F.rms_norm(x, (x.size(-1),))
+
+
+#### Config #####
+@dataclass
+class AttentionConfig:
+    n_layers: int
+    n_head: int
+    n_kv_head: int
+    sequence_len: Optional[int] = None
+    positional_encoding: bool = True
+    aggregation_method: str = "mean"  # or 'last', or 'first'
+
+
+#### Attention Block #####
+
+# Composed of SelfAttentionLayer and MLP with residual connections
+
+
+class Block(nn.Module):
+    def __init__(self, config: AttentionConfig, layer_idx: int):
+        super().__init__()
+
+        self.layer_idx = layer_idx
+        self.attn = SelfAttentionLayer(config, layer_idx)
+        self.mlp = MLP(config)
+
+    def forward(self, x, cos_sin):
+        x = x + self.attn(norm(x), cos_sin)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+##### Components of the Block #####
+
+
+class SelfAttentionLayer(nn.Module):
+    def __init__(self, config: AttentionConfig, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.enable_gqa = (
+            self.n_head != self.n_kv_head
+        )  # Group Query Attention (GQA): duplicate key/value heads to match query heads if desired
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+
+        self.apply_positional_encoding = config.positional_encoding
+
+    def forward(self, x, cos_sin=None):
+        B, T, C = x.size()
+
+        # Project the input to get queries, keys, and values
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        if self.apply_positional_encoding:
+            assert cos_sin is not None, "Rotary embeddings require precomputed cos/sin tensors"
+            cos, sin = cos_sin
+            q, k = (
+                apply_rotary_emb(q, cos, sin),
+                apply_rotary_emb(k, cos, sin),
+            )  # QK rotary embedding
+
+        q, k = norm(q), norm(k)  # QK norm
+        q, k, v = (
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+        )  # make head be batch dim, i.e. (B, T, H, D) -> (B, H, T, D)
+
+        # is_causal=False for non-autoregressive models (BERT-like)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=self.enable_gqa)
+
+        # Re-assemble the heads side by side and project back to residual stream
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x

torchTextClassifiers/model/components/text_embedder.py

@@ -1,21 +1,107 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
 import torch
 from torch import nn
 
+from torchTextClassifiers.model.components.attention import AttentionConfig, Block, norm
+
+
+@dataclass
+class TextEmbedderConfig:
+    vocab_size: int
+    embedding_dim: int
+    padding_idx: int
+    attention_config: Optional[AttentionConfig] = None
+
 
 class TextEmbedder(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int, padding_idx: int):
+    def __init__(self, text_embedder_config: TextEmbedderConfig):
         super().__init__()
 
-        self.vocab_size = vocab_size
-        self.embedding_dim = embedding_dim
-        self.padding_idx = padding_idx
+        self.config = text_embedder_config
+
+        self.attention_config = text_embedder_config.attention_config
+        if self.attention_config is not None:
+            self.attention_config.n_embd = text_embedder_config.embedding_dim
+
+        self.vocab_size = text_embedder_config.vocab_size
+        self.embedding_dim = text_embedder_config.embedding_dim
+        self.padding_idx = text_embedder_config.padding_idx
 
         self.embedding_layer = nn.Embedding(
-            embedding_dim=embedding_dim,
-            num_embeddings=vocab_size,
+            embedding_dim=self.embedding_dim,
+            num_embeddings=self.vocab_size,
             padding_idx=self.padding_idx,
         )
 
+        if self.attention_config is not None:
+            self.transformer = nn.ModuleDict(
+                {
+                    "h": nn.ModuleList(
+                        [
+                            Block(self.attention_config, layer_idx)
+                            for layer_idx in range(self.attention_config.n_layers)
+                        ]
+                    ),
+                }
+            )
+
+            head_dim = self.attention_config.n_embd // self.attention_config.n_head
+
+            if head_dim * self.attention_config.n_head != self.attention_config.n_embd:
+                raise ValueError("embedding_dim must be divisible by n_head.")
+
+            if self.attention_config.positional_encoding:
+                if head_dim % 2 != 0:
+                    raise ValueError(
+                        "embedding_dim / n_head must be even for rotary positional embeddings."
+                    )
+
+                if self.attention_config.sequence_len is None:
+                    raise ValueError(
+                        "sequence_len must be specified in AttentionConfig when positional_encoding is True."
+                    )
+
+                self.rotary_seq_len = self.attention_config.sequence_len * 10
+                cos, sin = self._precompute_rotary_embeddings(
+                    seq_len=self.rotary_seq_len, head_dim=head_dim
+                )
+
+                self.register_buffer(
+                    "cos", cos, persistent=False
+                )  # persistent=False means it's not saved to the checkpoint
+                self.register_buffer("sin", sin, persistent=False)
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+
+        # zero out c_proj weights in all blocks
+        if self.attention_config is not None:
+            for block in self.transformer.h:
+                torch.nn.init.zeros_(block.mlp.c_proj.weight)
+                torch.nn.init.zeros_(block.attn.c_proj.weight)
+            # init the rotary embeddings
+            head_dim = self.attention_config.n_embd // self.attention_config.n_head
+            cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+            self.cos, self.sin = cos, sin
+        # Cast the embeddings from fp32 to bf16: optim can tolerate it and it saves memory: both in the model and the activations
+        if self.embedding_layer.weight.device.type == "cuda":
+            self.embedding_layer.to(dtype=torch.bfloat16)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # https://arxiv.org/pdf/2310.17813
+            fan_out = module.weight.size(0)
+            fan_in = module.weight.size(1)
+            std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in))
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=1.0)
+
     def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
         """Converts input token IDs to their corresponding embeddings."""
 
@@ -36,6 +122,19 @@ def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torc
             encoded_text
         )  # (batch_size, seq_len, embedding_dim)
 
+        token_embeddings = norm(token_embeddings)
+
+        if self.attention_config is not None:
+            if self.attention_config.positional_encoding:
+                cos_sin = self.cos[:, :seq_len], self.sin[:, :seq_len]
+            else:
+                cos_sin = None
+
+            for block in self.transformer.h:
+                token_embeddings = block(token_embeddings, cos_sin)
+
+            token_embeddings = norm(token_embeddings)
+
         text_embedding = self._get_sentence_embedding(
             token_embeddings=token_embeddings, attention_mask=attention_mask
         )
@@ -57,9 +156,28 @@ def _get_sentence_embedding(
 
         # average over non-pad token embeddings
         # attention mask has 1 for non-pad tokens and 0 for pad token positions
-        # TODO: add attention logic at some point
 
         # mask pad-tokens
+
+        if self.attention_config is not None:
+            if self.attention_config.aggregation_method is not None:
+                if self.attention_config.aggregation_method == "first":
+                    return token_embeddings[:, 0, :]
+                elif self.attention_config.aggregation_method == "last":
+                    lengths = attention_mask.sum(dim=1).clamp(min=1)  # last non-pad token index + 1
+                    return token_embeddings[
+                        torch.arange(token_embeddings.size(0)),
+                        lengths - 1,
+                        :,
+                    ]
+                else:
+                    if self.attention_config.aggregation_method != "mean":
+                        raise ValueError(
+                            f"Unknown aggregation method: {self.attention_config.aggregation_method}. Supported methods are 'mean', 'first', 'last'."
+                        )
+
+        assert self.attention_config is None or self.attention_config.aggregation_method == "mean"
+
         mask = attention_mask.unsqueeze(-1).float()  # (batch_size, seq_len, 1)
         masked_embeddings = token_embeddings * mask  # (batch_size, seq_len, embedding_dim)
 
@@ -79,3 +197,24 @@ def __call__(self, *args, **kwargs):
                 f"(got shape {tuple(out.shape)})"
             )
         return out
+
+    def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
+        # autodetect the device from model embeddings
+        if device is None:
+            device = next(self.parameters()).device
+
+        # stride the channels
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        # stride the time steps
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        # calculate the rotation frequencies at each (time, channel) pair
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()  # keep them in bfloat16
+        cos, sin = (
+            cos[None, :, None, :],
+            sin[None, :, None, :],
+        )  # add batch and head dims for later broadcasting
+
+        return cos, sin

torchTextClassifiers/model/model.py

@@ -69,6 +69,10 @@ def __init__(
 
         self.num_classes = self.classification_head.num_classes
 
+        torch.nn.init.zeros_(self.classification_head.net.weight)
+        if self.text_embedder is not None:
+            self.text_embedder.init_weights()
+
     def _validate_component_connections(self):
         def _check_text_categorical_connection(self, text_embedder, cat_var_net):
             if cat_var_net.forward_type == CategoricalForwardType.SUM_TO_TEXT:

torchTextClassifiers/torchTextClassifiers.py

@@ -23,10 +23,12 @@
 from torchTextClassifiers.dataset import TextClassificationDataset
 from torchTextClassifiers.model import TextClassificationModel, TextClassificationModule
 from torchTextClassifiers.model.components import (
+    AttentionConfig,
     CategoricalForwardType,
     CategoricalVariableNet,
     ClassificationHead,
     TextEmbedder,
+    TextEmbedderConfig,
 )
 from torchTextClassifiers.tokenizers import BaseTokenizer, TokenizerOutput
 
@@ -48,6 +50,7 @@ class ModelConfig:
     categorical_vocabulary_sizes: Optional[List[int]] = None
     categorical_embedding_dims: Optional[Union[List[int], int]] = None
     num_classes: Optional[int] = None
+    attention_config: Optional[AttentionConfig] = None
 
     def to_dict(self) -> Dict[str, Any]:
         return asdict(self)
@@ -139,10 +142,14 @@ def __init__(
             )
             self.embedding_dim = self.tokenizer.output_dim
         else:
-            self.text_embedder = TextEmbedder(
+            text_embedder_config = TextEmbedderConfig(
                 vocab_size=self.vocab_size,
                 embedding_dim=self.embedding_dim,
                 padding_idx=tokenizer.padding_idx,
+                attention_config=model_config.attention_config,
+            )
+            self.text_embedder = TextEmbedder(
+                text_embedder_config=text_embedder_config,
             )
 
         classif_head_input_dim = self.embedding_dim