diff --git a/configs/train/model_shell/autoregressive.yaml b/configs/train/model_shell/autoregressive.yaml index e55e3f02..52a015f6 100644 --- a/configs/train/model_shell/autoregressive.yaml +++ b/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" diff --git a/configs/train/model_shell/autoregressive_byte_level.yaml b/configs/train/model_shell/autoregressive_byte_level.yaml index 1f74f76c..14024658 100644 --- a/configs/train/model_shell/autoregressive_byte_level.yaml +++ b/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" diff --git a/models/autoregressive_model_shell.py b/models/autoregressive_model_shell.py deleted file mode 100644 index 36d9dfff..00000000 --- a/models/autoregressive_model_shell.py +++ /dev/null @@ -1,132 +0,0 @@ -""" -The Model Shell holds the tokenizer, core-model and model head. -""" - -import torch -from torch import nn - -from models.components.lm_heads import NextTokenHead -from models.components.tokenizers import build_tokenizer -from models.utils import print_model_stats -from models.weight_initialization import build_weight_init - - -class AutoregressiveModelShell(nn.Module): - """Code that wraps the model, embedder, and head together.""" - - def __init__( - self, - cfg, - core_model, - ): - super().__init__() - - # move to class - self.cfg = cfg - self.core_model = core_model - - # build the tokenizer - self.tokenizer = build_tokenizer( - tokenizer_type=self.cfg["model_shell"]["tokenizer"], - vocab_size=self.cfg["model_shell"]["vocab_size"], - dataset_name=self.cfg["model_shell"]["tokenizer_dataset_name"], - ) - - # build the embedder - self.token_embedder = nn.Embedding( - num_embeddings=self.cfg["model_shell"]["vocab_size"], - embedding_dim=self.cfg["core_model"]["hidden_dim"], - ) - - # build the language model head - self.lm_head = NextTokenHead( - hidden_dim=self.cfg["core_model"]["hidden_dim"], - vocab_size=self.cfg["model_shell"]["vocab_size"], - ) - - # share the weights between the token embeddings and the final logit layer - self.token_embedder.weight = ( - self.lm_head.linear.weight - ) # https://paperswithcode.com/method/weight-tying - - # report number of parameters - print_model_stats(self) - - # weight init - self.weight_init_func = build_weight_init( - weight_init_type=self.cfg["model_shell"]["weight_init"], - ) - self.init_weights() - - def init_weights(self): - """ - Initialize the weights of the model - """ - self.apply(self.weight_init_func) - - def forward(self, token_ids): - """ - The default forward pass is used for training and accepts the - token_ids as input. When the model is in eval mode, only the - last token is passed into the NextTokenHead. - """ - - _, s = token_ids.size() - - # check that the sequence length is not longer than the context window - assert s <= self.cfg["model_shell"]["context_window"], ( - f"Cannot forward sequence of length {s}, " - f"block size is only {self.cfg['model_shell']['context_window']}" - ) - - # embed token_ids - x = self.token_embedder(token_ids) - - # forward through the core model - x_return = self.core_model(x) - if isinstance(x, tuple): - x, loss = x_return - else: - x, loss = x_return, None - - # get logits - logits = self.lm_head(x) - - return logits, loss - - def inference(self, sequence): - """ - Similar to the forward pass, but takes in a string - (or batch of strings) and only return the logits - for the next token. - Args: - text_string: a string or list of strings - Returns: - logits for the next token - """ - if isinstance(sequence, str): - sequence = [sequence] - - token_ids = self.tokenizer.encode_batch(sequence) - - # pad token_ids and format as tensor - tokens, _ = self.tokenizer.pad_batch(token_ids) - # ignore mask for now... - - _, s = tokens.size() - - # check that the sequence length is not longer than the context window - assert s <= self.cfg["model_shell"]["context_window"], ( - f"Cannot forward sequence of length {s}, " - f"block size is only {self.cfg['model_shell']['context_window']}" - ) - - # embed token_ids - x = self.token_embedder(tokens) - - # forward through the core model - x = self.core_model(x) - - # forward only the last token through the lm_head - logits = self.lm_head(x[:, -1, :]) - return logits diff --git a/models/build_models.py b/models/build_models.py index f97bf7e6..ded4392b 100644 --- a/models/build_models.py +++ b/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 diff --git a/models/components/lm_heads.py b/models/components/lm_heads.py index f39e315e..45ec793f 100644 --- a/models/components/lm_heads.py +++ b/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}") diff --git a/models/core_models/layer_skip.py b/models/core_models/layer_skip.py index 07b4592f..39a124cf 100644 --- a/models/core_models/layer_skip.py +++ b/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... diff --git a/models/shells/__init__.py b/models/shells/__init__.py new file mode 100644 index 00000000..577dd07e --- /dev/null +++ b/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.") diff --git a/models/autoregressive_byte_model_shell.py b/models/shells/autoregressive_byte_model_shell.py similarity index 55% rename from models/autoregressive_byte_model_shell.py rename to models/shells/autoregressive_byte_model_shell.py index d589ae81..c14113a8 100644 --- a/models/autoregressive_byte_model_shell.py +++ b/models/shells/autoregressive_byte_model_shell.py @@ -6,100 +6,19 @@ from torch import nn from models.components.layers import BidirectionalTransformerBlock -from models.components.lm_heads import NextTokenHead -from models.components.tokenizers import build_tokenizer -from models.utils import print_model_stats +from models.shells.shell import Shell -class AutoregressiveByteModelShell(nn.Module): +class AutoregressiveByteModelShell(Shell): """ A model shell for byte-level learning. """ - def __init__( - self, - cfg, - core_model, - ): - super().__init__() - - # move to class - self.cfg = cfg - self.core_model = core_model - - # build the tokenizer - self.byte_tokenizer = build_tokenizer( - tokenizer_type=self.cfg["model_shell"]["pooling_tokenizer"], - vocab_size=self.cfg["model_shell"]["pooling_vocab_size"], - dataset_name=self.cfg["model_shell"]["tokenizer_dataset_name"], - ) - - self.pooling_tokenizer = build_tokenizer( - tokenizer_type=self.cfg["model_shell"]["tokenizer"], - vocab_size=self.cfg["model_shell"]["vocab_size"], - dataset_name=self.cfg["model_shell"]["tokenizer_dataset_name"], - ) - - # build the embedder - self.token_embedder = nn.Embedding( - num_embeddings=self.cfg["model_shell"]["vocab_size"], - embedding_dim=self.cfg["model_shell"]["embedding_dim"], - ) - - self.byte_token_processor = ByteLevelProcessor( - embedding_dim=self.cfg["model_shell"]["embedding_dim"], - hidden_dim=self.cfg["core_model"]["hidden_dim"], - pooling_tokenizer=self.pooling_tokenizer, - byte_tokenizer=self.byte_tokenizer, - token_embedder=self.token_embedder, - ) - - # build the language model head - self.lm_head = NextTokenHead( - hidden_dim=self.cfg["core_model"]["hidden_dim"], - vocab_size=self.cfg["model_shell"]["vocab_size"], - ) - - # share the weights between the token embeddings and the final logit layer - # self.token_embedder.weight = ( - # self.lm_head.linear.weight - # ) # https://paperswithcode.com/method/weight-tying - - # report number of parameters - print_model_stats(self) - - def forward(self, token_ids): + def embed(self, token_ids): """ - The default forward pass is used for training and accepts the - token_ids as input. When the model is in eval mode, only the - last token is passed into the NextTokenHead. + Embed the token ids. """ - - _, s = token_ids.size() - - # check that the sequence length is not longer than the context window - assert s <= self.cfg["model_shell"]["context_window"], ( - f"Cannot forward sequence of length {s}, block size is only" - f" {self.cfg['model_shell']['context_window']}" - ) - - # embed token_ids - # x = self.token_embedder(token_ids) - - # process to sub-word tokens - x = self.byte_token_processor(token_ids) - - # forward through the core model - x_return = self.core_model(x) - if isinstance(x, tuple): - x, loss = x_return - else: - x, loss = x_return, None - - # get logits - logits = self.lm_head(x) - - return logits, loss + return self.tokenizer(token_ids) def inference(self, sequence): """ @@ -114,22 +33,19 @@ def inference(self, sequence): if isinstance(sequence, str): sequence = [sequence] - byte_ids = self.byte_tokenizer.encode_batch(sequence) - - token_ids = self.byte_token_processor.token_to_byte_ids(byte_ids) - - tokens, _ = self.pooling_tokenizer.encode_batch(token_ids) + tokens = self.tokenizer.pooling_tokenizer.encode_batch(sequence) + tokens = torch.tensor(tokens) _, s = tokens.size() # check that the sequence length is not longer than the context window - assert s <= self.cfg["model_shell"]["context_window"], ( + assert s <= self.context_window, ( f"Cannot forward sequence of length {s}," - f" block size is only {self.cfg['model_shell']['context_window']}" + f"max window size is only {self.context_window}" ) # process to sub-word tokens - x = self.byte_token_processor(tokens) + x = self.tokenizer(tokens) # forward through the core model x_return = self.core_model(x) @@ -213,7 +129,6 @@ def forward(self, batch_of_pooled_token_ids): 12, self.embedding_dim, ), - device=torch.device("cuda"), # self.device ) for i, token_batch in enumerate(batch_of_pooled_token_ids): @@ -221,11 +136,11 @@ def forward(self, batch_of_pooled_token_ids): for j, token_id in enumerate(token_batch): # decode into string - token_string = self.pooling_tokenizer.decode([token_id]) + token_string = self.pooling_tokenizer.decode([token_id.item()]) # encode into character ids byte_ids = self.byte_tokenizer.encode(token_string) # convert to tensor - byte_ids = torch.tensor(byte_ids).to("cuda") + byte_ids = torch.tensor(byte_ids) num_ids = len(byte_ids) if num_ids > 12: byte_ids = byte_ids[:12] diff --git a/models/shells/autoregressive_model_shell.py b/models/shells/autoregressive_model_shell.py new file mode 100644 index 00000000..a4744d5c --- /dev/null +++ b/models/shells/autoregressive_model_shell.py @@ -0,0 +1,52 @@ +""" +The Model Shell holds the tokenizer, core-model and model head. +""" + +from models.shells.shell import Shell + + +class AutoregressiveModelShell(Shell): + """Code that wraps the model, embedder, and head together.""" + + def embed(self, token_ids): + """ + Embed the token ids. + """ + return self.token_embedder(token_ids) + + def inference(self, sequence): + """ + Similar to the forward pass, but takes in a string + (or batch of strings) and only return the logits + for the next token. + Args: + text_string: a string or list of strings + Returns: + logits for the next token + """ + if isinstance(sequence, str): + sequence = [sequence] + + token_ids = self.tokenizer.encode_batch(sequence) + + # pad token_ids and format as tensor + tokens, _ = self.tokenizer.pad_batch(token_ids) + # ignore mask for now... + + _, s = tokens.size() + + # check that the sequence length is not longer than the context window + assert s <= self.context_window, ( + f"Cannot forward sequence of length {s}, " + f"block size is only {self.context_window}" + ) + + # embed token_ids + x = self.token_embedder(tokens) + + # forward through the core model + x = self.core_model(x) + + # forward only the last token through the lm_head + logits = self.lm_head(x[:, -1, :]) + return logits diff --git a/models/shells/layerskip_shell.py b/models/shells/layerskip_shell.py new file mode 100644 index 00000000..f6d1b6d0 --- /dev/null +++ b/models/shells/layerskip_shell.py @@ -0,0 +1,71 @@ +""" +The Model Shell for layer skipping, which has a more complex forward pass. +""" + +from models.shells.shell import Shell + + +class LayerSkipShell(Shell): + """ + A model shell for layer skipping. + """ + + def forward(self, token_ids): + """ + Forward pass. + """ + # embed token_ids + x = self.embed(token_ids) + + # forward through the core model + xs = self.core_model(x) + + # forward through the lm_head + logits = self.lm_head(xs) + + return logits + + def embed(self, token_ids): + """ + Embed the token ids. + """ + return self.token_embedder(token_ids) + + def inference(self, sequence): + """ + Similar to the forward pass, but takes in a string + (or batch of strings) and only return the logits + for the next token. + + For the lay + Args: + text_string: a string or list of strings + Returns: + logits for the next token + """ + if isinstance(sequence, str): + sequence = [sequence] + + token_ids = self.tokenizer.encode_batch(sequence) + + # pad token_ids and format as tensor + tokens, _ = self.tokenizer.pad_batch(token_ids) + # ignore mask for now... + + _, s = tokens.size() + + # check that the sequence length is not longer than the context window + assert s <= self.context_window, ( + f"Cannot forward sequence of length {s}, " + f"block size is only {self.context_window}" + ) + + # embed token_ids + x = self.token_embedder(tokens) + + # forward through the core model + x = self.core_model(x) + + # forward only the last token through the lm_head + logits = self.lm_head(x[:, -1, :]) + return logits diff --git a/models/shells/shell.py b/models/shells/shell.py new file mode 100644 index 00000000..fb370360 --- /dev/null +++ b/models/shells/shell.py @@ -0,0 +1,96 @@ +"""Defines the interface for model shells""" + +from torch import nn + +from models.utils import print_model_stats + + +class Shell(nn.Module): + """ + A model shell. + """ + + def __init__( + self, + tokenizer, + token_embedder, + lm_head, + core_model, + weight_init_func=None, + context_window=1024, + ): + super().__init__() + + # move to class + self.core_model = core_model + + # build components + self.tokenizer = tokenizer + self.token_embedder = token_embedder + self.lm_head = lm_head + + # report number of parameters + print_model_stats(self) + + # weight init + self.weight_init_func = weight_init_func + if self.weight_init_func is not None: + self._init_weights() + + self.context_window = context_window + + def _tie_weights(self): + """Tie the weights of the model embeddings and the final logit layer. + + see: https://paperswithcode.com/method/weight-tying""" + self.token_embedder.weight = self.lm_head.linear.weight + + def _init_weights(self): + """ + Initialize the weights of the model. + Does nothing if the weight_init_func is None - in this case + The weights are initialized by the PyTorch default initializer. + """ + self.apply(self.weight_init_func) + + def embed(self, token_ids): + """ + Embed the token ids. + """ + return self.token_embedder(token_ids) + + def forward(self, token_ids): + """ + The default forward pass is used for training and accepts the + token_ids as input. When the model is in eval mode, only the + last token is passed into the NextTokenHead. + """ + + _, s = token_ids.size() + + # check that the sequence length is not longer than the context window + assert s <= self.context_window, ( + f"Cannot forward sequence of length {s}, " + f"max window size is only {self.context_window}" + ) + + # embed token_ids + x = self.embed(token_ids) + + # forward through the core model + x_return = self.core_model(x) + if isinstance(x, tuple): + x, loss = x_return + else: + x, loss = x_return, None + + # get logits + logits = self.lm_head(x) + + return logits, loss + + def inference(self, sequence): + """ + Inference pass. + """ + raise NotImplementedError diff --git a/models/shells/test_shell.py b/models/shells/test_shell.py new file mode 100644 index 00000000..86996779 --- /dev/null +++ b/models/shells/test_shell.py @@ -0,0 +1,108 @@ +"""Test the model shells.""" + +# pylint: disable=missing-function-docstring +# pylint: disable=unused-import +import pytest +import torch + +from models.core_models import GenericTransformer +from models.shells import ( + AutoregressiveByteModelShell, + AutoregressiveModelShell, + build_shell, +) + +AUTO_TEST_CONFIG = { + "model_shell": { + "shell_type": "autoregressive", + "tokenizer": "BPE", + "tokenizer_dataset_name": "debug", + "embedding_dim": 512, + "vocab_size": 512, + "context_window": 512, + "weight_init": "standard", + "head_type": "next_token", + }, + "core_model": { + "core_model_type": "modern", + "hidden_dim": 512, + "depth": 2, + "num_heads": 8, + "ffn_type": "swiglu", + "normalization": "rmsnorm", + "attn_group_size": 1, + "ffn_dim": 512, + "ffn_activation": "gelu", + "dropout": 0.1, + "bias": False, + "attn_type": "rope", + }, +} + +BYTE_TEST_CONFIG = { + "model_shell": { + "shell_type": "autoregressive_byte", + "pooling_tokenizer": "BPE", + "byte_tokenizer": "BPE", + "tokenizer_dataset_name": "debug", + "pooling_vocab_size": 512, + "embedding_dim": 512, + "vocab_size": 512, + "context_window": 512, + "weight_init": "standard", + "head_type": "next_token", + }, + "core_model": { + "core_model_type": "modern", + "hidden_dim": 512, + "depth": 2, + "num_heads": 8, + "ffn_type": "swiglu", + "normalization": "rmsnorm", + "attn_group_size": 1, + "ffn_dim": 512, + "ffn_activation": "gelu", + "dropout": 0.1, + "bias": False, + "attn_type": "rope", + }, +} + +INPUT_STRING = "This is a test string to check the byte level model shell." + + +def test_autoregressive_model_shell(): + """Test the autoregressive model shell.""" + core_model = GenericTransformer(AUTO_TEST_CONFIG) + model = build_shell(AUTO_TEST_CONFIG, core_model) + assert isinstance(model, AutoregressiveModelShell) + # test forward pass with random input + + input_tensor = torch.randint(0, 512, (1, 512)) + output, _ = model(input_tensor) + assert output.shape == (1, 512, 512) + # test inference + + model.eval() + with torch.no_grad(): + output = model.inference(INPUT_STRING) + assert output.shape == (1, 512) + + +def test_autoregressive_byte_model_shell(): + """Test the autoregressive byte model shell.""" + core_model = GenericTransformer(BYTE_TEST_CONFIG) + model = build_shell(BYTE_TEST_CONFIG, core_model) + assert isinstance(model, AutoregressiveByteModelShell) + # test forward pass with random input + + input_tensor = torch.randint(0, 512, (1, 512)) + output, _ = model(input_tensor) + assert output.shape == (1, 512, 512) + + # test inference + + model.eval() + with torch.no_grad(): + output = model.inference(INPUT_STRING) + assert output.shape == (1, 512) diff --git a/models/test_models.py b/models/test_models.py index b1b5ba32..01ac1d22 100644 --- a/models/test_models.py +++ b/models/test_models.py @@ -20,6 +20,8 @@ "vocab_size": 512, "context_window": 512, "weight_init": "standard", + "head_type": "next_token", + "embedding_dim": 512, }, "core_model": { "core_model_type": "modern", @@ -69,7 +71,7 @@ def test_build_model(): model = build_model(sample_cfg) # pass the tokens - _ = model(test_tokens) + _, _ = model(test_tokens) # pass the string assert model.tokenizer is not None