From b832acc45f7efcc78dbf98130d6d1269bfa07d62 Mon Sep 17 00:00:00 2001 From: Dylan Hillier Date: Wed, 1 May 2024 13:46:50 +0800 Subject: [PATCH 1/4] wip layer skip! --- models/autoregressive_model_shell.py | 42 +++++++---------- models/core_models/layer_skip.py | 67 ++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 26 deletions(-) create mode 100644 models/core_models/layer_skip.py diff --git a/models/autoregressive_model_shell.py b/models/autoregressive_model_shell.py index cc6bd822..7a4d2aaa 100644 --- a/models/autoregressive_model_shell.py +++ b/models/autoregressive_model_shell.py @@ -2,40 +2,34 @@ The Model Shell holds the tokenizer, core-model and model head. """ -import torch +import torch import torch.nn as nn +from models.components.LMHeads import NextTokenHead from models.components.tokenizers import build_tokenizer -from models.components.LMHeads import ( - NextTokenHead -) - -from models.utils import ( - print_model_stats -) - +from models.utils import print_model_stats class AutoregressiveModelShell(nn.Module): def __init__( - self, - cfg, - core_model, - ): + self, + cfg, + core_model, + ): super().__init__() # move to class self.cfg = cfg self.core_model = core_model - # build the tokenizer + # 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 + # build the embedder self.token_embedder = nn.Embedding( num_embeddings=self.cfg["model_shell"]["vocab_size"], embedding_dim=self.cfg["core_model"]["hidden_dim"], @@ -50,17 +44,15 @@ def __init__( # 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 - + ) # https://paperswithcode.com/method/weight-tying # report number of parameters print_model_stats(self) - 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 + 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. """ @@ -71,7 +63,6 @@ def forward(self, token_ids): s <= self.cfg["model_shell"]["context_window"] ), f"Cannot forward sequence of length {s}, block size is only {self.cfg['model_shell']['context_window']}" - # embed token_ids x = self.token_embedder(token_ids) @@ -86,11 +77,11 @@ def forward(self, token_ids): logits = self.lm_head(x) return logits, loss - + def inference(self, text_string, output_tokens): """ - Similar to the forward pass, but takes in a string - (or batch of strings) and only return the logits + 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 @@ -114,13 +105,12 @@ def inference(self, text_string, output_tokens): s <= self.cfg["model_shell"]["context_window"] ), f"Cannot forward sequence of length {s}, 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 = self.core_model(x) - # forward only the last token through the lm_head + # forward only the last token through the lm_head logits = self.lm_head(x[:, -1, :]) return logits diff --git a/models/core_models/layer_skip.py b/models/core_models/layer_skip.py new file mode 100644 index 00000000..7b786ae8 --- /dev/null +++ b/models/core_models/layer_skip.py @@ -0,0 +1,67 @@ +"""A normal transformer but skips some layers and has early exits.""" + +from models.core_models import generic + +import torch + + +class LayerSkipTransformer(generic.GenericTransformerBlock): + """ + A transformer that skips some layers and has early exits. + """ + + def __init__( + self, + hidden_dim, + ffn_type, + ffn_dim, + bias, + num_heads, + normalization="layernorm", + attn_type="causal", + ffn_activation=None, + max_layer_dropout=1.0, + ): + super().__init__( + hidden_dim, + ffn_type, + ffn_dim, + bias, + num_heads, + normalization, + attn_type, + ffn_activation, + ) + + self.dropouts = torch.nn.ModuleList( + [torch.nn.Dropout() for _ in range(len(self.layers))] + ) + self.max_layer_dropout = max_layer_dropout + + def set_dropouts(self, iteration, max_iteration): + """ + Set the dropout probability for a given layer + """ + scale_t = torch.exp(iteration * torch.log(2) / max_iteration) - 1 + for i, dropout in enumerate(self.dropouts): + depth_layerwise = torch.exp(i * torch.log(2) // (len(self.layers) - 1)) - 1 + dropout.p = scale_t * depth_layerwise * self.max_layer_dropout + + def forward(self, x, attention_mask=None): + """ + A simple, residual forward + """ + if self.pos_encoder is not None: + x = x + self.pos_encoder(x) + x = self.transformer.drop(x) + early_exits = [] + for i, block in enumerate(self.transformer.h): + layer_dropout_update_mask = torch.ones(x.size(0), device=x.device) + layer_dropout_update_mask = self.dropouts[i](layer_dropout_update_mask) + x[layer_dropout_update_mask] = block(x, attention_mask)[ + layer_dropout_update_mask + ] + early_exits.append(x) + return torch.stack( + early_exits, dim=1 + ) # very hacky way to pass the early exits all through the head... may cause memory issues... From d45186d437a284a630cd2ff89c94d4be28c09b54 Mon Sep 17 00:00:00 2001 From: Dylan Hillier Date: Tue, 7 May 2024 17:10:28 +0800 Subject: [PATCH 2/4] refactorings to the shells --- models/autoregressive_model_shell.py | 132 ------------------ models/build_models.py | 30 +--- models/components/lm_heads.py | 23 ++- models/shells/__init__.py | 65 +++++++++ .../autoregressive_byte_model_shell.py | 103 +++----------- models/shells/autoregressive_model_shell.py | 63 +++++++++ models/shells/layerskip_shell.py | 0 models/shells/shell.py | 87 ++++++++++++ 8 files changed, 258 insertions(+), 245 deletions(-) delete mode 100644 models/autoregressive_model_shell.py create mode 100644 models/shells/__init__.py rename models/{ => shells}/autoregressive_byte_model_shell.py (60%) create mode 100644 models/shells/autoregressive_model_shell.py create mode 100644 models/shells/layerskip_shell.py create mode 100644 models/shells/shell.py 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/shells/__init__.py b/models/shells/__init__.py new file mode 100644 index 00000000..20680bd0 --- /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"]["lm_head"], + 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["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["shell"]["type"] == "autoregressive_byte": + pooling_tokenizer = build_tokenizer( + tokenizer_type=cfg["model_shell"]["pooling_tokenizer"], + vocab_size=cfg["model_shell"]["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( + processor=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 60% rename from models/autoregressive_byte_model_shell.py rename to models/shells/autoregressive_byte_model_shell.py index d589ae81..77ab2105 100644 --- a/models/autoregressive_byte_model_shell.py +++ b/models/shells/autoregressive_byte_model_shell.py @@ -6,100 +6,35 @@ 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, + processor, + token_embedder, + lm_head, core_model, + weight_init_func=None, ): - 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"], + super().__init__( + tokenizer=processor, + lm_head=lm_head, + token_embedder=token_embedder, + core_model=core_model, + weight_init_func=weight_init_func, ) - 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,11 +49,11 @@ def inference(self, sequence): if isinstance(sequence, str): sequence = [sequence] - byte_ids = self.byte_tokenizer.encode_batch(sequence) + byte_ids = self.tokenizer.byte_tokenizer.encode_batch(sequence) - token_ids = self.byte_token_processor.token_to_byte_ids(byte_ids) + token_ids = self.tokenizer.token_to_byte_ids(byte_ids) - tokens, _ = self.pooling_tokenizer.encode_batch(token_ids) + tokens, _ = self.tokenizer.pooling_tokenizer.encode_batch(token_ids) _, s = tokens.size() @@ -129,7 +64,7 @@ def inference(self, sequence): ) # process to sub-word tokens - x = self.byte_token_processor(tokens) + x = self.processor(tokens) # forward through the core model x_return = self.core_model(x) diff --git a/models/shells/autoregressive_model_shell.py b/models/shells/autoregressive_model_shell.py new file mode 100644 index 00000000..fc497754 --- /dev/null +++ b/models/shells/autoregressive_model_shell.py @@ -0,0 +1,63 @@ +""" +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 __init__( + self, tokenizer, token_embedder, lm_head, core_model, weight_init_func=None + ): + super().__init__( + tokenizer=tokenizer, + token_embedder=token_embedder, + lm_head=lm_head, + core_model=core_model, + weight_init_func=weight_init_func, + ) + + 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.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/shells/layerskip_shell.py b/models/shells/layerskip_shell.py new file mode 100644 index 00000000..e69de29b diff --git a/models/shells/shell.py b/models/shells/shell.py new file mode 100644 index 00000000..fcdc0813 --- /dev/null +++ b/models/shells/shell.py @@ -0,0 +1,87 @@ +"""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 + ): + 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 + self._init_weights() + + 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.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.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 From 3838ce6a5fb3b9f9a30cbb79909c6f46552bdb82 Mon Sep 17 00:00:00 2001 From: Dylan Hillier Date: Tue, 7 May 2024 18:24:23 +0800 Subject: [PATCH 3/4] fixes the tests, not sure if it works --- configs/train/model_shell/autoregressive.yaml | 1 + .../autoregressive_byte_level.yaml | 1 + models/shells/__init__.py | 10 +- .../shells/autoregressive_byte_model_shell.py | 34 ++---- models/shells/autoregressive_model_shell.py | 15 +-- models/shells/shell.py | 17 ++- models/shells/test_shell.py | 108 ++++++++++++++++++ models/test_models.py | 4 +- 8 files changed, 140 insertions(+), 50 deletions(-) create mode 100644 models/shells/test_shell.py 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/shells/__init__.py b/models/shells/__init__.py index 20680bd0..577dd07e 100644 --- a/models/shells/__init__.py +++ b/models/shells/__init__.py @@ -17,7 +17,7 @@ def build_shell(cfg, core_model): """ lm_head = build_lm_head( - head_type=cfg["model_shell"]["lm_head"], + head_type=cfg["model_shell"]["head_type"], hidden_dim=cfg["core_model"]["hidden_dim"], vocab_size=cfg["model_shell"]["vocab_size"], ) @@ -26,7 +26,7 @@ def build_shell(cfg, core_model): embedding_dim=cfg["model_shell"]["embedding_dim"], ) - if cfg["shell"]["type"] == "autoregressive": + if cfg["model_shell"]["shell_type"] == "autoregressive": tokenizer = build_tokenizer( tokenizer_type=cfg["model_shell"]["tokenizer"], vocab_size=cfg["model_shell"]["vocab_size"], @@ -38,10 +38,10 @@ def build_shell(cfg, core_model): lm_head=lm_head, core_model=core_model, ) - elif cfg["shell"]["type"] == "autoregressive_byte": + elif cfg["model_shell"]["shell_type"] == "autoregressive_byte": pooling_tokenizer = build_tokenizer( tokenizer_type=cfg["model_shell"]["pooling_tokenizer"], - vocab_size=cfg["model_shell"]["vocab_size"], + vocab_size=cfg["model_shell"]["pooling_vocab_size"], dataset_name=cfg["model_shell"]["tokenizer_dataset_name"], ) byte_tokenizer = build_tokenizer( @@ -57,7 +57,7 @@ def build_shell(cfg, core_model): token_embedder=token_embedder, ) return AutoregressiveByteModelShell( - processor=processor, + tokenizer=processor, token_embedder=token_embedder, lm_head=lm_head, core_model=core_model, diff --git a/models/shells/autoregressive_byte_model_shell.py b/models/shells/autoregressive_byte_model_shell.py index 77ab2105..c14113a8 100644 --- a/models/shells/autoregressive_byte_model_shell.py +++ b/models/shells/autoregressive_byte_model_shell.py @@ -14,22 +14,6 @@ class AutoregressiveByteModelShell(Shell): A model shell for byte-level learning. """ - def __init__( - self, - processor, - token_embedder, - lm_head, - core_model, - weight_init_func=None, - ): - super().__init__( - tokenizer=processor, - lm_head=lm_head, - token_embedder=token_embedder, - core_model=core_model, - weight_init_func=weight_init_func, - ) - def embed(self, token_ids): """ Embed the token ids. @@ -49,22 +33,19 @@ def inference(self, sequence): if isinstance(sequence, str): sequence = [sequence] - byte_ids = self.tokenizer.byte_tokenizer.encode_batch(sequence) - - token_ids = self.tokenizer.token_to_byte_ids(byte_ids) - - tokens, _ = self.tokenizer.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.processor(tokens) + x = self.tokenizer(tokens) # forward through the core model x_return = self.core_model(x) @@ -148,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): @@ -156,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 index fc497754..a4744d5c 100644 --- a/models/shells/autoregressive_model_shell.py +++ b/models/shells/autoregressive_model_shell.py @@ -8,17 +8,6 @@ class AutoregressiveModelShell(Shell): """Code that wraps the model, embedder, and head together.""" - def __init__( - self, tokenizer, token_embedder, lm_head, core_model, weight_init_func=None - ): - super().__init__( - tokenizer=tokenizer, - token_embedder=token_embedder, - lm_head=lm_head, - core_model=core_model, - weight_init_func=weight_init_func, - ) - def embed(self, token_ids): """ Embed the token ids. @@ -47,9 +36,9 @@ def inference(self, sequence): _, 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"block size is only {self.context_window}" ) # embed token_ids diff --git a/models/shells/shell.py b/models/shells/shell.py index fcdc0813..fb370360 100644 --- a/models/shells/shell.py +++ b/models/shells/shell.py @@ -11,7 +11,13 @@ class Shell(nn.Module): """ def __init__( - self, tokenizer, token_embedder, lm_head, core_model, weight_init_func=None + self, + tokenizer, + token_embedder, + lm_head, + core_model, + weight_init_func=None, + context_window=1024, ): super().__init__() @@ -28,7 +34,10 @@ def __init__( # weight init self.weight_init_func = weight_init_func - self._init_weights() + 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. @@ -60,9 +69,9 @@ def forward(self, 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"], ( + 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}" ) # embed token_ids 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 From d2d5038958c7ea496b35479d6fc42ed9804c0e5b Mon Sep 17 00:00:00 2001 From: Dylan Hillier Date: Thu, 9 May 2024 15:42:07 +0800 Subject: [PATCH 4/4] some changes --- models/shells/layerskip_shell.py | 71 ++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/models/shells/layerskip_shell.py b/models/shells/layerskip_shell.py index e69de29b..f6d1b6d0 100644 --- a/models/shells/layerskip_shell.py +++ 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