Layer skip

configs/train/model_shell/autoregressive.yaml

@@ -4,3 +4,4 @@ tokenizer_dataset_name: "simple_en_wiki"
 vocab_size: 50304
 context_window: 512
 weight_init: gpt2
+head_type: "next_token"

configs/train/model_shell/autoregressive_byte_level.yaml

@@ -7,3 +7,4 @@ pooling_vocab_size: 512
 tokenizer_dataset_name: "simple_en_wiki"
 context_window: 512
 weight_init: gpt2
+head_type: "next_token"

models/build_models.py

@@ -2,14 +2,13 @@
 Contains the functions to build the actual model.
 """
 
-from models.autoregressive_byte_model_shell import AutoregressiveByteModelShell
-from models.autoregressive_model_shell import AutoregressiveModelShell
 from models.core_models import (
     GenericTransformer,
     ModernFFNSharingTransformer,
     ModernTransformer,
     StandardTransformer,
 )
+from models.shells import build_shell
 
 
 def build_model(cfg=None, checkpoint=None):
@@ -60,28 +59,6 @@ def build_core_model(cfg):
     return CORE_MODEL_DICT[cfg["core_model"]["core_model_type"]](cfg=cfg)
 
 
-MODEL_SHELL_DICT = {
-    "autoregressive": AutoregressiveModelShell,
-    "autoregressive_byte_encoding": AutoregressiveByteModelShell,
-}
-
-
-def build_shell(cfg, core_model):
-    """
-    Given the model shell config, build it.
-    Args:
-        cfg: cfg
-        tokenizer: tokenizer_instance
-        core_model: core_model_instance
-    Returns:
-        model_shell: model_shell_instance
-    """
-    return MODEL_SHELL_DICT[cfg["model_shell"]["shell_type"]](
-        cfg=cfg,
-        core_model=core_model,
-    )
-
-
 def initialize_model(cfg):
     """
     Initializes the model from the configuration
@@ -94,9 +71,6 @@ def initialize_model(cfg):
     core_model = build_core_model(cfg)
 
     # build the model shell
-    model = build_shell(
-        cfg=cfg,
-        core_model=core_model,
-    )
+    model = build_shell(cfg, core_model)
 
     return model

models/components/lm_heads.py

@@ -7,7 +7,20 @@
 from models.components.layers import LayerNorm
 
 
-class NextTokenHead(nn.Module):
+class LMHead(nn.Module):
+    """Interface for the Language Model head."""
+
+    def forward(self, x):
+        """Forward pass."""
+        raise NotImplementedError
+
+    def inference(self, x):
+        """Inference."""
+        x = x[:, -1, :]
+        return self(x)
+
+
+class NextTokenHead(LMHead):
     """Next token prediction head."""
 
     def __init__(self, hidden_dim, vocab_size):
@@ -20,3 +33,11 @@ def forward(self, x):
         x = self.ln(x)
         logits = self.linear(x)
         return logits
+
+
+def build_lm_head(head_type, hidden_dim, vocab_size):
+    """Build a language model head."""
+    if head_type == "next_token":
+        return NextTokenHead(hidden_dim, vocab_size)
+    else:
+        raise ValueError(f"Unknown head type: {head_type}")

models/core_models/layer_skip.py

@@ -64,4 +64,4 @@ def forward(self, x, attention_mask=None):
             early_exits.append(x)
         return torch.stack(
             early_exits, dim=1
-        )  # very hacky way to pass the early exits all through the head...
+        )  # very hacky way to pass the early exits all through the head... may cause memory issues...

models/shells/__init__.py

@@ -0,0 +1,65 @@
+"""Contains the shell classes for the models."""
+
+from torch import nn
+
+from models.components.lm_heads import build_lm_head
+from models.components.tokenizers import build_tokenizer
+from models.shells.autoregressive_byte_model_shell import (
+    AutoregressiveByteModelShell,
+    ByteLevelProcessor,
+)
+from models.shells.autoregressive_model_shell import AutoregressiveModelShell
+
+
+def build_shell(cfg, core_model):
+    """
+    Build the shell.
+    """
+
+    lm_head = build_lm_head(
+        head_type=cfg["model_shell"]["head_type"],
+        hidden_dim=cfg["core_model"]["hidden_dim"],
+        vocab_size=cfg["model_shell"]["vocab_size"],
+    )
+    token_embedder = nn.Embedding(
+        num_embeddings=cfg["model_shell"]["vocab_size"],
+        embedding_dim=cfg["model_shell"]["embedding_dim"],
+    )
+
+    if cfg["model_shell"]["shell_type"] == "autoregressive":
+        tokenizer = build_tokenizer(
+            tokenizer_type=cfg["model_shell"]["tokenizer"],
+            vocab_size=cfg["model_shell"]["vocab_size"],
+            dataset_name=cfg["model_shell"]["tokenizer_dataset_name"],
+        )
+        return AutoregressiveModelShell(
+            tokenizer=tokenizer,
+            token_embedder=token_embedder,
+            lm_head=lm_head,
+            core_model=core_model,
+        )
+    elif cfg["model_shell"]["shell_type"] == "autoregressive_byte":
+        pooling_tokenizer = build_tokenizer(
+            tokenizer_type=cfg["model_shell"]["pooling_tokenizer"],
+            vocab_size=cfg["model_shell"]["pooling_vocab_size"],
+            dataset_name=cfg["model_shell"]["tokenizer_dataset_name"],
+        )
+        byte_tokenizer = build_tokenizer(
+            tokenizer_type=cfg["model_shell"]["byte_tokenizer"],
+            vocab_size=cfg["model_shell"]["vocab_size"],
+            dataset_name=cfg["model_shell"]["tokenizer_dataset_name"],
+        )
+        processor = ByteLevelProcessor(
+            embedding_dim=cfg["model_shell"]["embedding_dim"],
+            hidden_dim=cfg["core_model"]["hidden_dim"],
+            pooling_tokenizer=pooling_tokenizer,
+            byte_tokenizer=byte_tokenizer,
+            token_embedder=token_embedder,
+        )
+        return AutoregressiveByteModelShell(
+            tokenizer=processor,
+            token_embedder=token_embedder,
+            lm_head=lm_head,
+            core_model=core_model,
+        )
+    raise ValueError(f"Shell type {cfg['shell']['type']} not recognized.")