diff --git a/README.md b/README.md index 1d39a22b..3b2009a4 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ This GitHub repository presents our research on Super Tiny Language Models (STLM Our mission is to enhance the accessibility and practicality of high-performing language models across a wide range of applications by drastically reducing their computational and energy demands. We believe that our approach has the potential to contribute to more sustainable and inclusive AI development. -For a comprehensive understanding of our research methodology and initial findings, we strongly encourage you to read our paper: [Super Tiny Language Models](https://arxiv.org/abs/2405.14159) +For a comprehensive understanding of our research methodology and initial findings, we strongly encourage you to read our paper: [Super Tiny Language Models](https://arxiv.org/abs/2405.14159). If you are new to language models we provide some reasources to read up on them in the [STARTER.md](STARTER.md) file. Please note that this repository is an evolving work in progress, reflecting the ongoing nature of our research. It is subject to frequent updates and improvements as we continue to explore and refine our work on STLMs. We welcome the community's engagement with our work, value your feedback, and appreciate any contributions to this challenging but promising endeavor. diff --git a/STARTER.md b/STARTER.md new file mode 100644 index 00000000..c2e12dc0 --- /dev/null +++ b/STARTER.md @@ -0,0 +1,34 @@ +# Resources for Getting Started with Language Models +## Important Papers +While academic papers can be difficult to read, don't worry about trying to understand everything. Oftentimes the fine details are not relevant, so try to focus on understanding the main ideas. Here are some papers that are important to understand the development of language models: +- [the original transformer paper](https://arxiv.org/abs/1706.03762): This paper is an *encoder-decoder* architecture, which is not popular these days. Main things to try to understand from the paper are the introduction of the attention mechanism. +- [BERT](https://arxiv.org/abs/1810.04805): This paper introduces a 'masked language modelling objective' - reconstructing a partially obscured version of the original text. It is an *encoder* only model +- [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf): An example of a *decoder* only model. It uses a causal language modelling objective - predicting the next word in a sentence. More importantly, it shows that just scaling up the model size/data can lead to impressive zero-shot performance on a variety of tasks +- [GPT-3](https://arxiv.org/abs/2005.14165): This paper introduces a few-shot learning objective - predicting the next word in a sentence given a few examples of the task. + +## Blogs +- [The Illustrated Transformer](http://jalammar.github.io/illustrated-transformer/): A great blog post that explains the transformer architecture in a very visual way. +- [Illustrated GPT-2](http://jalammar.github.io/illustrated-gpt2/): Another great blog post that explains the GPT-2 architecture in a very visual way. + +## Huggingface +Huggingface is the de-facto standard for working with and sharing language models. It's great for downstream tasks using pre-trained models, and they provide helpful tools for understanding the concepts behind language models as well as tools for finetuning, a massive collection of datasets, and a large collection of pre-trained models. Here are some resources to get started with Huggingface: +- [Huggingface Transformers Documentation](https://huggingface.co/transformers/): The official documentation for the Huggingface Transformers library. +- [Huggingface Datasets Documentation](https://huggingface.co/docs/datasets/): The official documentation for the Huggingface Datasets library. +- [Huggingface Model Hub](https://huggingface.co/models): The official model hub for Huggingface. You can find a large collection of pre-trained models here. + +## Videos +While we can't vouch for them [Stanford](https://www.youtube.com/playlist?list=PLoROMvodv4rMFqRtEuo6SGjY4XbRIVRd4) and [MIT]() have some great free lectures on deep learning and NLP. + +## Repositories +We highly recommend trying to train a language model from scratch to better understand the concepts. Here are some repositories that can help you get started: +- [MinGPT](https://github.com/karpathy/minGPT) - Very simple PyTorch implementation of GPT-2, also the basis of this repository + +## Key Concepts +- **Tokenization**: The process of breaking up text into smaller pieces, usually words or subwords. These have integer representations that can be used as input to a model. +- **Embeddings**: A way to represent words as vectors. These vectors are learned during training and are used as input to the model. Additionally the final outputs of the model also act as embeddings. +- **Logits**: The raw output of a model before it is converted to probabilities. +- **Feed Forward Neural Network**: A neural network with (typically) multiple, fully connected layers that each perform a linear transformation followed by a non-linear activation function. These are applied to each token independently. +- **Attention**: The attention mechanism learns to update the representation of each token based on other tokens. A score is computed that determines how much influence a given token has on another token. Then the "value vectors" of the tokens are combined based on these scores to determine the update +- **Heads**: Attention heads basically divide the embedding space into multiple subspaces and then apply the attention mechanism to each subspace. This allows the model to learn different types of relationships between tokens. +- **Transformers**: A model that has interwoven attention and feed forward neural networks. The transformer architecture is composed of multiple layers of these blocks. The transformer architecture is the basis for many modern language models. +- **Weight Tying**: This is a technique where two layers of a neural network share the same weights -- this forces them to always compute the same function. diff --git a/models/components/tokenizers/utils.py b/models/components/tokenizers/utils.py index 946dba6f..db88c384 100644 --- a/models/components/tokenizers/utils.py +++ b/models/components/tokenizers/utils.py @@ -6,7 +6,7 @@ import unicodedata from collections import Counter -import hydra # to get the absolute path to the tokenizer +from hydra.utils import to_absolute_path def get_tokenizer_path(tokenizer_type, vocab_size, dataset_name): @@ -16,7 +16,7 @@ def get_tokenizer_path(tokenizer_type, vocab_size, dataset_name): tokenizer_folder = os.path.join( "models", "components", "tokenizers", "tokenizer_models" ) - tokenizer_folder = hydra.utils.to_absolute_path(tokenizer_folder) + tokenizer_folder = to_absolute_path(tokenizer_folder) tokenizer_full_path = os.path.join( tokenizer_folder, f"{tokenizer_type}_{dataset_name}_{vocab_size}.model" ) diff --git a/models/experimental/byte_level/byte_model_shell.py b/models/experimental/byte_level/byte_model_shell.py index b392eb39..0ed2ed4d 100644 --- a/models/experimental/byte_level/byte_model_shell.py +++ b/models/experimental/byte_level/byte_model_shell.py @@ -6,8 +6,7 @@ import torch from models import core_models, embedding_models, model_heads -from models.model_shell import ModelShell - +from models.model_shell import ModelShell class ByteModelShell(ModelShell): @@ -15,6 +14,7 @@ class ByteModelShell(ModelShell): Slight deviation from the standard Model Shell to allow for a re-constructive auxiliary loss to the input. """ + def __init__( self, embedding_model: embedding_models.EmbedderInterface, @@ -37,12 +37,10 @@ def forward(self, token_ids): # to get B, S, H (with pos encoding if necessary) x = self.embedding_model(token_ids) - # calculate the reconstruction loss + # calculate the reconstruction loss logits = self.model_head(x)[0] loss = torch.nn.functional.cross_entropy( - logits.view(-1, logits.size(-1)), - token_ids.view(-1), - ignore_index=257 + logits.view(-1, logits.size(-1)), token_ids.view(-1), ignore_index=257 ) # pass the embeddings through the core model @@ -51,6 +49,4 @@ def forward(self, token_ids): # pass the core model output through the model head x = self.model_head(x)[0] - return x, loss - - + return x, loss diff --git a/models/experimental/hugging_face.py b/models/experimental/hugging_face.py index a7ea4498..ec22f73d 100644 --- a/models/experimental/hugging_face.py +++ b/models/experimental/hugging_face.py @@ -10,10 +10,11 @@ from trainers.base_trainer import BaseTrainer from trainers.dataloader import BaseDataloader -def build_model(model_cfg): - ''' + +def build_hf_model(model_cfg): + """ Helper function to build a model from the huggingface model hub. - ''' + """ ## get the model string model_str = model_cfg["model_string"] @@ -79,7 +80,7 @@ def __init__(self, model_cfg): self.model_cfg = model_cfg model_string = model_cfg["model_string"] self.tokenizer = HFTokenizerWrapper(model_string) - self.embeddings = build_model(model_cfg).get_input_embeddings() + self.embeddings = build_hf_model(model_cfg).get_input_embeddings() def decode(self, token_ids): """ @@ -128,7 +129,7 @@ class HFTransformerCore(torch.nn.Module): def __init__(self, model_cfg): super().__init__() - self.model = build_model(model_cfg = model_cfg) + self.model = build_hf_model(model_cfg=model_cfg) ## freeze the parameters print("Note: Freezing the parameters of the hf_core model.") @@ -140,13 +141,14 @@ def forward(self, x): Calls the huggingface model in question, and returns the last hidden state. """ ## get the hidden states - hidden_states = self.model(inputs_embeds = x, output_hidden_states = True).hidden_states + hidden_states = self.model( + inputs_embeds=x, output_hidden_states=True + ).hidden_states ## return the last hidden state if isinstance(hidden_states, tuple): return hidden_states[-1] - class HFLMHead(torch.nn.Module): """ @@ -155,8 +157,8 @@ class HFLMHead(torch.nn.Module): def __init__(self, model_cfg): super().__init__() - self.lm_head = build_model(model_cfg = model_cfg).get_output_embeddings() - + self.lm_head = build_hf_model(model_cfg=model_cfg).get_output_embeddings() + def forward(self, x): """ Passes the input through the language model head to get logits. diff --git a/models/experimental/next_thought/core_models.py b/models/experimental/next_thought/core_models.py index de760aa6..7d7b4c06 100644 --- a/models/experimental/next_thought/core_models.py +++ b/models/experimental/next_thought/core_models.py @@ -1,15 +1,16 @@ """ The core next-thought model. """ -import torch +import torch class BaselineCoreModel(torch.nn.Module): """ - An extremely simplistic core model for + An extremely simplistic core model for next thought prediction. """ + def __init__(self, model_cfg): super().__init__() @@ -34,12 +35,13 @@ def forward(self, x): x: torch.tensor(B, S, H) """ return self.model(x) - + class Conv1dCoreModel(torch.nn.Module): """ A core model for next thought prediction using Conv1d layers. """ + def __init__(self, model_cfg): super().__init__() @@ -48,9 +50,6 @@ def __init__(self, model_cfg): self.conv2 = torch.nn.Linear(300, 300) self.conv3 = torch.nn.Linear(3, 3) - - - def forward(self, x): """ Pass an input through the model @@ -73,4 +72,4 @@ def forward(self, x): x = x.view(x.size(0), 1600, 3) x = self.conv3(x) x = x.view(x.size(0), 4800) - return x \ No newline at end of file + return x diff --git a/models/experimental/next_thought/embedding_models.py b/models/experimental/next_thought/embedding_models.py index c958fd54..fa80f109 100644 --- a/models/experimental/next_thought/embedding_models.py +++ b/models/experimental/next_thought/embedding_models.py @@ -1,30 +1,29 @@ """ The Embedding model for a VAE style sequence to sequence model. """ -import torch -from models.embedding_models import GenericEmbedder -from models.components.layers.transformer_blocks import GenericTransformerBlock +import torch +from models.components.layers.transformer_blocks import GenericTransformerBlock from models.components.positional_encoding import build_positional_encodings from models.components.tokenizers import build_tokenizer - +from models.embedding_models import EmbedderInterface # import local components from models.experimental.next_thought.layers import AttentionPoolingRemoval - -class HierarchicalEncoder(GenericEmbedder): +class HierarchicalEncoder(EmbedderInterface): """ Accepts an arbitrary length sequence as input, uses the QK^T matrix to, at every layer, - pick the top n-percent of nodes to pool into - a single token (the one paying most attention + pick the top n-percent of nodes to pool into + a single token (the one paying most attention to the other should be pooled into the other token). """ + def __init__(self, model_cfg): - super().__init__(model_cfg=model_cfg) + super().__init__() # build the tokenizer self.tokenizer = build_tokenizer( tokenizer_type=model_cfg["embedder"]["tokenizer_type"], @@ -41,7 +40,6 @@ def __init__(self, model_cfg): # build the positional encodings self.positional_encodings = build_positional_encodings(model_cfg=model_cfg) - self.standard_transformer = torch.nn.ModuleList( [ GenericTransformerBlock( @@ -56,25 +54,23 @@ def __init__(self, model_cfg): self.pooling_transformer = torch.nn.ModuleList( [ - AttentionPoolingRemoval( hidden_size_in=model_cfg["embedder"]["pooling_dims"][i], - hidden_size_out=model_cfg["embedder"]["pooling_dims"][i+1], + hidden_size_out=model_cfg["embedder"]["pooling_dims"][i + 1], num_attention_heads=12, pct_pool_per_layer=model_cfg["embedder"]["pct_pool_per_layer"][i], - ) for i in range(len(model_cfg["embedder"]["pooling_dims"]) - 1) + ) + for i in range(len(model_cfg["embedder"]["pooling_dims"]) - 1) ] ) - def forward(self, token_ids): - # embed the input + # embed the input x = self.embedding(token_ids) - # apply positional encoding + # apply positional encoding x = x + self.positional_encoding(x) - # first pass through normal attention blocks for layer in self.standard: x = layer(x) @@ -84,4 +80,4 @@ def forward(self, token_ids): x = layer(x) # mean pool final representation x = x.mean(dim=-2) - return x \ No newline at end of file + return x diff --git a/models/model_shell.py b/models/model_shell.py index 330f6de3..219decf9 100644 --- a/models/model_shell.py +++ b/models/model_shell.py @@ -8,7 +8,6 @@ from models import core_models, embedding_models, model_heads - class ModelShell(torch.nn.Module): """ Unify the embedding model, core model and LM head @@ -71,7 +70,9 @@ def inference(self, model_input): # check if input is string if isinstance(model_input, str): # use inference function of the embedding model - model_input = self.embedding_model.tokenize_input(model_input, truncate=True, add_eot=False) + model_input = self.embedding_model.tokenize_input( + model_input, truncate=True, add_eot=False + ) x = torch.tensor(model_input, device=self.device, dtype=torch.long).unsqueeze(0) x = self.embedding_model(model_input) @@ -94,9 +95,16 @@ def loglikelihood(self, prefixes, continuations): Returns: ll: torch.tensor(B) """ - total_strings = [f"{prefix} {cont}" for prefix, cont in zip(prefixes, continuations)] - input_tokens = [self.embedding_model.tokenize_input(string, truncate=True) for string in total_strings] - padded_batch, mask = self.embedding_model.pad_batch(input_tokens, direction="right") + total_strings = [ + f"{prefix} {cont}" for prefix, cont in zip(prefixes, continuations) + ] + input_tokens = [ + self.embedding_model.tokenize_input(string, truncate=True) + for string in total_strings + ] + padded_batch, mask = self.embedding_model.pad_batch( + input_tokens, direction="right" + ) input_tensor = torch.tensor(padded_batch, device=self.device, dtype=torch.long) logits, _ = self.forward(input_tensor) logits = logits[:, :-1].reshape(-1, logits.size(-1)) @@ -105,4 +113,4 @@ def loglikelihood(self, prefixes, continuations): mask = mask[:, 1:].reshape(-1).to(ll.device) ll = ll * mask ll = ll.view(input_tensor.size(0), -1).sum(dim=1) - return -ll \ No newline at end of file + return -ll diff --git a/stlm_hf_integration/__init__.py b/stlm_hf_integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/stlm_hf_integration/configs.py b/stlm_hf_integration/configs.py new file mode 100644 index 00000000..336eb6b9 --- /dev/null +++ b/stlm_hf_integration/configs.py @@ -0,0 +1,15 @@ +"""Wrapper from huggingface""" + +from transformers import PretrainedConfig + + +class STLMConfig(PretrainedConfig): + model_type = "stlm" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + +if __name__ == "__main__": + stlmconfig = STLMConfig() + stlmconfig.save_pretrained("stlm") diff --git a/stlm_hf_integration/flatten.py b/stlm_hf_integration/flatten.py new file mode 100644 index 00000000..7219a68a --- /dev/null +++ b/stlm_hf_integration/flatten.py @@ -0,0 +1,228 @@ +"""BAD CODE lol +https://huggingface.co/docs/transformers/en/custom_models#using-a-model-with-custom-code""" + +import os +import re + +import networkx as nx +from networkx.algorithms.dag import topological_sort + + +def collect_python_files(directory, ignore_dirs): + """Recursively collect all Python files in the given directory, excluding specified directories.""" + python_files = set() + for root, _, files in os.walk(directory): + if any(ignore_dir in root for ignore_dir in ignore_dirs): + continue + for file in files: + if file.endswith(".py"): + # relative_path = os.path.relpath(os.path.join(root, file), directory) + # module_name = relative_path.replace(os.sep, ".")[ + # :-3 + # ] # Remove .py extension + python_files.add(os.path.join(root, file)) + return python_files + + +def get_module_name(file_path, base_directory): + """Get the module name from the file path relative to the base directory.""" + relative_path = os.path.relpath(file_path, base_directory) + module_name = relative_path.replace(os.sep, ".")[:-3] # Remove .py extension + return module_name + + +def extract_imported_modules(content): + """Extract imported module names from the file content.""" + imports = re.findall(r"(?:from\s+(\S+)\s+import\s+|\bimport\s+)(\S+)", content) + modules = set() + for imp in imports: + right_imports = imp[1].split(",") + for right_import in right_imports: + imported_module = f"{imp[0]}.{right_import}" if imp[0] else right_import + # Split by '.' to handle submodules and handle direct imports + imported_modules = imported_module.split(".") + for i in range(len(imported_modules)): + module = ".".join(imported_modules[: i + 1]) + modules.add(module) + return modules + + +def build_dependency_graph(file_paths, base_directory): + """Build a dependency graph from the collected Python file paths.""" + graph = nx.DiGraph() + + # Map file paths to module names + module_map = { + file_path: get_module_name(file_path, base_directory) + for file_path in file_paths + } + + # Add nodes to the graph with their corresponding file paths + for file_path, module_name in module_map.items(): + graph.add_node(file_path, module_name=module_name) + with open(file_path, "r") as file: + content = file.read() + + # Extract all imported modules + imported_modules = extract_imported_modules(content) + + # Process each import + for imported_module in imported_modules: + if imported_module in module_map.values(): + # Find the file path of the imported module + imported_file_path = next( + path for path, name in module_map.items() if name == imported_module + ) + # Add edge from imported file to current file + graph.add_edge(imported_file_path, file_path) + return graph + + +def process_imports(file_content, imported_model_modules, all_imports): + """Adjust import statements to work in a single-file context and inline content.""" + lines = file_content.split("\n") + processed_lines = [] + + in_multiline_import = False + multiline_import = [] + + for line in lines: + stripped_line = line.strip() + if in_multiline_import: + multiline_import.append(line) + if stripped_line.endswith(")"): + full_import = "\n".join(multiline_import) + in_multiline_import = False + if full_import.startswith("from models") or full_import.startswith( + "import models" + ): + continue + all_imports.add(full_import) + continue + elif stripped_line.startswith(("from ", "import ")): + if stripped_line.endswith("\\") or stripped_line.endswith("("): + in_multiline_import = True + multiline_import = [line] + continue + if stripped_line.startswith("from models") or stripped_line.startswith( + "import models" + ): + continue + all_imports.add(line) + continue + if "core_models" not in imported_model_modules: + print(imported_model_modules) + pass + line = simplify_reference(line, imported_model_modules) + processed_lines.append(line) + + return "\n".join(processed_lines) + + +def flatten_models_code( + base_directory, output_file, ignore_dirs, include_subfolders=[] +): + """Flatten the models code into a single file, excluding specified directories.""" + ignore_dirs = [os.path.join(base_directory, d) for d in ignore_dirs] + include_subfolders = [os.path.join(base_directory, d) for d in include_subfolders] + python_files = collect_python_files(base_directory, ignore_dirs) + for include_subfolder in include_subfolders: + include_files = collect_python_files(include_subfolder, []) + python_files.update(include_files) + dep_graph = build_dependency_graph(python_files, "") + python_files = list(topological_sort(dep_graph)) + print(python_files) + + all_imports = set() + # collect references to models in the imported modules + imported_models_modules = [] + for file_path in python_files: + imported_models_modules.extend(collect_models_imported_modules(file_path)) + + # First pass: collect imports and file contents + file_contents = [] + processed_files = set() + for file_path in python_files: + if file_path not in processed_files: + with open(file_path, "r") as infile: + file_content = infile.read() + processed_files.add(file_path) + processed_content = process_imports( + file_content, + imported_models_modules, + all_imports, + ) + file_contents.append(f"# {file_path}\n{processed_content}\n\n") + + # These imports: from trainers.base_trainer import BaseTrainer + # from trainers.dataloader import BaseDataloader + # from trainers.utils import load_data + # all need to be replaced with mock functions to make the code run, as they are not present in the models directory, + # but also won't be called anyhow... (hopefully) (if it gets to that point maybe come up with a method that isn't flattening + # our code base into a single file...) (or just remove them from the output file manually) + all_imports.add("class BaseTrainer: pass") + all_imports.add("class BaseDataloader: pass") + all_imports.add("def load_data(): pass") + all_imports.add("def to_absolute_path(*args, **kwargs): pass") + all_imports.remove("from hydra.utils import to_absolute_path") + all_imports.remove("from trainers.base_trainer import BaseTrainer") + all_imports.remove("from trainers.dataloader import BaseDataloader") + all_imports.remove("from trainers.utils import load_data") + + # Second pass: write imports and contents to output file + with open(output_file, "w") as outfile: + if all_imports: + outfile.write("\n".join(sorted(all_imports)) + "\n\n") + for content in file_contents: + outfile.write(content) + + +def collect_models_imported_modules(file_path): + """Collect the modules imported from models in a given file.""" + with open(file_path, "r") as infile: + lines = infile.readlines() + imported_models_modules = [] + for line in lines: + if line.startswith("from models") or line.startswith("import models"): + match = re.search(r"from models(.*) import (.+)", line) + if match: + imported_models_modules.append(match.group(2)) + match = re.search(r"import models(.*)", line) + if match: + imported_models_modules.append(match.group(1)) + for module in imported_models_modules: + if "," in module: + imported_models_modules.extend(module.split(", ")) + return imported_models_modules + + +def simplify_reference(line, imported_models_modules): + """If the line references something imported from models, simplify it.""" + # Create a regex pattern to match fully qualified names of imported modules + pattern = ( + r"\b(" + + "|".join(re.escape(module) for module in imported_models_modules) + + r")\.(\w+)\b" + ) + + def replacement(match): + # The match group 2 is the class or function name + return match.group(2) + + # Replace fully qualified names in the line with just the class/function names + simplified_line = re.sub(pattern, replacement, line) + + return simplified_line + + +# Specify the models directory and the output file +models_directory = "models" +output_file = "stlm_hf_integration/stlm_modelling.py" + +# ignore_dirs = ["experimental", "build_"] # List of directories to ignore +ignore_dirs = [] # use this arg at your peril.. imports *will* break +include_subfolders = ["experimental/next_thought"] # List of subfolders to include + +flatten_models_code(models_directory, output_file, ignore_dirs, include_subfolders) + +print(f"Flattened models code has been written to {output_file}") diff --git a/stlm_hf_integration/model.py b/stlm_hf_integration/model.py new file mode 100644 index 00000000..ef4b9d87 --- /dev/null +++ b/stlm_hf_integration/model.py @@ -0,0 +1,19 @@ +"""Wrapper from huggingface""" + +from stlm_modelling import ModelShell, build_model +from transformers import PreTrainedModel + +from stlm_hf_integration.configs import STLMConfig + + +class STLMModel(PreTrainedModel): + """Big model""" + + config_class = STLMConfig + + def __init__(self, config): + super().__init__(config) + self.model: ModelShell = build_model(config) + + def forward(self, token_ids): + return self.model(token_ids) diff --git a/trainers/base_trainer.py b/trainers/base_trainer.py index d706e589..d4a20a0e 100644 --- a/trainers/base_trainer.py +++ b/trainers/base_trainer.py @@ -1,28 +1,24 @@ """Trainer class for training models with Next Token Prediction""" import time +from contextlib import nullcontext +from copy import deepcopy +from itertools import islice +import numpy as np import torch import wandb from omegaconf import OmegaConf +from torch.nn.parallel import DistributedDataParallel as DDP from torch.profiler import ProfilerActivity, profile, record_function -from copy import deepcopy -from contextlib import nullcontext +from torch.utils.data import SequentialSampler +from torch.utils.data.distributed import DistributedSampler from models import model_shell from trainers import dataloader as train_dataloader from trainers import utils - -from trainers.loss_fn import ( - compute_perplexity -) from trainers.evaluator import train_eval - -import numpy as np -from itertools import islice -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.data.distributed import DistributedSampler -from torch.utils.data import SequentialSampler +from trainers.loss_fn import compute_perplexity from trainers.utils import aggregate_value @@ -40,13 +36,13 @@ def __init__( optimizer, dataloader: train_dataloader.BaseDataloader, loss_fn, - gpu_id, + gpu_id, lr_scheduler=None, dropout_scheduler=None, ) -> None: self.model = model self.DDP_model = DDP(self.model, device_ids=[gpu_id]) - self.gpu_id = gpu_id + self.gpu_id = gpu_id self.optimizer = optimizer self.lr_scheduler = lr_scheduler self.dropout_scheduler = dropout_scheduler @@ -54,22 +50,31 @@ def __init__( self.train_val_dataloaders = {} self.loss_fn = loss_fn self.cfg = cfg - assert self.cfg["trainer"]["training"]["gradient_accumulation_steps"] % torch.cuda.device_count() == 0, "Gradient Accumulation Steps must be divisible by the number of GPUs" - self.gradient_accumulation_steps = cfg["trainer"]["training"][ - "gradient_accumulation_steps" - ] // torch.cuda.device_count() ## divide by number of GPUs to maximise throughput + assert ( + self.cfg["trainer"]["training"]["gradient_accumulation_steps"] + % torch.cuda.device_count() + == 0 + ), "Gradient Accumulation Steps must be divisible by the number of GPUs" + self.gradient_accumulation_steps = ( + cfg["trainer"]["training"]["gradient_accumulation_steps"] + // torch.cuda.device_count() + ) ## divide by number of GPUs to maximise throughput self.scaler = None self.use_wandb = cfg["general"]["logging"]["wandb_log"] self.checkpoint_dir = cfg["general"]["paths"]["checkpoint_dir"] self.cached_sets = {"train": {}, "val": {}} - self.batch_size = cfg["trainer"]["training"]["batch_size"] ## new + self.batch_size = cfg["trainer"]["training"]["batch_size"] ## new # For training, always force the device to be cuda assert torch.cuda.is_available(), "CUDA must be available for training" self.ctx = self._setup_ctx() - if self.use_wandb and self.gpu_id == 0: ## ensures that only the first GPU logs to wandb + if ( + self.use_wandb and self.gpu_id == 0 + ): ## ensures that only the first GPU logs to wandb self._setup_logging() - if cfg.trainer.training.run_profiler and self.gpu_id == 0: ## ensures that only the first GPU runs the profiler + if ( + cfg.trainer.training.run_profiler and self.gpu_id == 0 + ): ## ensures that only the first GPU runs the profiler self.run_profile() raise SystemExit @@ -120,16 +125,16 @@ def _get_dataloader(self, split): ## return the dataloader if it has already been cached if split in self.train_val_dataloaders: return self.train_val_dataloaders[split] - + ## if the dataloader has not been created, create it # set the split data dataset = self.dataloader.split_dataloader(split) # create the dataset dataloader = torch.utils.data.DataLoader( dataset, - batch_size = self.batch_size, - shuffle = False, - num_workers = 0, + batch_size=self.batch_size, + shuffle=False, + num_workers=0, ) ## cache the dataloader @@ -146,11 +151,10 @@ def estimate_performance(self, eval_iters=None): perplexity = {} self.model.eval() for split in ["train", "val"]: - ## initialize the loss, perplexity losses = torch.zeros(eval_iters) perplexities = torch.zeros(eval_iters) - + ## initialize Pytorch's DataLoader dataloader = self._get_dataloader(split) @@ -189,16 +193,22 @@ def estimate_performance(self, eval_iters=None): ## aggregate the loss and perplexity across all GPUs avg_loss = aggregate_value(losses.mean().item(), self.cfg.general.device) loss[split] = avg_loss - avg_perplexity = aggregate_value(perplexities.mean().item(), self.cfg.general.device) + avg_perplexity = aggregate_value( + perplexities.mean().item(), self.cfg.general.device + ) perplexity[split] = avg_perplexity evaluator_results = {} for evaluator in self.cfg.trainer["eval"]: - evaluator_results[evaluator["evaluator"]] = train_eval(evaluator, self.model) + evaluator_results[evaluator["evaluator"]] = train_eval( + evaluator, self.model + ) # recurse over metrics to prepend the evaluator name as a prefix relabeled_results = {} for metric in evaluator_results[evaluator["evaluator"]]: - relabeled_results[f"{evaluator['evaluator']}/{metric}"] = evaluator_results[evaluator["evaluator"]][metric] + relabeled_results[ + f"{evaluator['evaluator']}/{metric}" + ] = evaluator_results[evaluator["evaluator"]][metric] evaluator_results[evaluator["evaluator"]] = relabeled_results self.model.train() return loss, perplexity, evaluator_results @@ -250,10 +260,10 @@ def run_profile(self): ) as prof: for i in range(10): if i <= 3: - self._run_step() ## set the 'epoch' to ensure shuffle + self._run_step() ## set the 'epoch' to ensure shuffle else: with record_function("_run_step"): - self._run_step() ## set the 'epoch' to ensure shuffle + self._run_step() ## set the 'epoch' to ensure shuffle # place profile in dictionary backwards_prof = prof.key_averages().table(sort_by="self_cpu_time_total") print(backwards_prof) @@ -310,11 +320,9 @@ def run_training_loop(self): f"step {iter_num}: train perplexity {perplexities['train']:.4f}," f" val perplexity {perplexities['val']:.4f}" ) - print( - f"step {iter_num}: benchmark results {benchmark_results}" - ) + print(f"step {iter_num}: benchmark results {benchmark_results}") - if self.gpu_id == 0: ## ensure only the first GPU logs + if self.gpu_id == 0: ## ensure only the first GPU logs if self.use_wandb: wandb.log( { @@ -325,21 +333,18 @@ def run_training_loop(self): "dropout": dropout, "train/perplexity": perplexities["train"], "val/perplexity": perplexities["val"], - **{ - k: v - for k, v in benchmark_results.items() - }, + **{k: v for k, v in benchmark_results.items()}, } ) # save checkpoints if ( not iter_num % self.cfg.trainer.training.checkpoint_interval and iter_num > 0 - and self.gpu_id == 0 ## ensure only the first GPU prints + and self.gpu_id == 0 ## ensure only the first GPU prints ): self._save_model(iter_num) - loss = self._run_step() ## set the 'epoch' to ensure shuffle + loss = self._run_step() ## set the 'epoch' to ensure shuffle end_time = time.time() if not iter_num % self.cfg.trainer.training.log_interval and iter_num > 0: lossf = loss.item() * self.gradient_accumulation_steps @@ -351,7 +356,9 @@ def run_training_loop(self): lossf = aggregate_value(lossf, self.cfg.general.device) ## print and log the result only on the first GPU after aggregation - print(f"All GPU(s): step {iter_num}: loss {lossf:.4f}, lr {lr:.1e}, dt {end_time-start_time:.1f}s") + print( + f"All GPU(s): step {iter_num}: loss {lossf:.4f}, lr {lr:.1e}, dt {end_time-start_time:.1f}s" + ) if self.gpu_id == 0 and self.use_wandb: wandb.log( { @@ -362,7 +369,7 @@ def run_training_loop(self): } ) # save the final model - if self.gpu_id == 0: ## ensure only the first GPU saves the model + if self.gpu_id == 0: ## ensure only the first GPU saves the model self._save_model(iter_num) def train(self, seed=42):