chore: import upstream snapshot with attribution

labml_nn/neox/__init__.py

@@ -0,0 +1,32 @@
+"""
+---
+title: GPT-NeoX
+summary: >
+    Simple GPT-NeoX implementation
+---
+
+# GPT-NeoX
+
+This is a simple implementation of [Eleuther GPT-NeoX](https://arxiv.org/abs/2204.06745) for inference and fine-tuning.
+
+
+* [Model definition](model.html)
+* [Tokenizer](tokenizer.html)
+* [Checkpoint downloading and loading helpers](checkpoint.html)
+* [Utilities](utils/index.html)
+* [LLM.int8() quantization](utils/llm_int8.html)
+
+### [Samples](samples/__init__.py)
+
+* [Generating text](samples/generate.html)
+* [Fine-tuning the biases with pipeline-parallel](samples/finetune.html)
+* [Generating text with LLM.int8()](samples/llm_int8.html)
+
+### [Evaluation](evaluation/__init__.py)
+
+* [Evaluating half precision model on a single GPU](evaluation/half_precision.html)
+* [Evaluating LLM.int8() model](evaluation/llm_int8.html)
+
+**Official [Eleuther](https://www.eleuther.ai)
+GPT-NoeX is source code is available at [eleutherai/gpt-neox](https://github.com/eleutherai/gpt-neox).**
+"""

labml_nn/neox/checkpoint.py

@@ -0,0 +1,164 @@
+"""
+---
+title: GPT-NeoX Checkpoints
+summary: >
+    Code to download checkpoints and helpers to load them.
+---
+
+# GPT-NeoX Checkpoints
+
+"""
+from pathlib import Path
+from typing import Dict, Union, Tuple, Optional
+
+import torch
+from torch import nn
+
+from labml import monit, lab, logger
+from labml.logger import Text, inspect
+from labml.utils.download import download_file
+
+# Parent url
+CHECKPOINTS_URL = 'https://mystic.the-eye.eu/public/AI/models/GPT-NeoX-20B/slim_weights/'
+
+_CHECKPOINTS_DOWNLOAD_PATH: Optional[Path] = None
+
+
+# Download path
+def get_checkpoints_download_path():
+    global _CHECKPOINTS_DOWNLOAD_PATH
+
+    if _CHECKPOINTS_DOWNLOAD_PATH is not None:
+        return _CHECKPOINTS_DOWNLOAD_PATH
+
+    _CHECKPOINTS_DOWNLOAD_PATH = lab.get_data_path() / 'neox_fast' / 'slim_weights'
+    if not _CHECKPOINTS_DOWNLOAD_PATH.exists():
+        _CHECKPOINTS_DOWNLOAD_PATH = lab.get_data_path() / 'neox' / 'slim_weights'
+    inspect(neox_checkpoint_path=_CHECKPOINTS_DOWNLOAD_PATH)
+
+    return _CHECKPOINTS_DOWNLOAD_PATH
+
+
+def get_files_to_download(n_layers: int = 44):
+    """
+    ### Get files to download
+
+    :return: a list of files to be downloaded
+    """
+    layers = (
+        # Embedding layer
+            [0] +
+            # Transformer layers
+            list(range(2, 2 + n_layers)) +
+            # Final normalization layer and readout layer
+            [47, 48]
+    )
+
+    return (
+        # Vocabulary and configs
+            ['20B_tokenizer.json', 'configs/20B.yml', 'latest'] +
+            # Layer checkpoints
+            [f'global_step150000/layer_{i :02d}-model_{p :02d}-model_states.pt' for i in layers for p in range(2)] +
+            # Empty states (not used)
+            [f'global_step150000/mp_rank_{i :02d}_model_states.pt' for i in range(8)]
+    )
+
+
+def download(n_layers: int = 44):
+    """
+    ## Download all checkpoint files
+    """
+
+    # Get files to download
+    files = get_files_to_download(n_layers)
+
+    # Iterate
+    for i, f in monit.enum('Download All', files):
+        # Log
+        logger.log(['Downloading ', (f'{i + 1 :3d}/{len(files)}', Text.meta), ': ', (f, Text.value)])
+        # Download
+        download_file(CHECKPOINTS_URL + f, get_checkpoints_download_path() / f)
+
+
+def load_checkpoint_files(files: Tuple[str, str]):
+    """
+    ### Load a pair of checkpoint files
+
+    :param files: pair of files to load
+    :return: the loaded parameter tensors
+    """
+    checkpoint_path = get_checkpoints_download_path() / 'global_step150000'
+    with monit.section('Load checkpoint'):
+        data = [torch.load(checkpoint_path / f) for f in files]
+
+    return data
+
+
+def merge_params_dim_0(param: Union[nn.Parameter, torch.Tensor], key: str, p1: Dict[str, torch.Tensor],
+                       p2: Dict[str, torch.Tensor]):
+    """
+    ### Load a parameter by merging the partitions along first dimension
+
+    :param param: is the parameter
+    :param key: is the name of the parameter
+    :param p1: first partition dictionary
+    :param p2: second partition dictionary
+    """
+    w1, w2 = p1[key], p2[key]
+    param.data[:w1.shape[0]] = w1
+    param.data[w1.shape[0]:] = w2
+
+
+def merge_params_dim_1(param: Union[nn.Parameter, torch.Tensor], key: str, p1: Dict[str, torch.Tensor],
+                       p2: Dict[str, torch.Tensor]):
+    """
+    ### Load a parameter by merging the partitions along second dimension
+
+    :param param: is the parameter
+    :param key: is the name of the parameter
+    :param p1: first partition dictionary
+    :param p2: second partition dictionary
+    """
+    w1, w2 = p1[key], p2[key]
+    param.data[:, :w1.shape[1]] = w1
+    param.data[:, w1.shape[1]:] = w2
+
+
+def merge_params_duplicate(param: Union[nn.Parameter, torch.Tensor], key: str, p1: Dict[str, torch.Tensor],
+                           p2: Dict[str, torch.Tensor]):
+    """
+    ### Load an un-partitioned parameter
+
+    This does a sanity check to make use both partitions are the same
+
+    :param param: is the parameter
+    :param key: is the name of the parameter
+    :param p1: first partition dictionary
+    :param p2: second partition dictionary
+    """
+    w1, w2 = p1[key], p2[key]
+
+    diff = sum((w1 - w2) ** 2).item()
+    assert diff < 1e-4, f'The partitions do not match: {key}'
+
+    param.data[:] = (w1 + w2) / 2.
+
+
+def merge_params_sum(param: Union[nn.Parameter, torch.Tensor], key: str, p1: Dict[str, torch.Tensor],
+                     p2: Dict[str, torch.Tensor]):
+    """
+    ### Load biases that are partitioned which gets added on reduce
+
+    :param param: is the parameter
+    :param key: is the name of the parameter
+    :param p1: first partition dictionary
+    :param p2: second partition dictionary
+    """
+    w1, w2 = p1[key], p2[key]
+
+    param.data[:] = w1 + w2
+
+
+#
+if __name__ == '__main__':
+    download()

labml_nn/neox/evaluation/__init__.py

@@ -0,0 +1,262 @@
+"""
+---
+title: Evaluation
+summary: >
+    Code to evaluate the model on NLP tasks through lm-evaluation-harness
+---
+
+# Evaluation
+
+This is the code to test the model on
+[EleutherAI/lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness).
+
+* [Evaluating half precision model on a single GPU](half_precision.html)
+"""
+import math
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from lm_eval import tasks, evaluator, utils
+from lm_eval.base import BaseLM
+from tokenizers import Tokenizer
+from torch import nn
+from tqdm import tqdm
+
+from labml import monit
+from labml_nn.neox.tokenizer import get_tokenizer
+
+
+class EvalHarnessAdapter(BaseLM):
+    """
+    ## Evaluation Harness Adapter
+
+    This is based on the [adapter from EleutherAI/gpt-neox](https://github.com/EleutherAI/gpt-neox/blob/main/eval_tasks/eval_adapter.py)
+    """
+
+    def __init__(self, tokenizer: Tokenizer, vocab_size: int, batch_size: int):
+        """
+        :param tokenizer: is the [Huggingface Tokenizer](huggingface/tokenizers)
+        :param vocab_size: is the size of the vocabulary
+         (this differs from the tokenizer vocab size since neox adds some extra to make the embedding layer
+         model parallel.)
+        :param batch_size: is the batch size
+        """
+        super().__init__()
+        self.tokenizer = tokenizer
+        self._eot_token_id = self.tokenizer.token_to_id("<|endoftext|>")
+        self._vocab_size = vocab_size
+
+        self._batch_size = batch_size
+
+    @property
+    def device(self):
+        raise RuntimeError()
+
+    @property
+    def vocab_size(self):
+        """Size of the vocabulary"""
+        return self._vocab_size
+
+    @property
+    def eot_token_id(self):
+        """End-of-text token"""
+        return self._eot_token_id
+
+    @property
+    def max_length(self):
+        """Maximum sequence length"""
+        return 2048
+
+    @property
+    def max_gen_toks(self):
+        """Maximum number of tokens to generate"""
+        return 128
+
+    @property
+    def batch_size(self):
+        """
+        Batch size
+        """
+        return self._batch_size
+
+    def tok_encode(self, string: str):
+        """
+        Encode a given text
+        """
+        return self.tokenizer.encode(string).ids
+
+    def tok_decode(self, tokens: List[int]):
+        """
+        Decode text from token ids
+        """
+        return self.tokenizer.decode(tokens)
+
+    def _model_call(self, inps: torch.Tensor):
+        raise NotImplementedError
+
+    def _model_generate(self, context, max_length, eos_token_id):
+        raise RuntimeError()
+
+    def greedy_until(self, requests):
+        raise RuntimeError()
+
+    @torch.no_grad()
+    def _loglikelihood_tokens(self, requests, disable_tqdm=False):
+        """
+        ### Get log-likelihoods of the next tokens
+
+        :param requests: List of requests containing the context and the expected continuation.
+        :param disable_tqdm: If True, disable tqdm progress bar.
+        """
+
+        # For results
+        res = []
+
+        # Reorder the requests in the descending order of the lengths,
+        # so that sequences with similar lengths are close
+        def _collate(x):
+            toks = x[1] + x[2]
+            return -len(toks), tuple(toks)
+
+        reord = utils.Reorderer(requests, _collate)
+
+        # Loop through requests with `batch_size` number of requests at a time
+        for chunk in utils.chunks(tqdm(reord.get_reordered(), disable=disable_tqdm), self.batch_size):
+            # To store the inputs for the batch
+            inps = []
+            # The continuations for the batch
+            continuations = []
+            # Lengths of the input sequences
+            inplens = []
+            # Padded length for the batch
+            padded_length = None
+            # Loop through each request in the chunk and collect them into PyTorch tensors with paddings
+            for _, context_enc, continuation_enc in chunk:
+                # Concatenate the context and continuation
+                inp = context_enc + continuation_enc
+                # Truncate from left if the size exceeds the `max_length`
+                inp = inp[-(self.max_length + 1):]
+                # Remove final token
+                inp = inp[:-1]
+                # Create a tensor
+                inp = torch.tensor(inp, dtype=torch.long)
+                # Input length
+                inplen = inp.shape[0]
+
+                # Determine the padded length.
+                # Shorter sequences will get padded.
+                if padded_length is None:
+                    padded_length = int(math.ceil(inplen / 32)) * 32
+                # padded_length = padded_length if padded_length is not None else inplen
+
+                # Padding
+                padding = torch.zeros(padded_length - inplen, dtype=torch.long)
+
+                # Add padding
+                inp = torch.cat([inp, padding], dim=0)
+
+                inps.append(inp)
+                continuations.append(continuation_enc)
+                inplens.append(inplen)
+
+            # Get model logits
+            logits = self._model_call(torch.stack(inps))
+
+            # Get log softmaxes
+            multi_logits = F.log_softmax(logits, dim=-1)
+
+            # Loop through the input/output pairs of the batch
+            for logits, inplen, cont_toks in zip(multi_logits, inplens, continuations):
+                # Get number of predicted tokens
+                contlen = len(cont_toks)
+                # Get logits of those
+                logits = logits[inplen - contlen: inplen]
+                # Get the tokens with the highest probabilities
+                greedy_tokens = logits.argmax(dim=-1)
+                # Get the target tokens
+                cont_toks = torch.tensor(cont_toks, dtype=torch.long).to(logits.device)
+                # Whether there's an exact match
+                max_equal = (greedy_tokens == cont_toks).all()
+                # Log-likelihoods of the target tokens
+                logits = torch.gather(logits, 1, cont_toks[:, None])
+                # Add the total log-likelihoods and whether there was a match to the results
+                res.append((float(logits.sum()), bool(max_equal)))
+
+        # Re-order and return results
+        return reord.get_original(res)
+
+    @torch.no_grad()
+    def run_eval(self, name: str, eval_tasks: List[str]):
+        """
+        ### Run given evaluations
+        """
+
+        # Run [EleutherAI/lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) evaluator
+        results = evaluator.evaluate(lm=self, task_dict=tasks.get_task_dict(eval_tasks))
+
+        # Add configs
+        results["config"] = {
+            "name": name,
+        }
+
+        #
+        return results
+
+
+class NoeXEvalHarnessAdapter(EvalHarnessAdapter):
+    """
+    ## Evaluation Harness Adapter
+
+    This is based on the [adapter from EleutherAI/gpt-neox](https://github.com/EleutherAI/gpt-neox/blob/main/eval_tasks/eval_adapter.py)
+    """
+
+    def __init__(self, model: nn.Module, tokenizer: Tokenizer, vocab_size: int, batch_size: int, device: torch.device):
+        """
+        :param model: is model
+        :param tokenizer: is the [Huggingface Tokenizer](huggingface/tokenizers)
+        :param vocab_size: is the size of the vocabulary
+         (this differs from the tokenizer vocab size since neox adds some extra to make the embedding layer
+         model parallel.)
+        :param batch_size: is the batch size
+        :param device: is the device of the model
+        """
+        super().__init__(tokenizer, vocab_size, batch_size)
+        self.model = model
+        self._device = device
+
+    def _model_call(self, inps: torch.Tensor):
+        """
+        Call the model
+        """
+        return self.model(inps.to(self._device))
+
+
+def run_eval_harness(model: nn.Module, name: str, eval_tasks: List[str], device: torch.device, batch_size: int = 8):
+    """
+    ## Run evaluation harness with a given model
+    """
+
+    # Load the tokenizer
+    with monit.section('Load tokenizer'):
+        tokenizer = get_tokenizer()
+
+    # All tasks if nothing is specified
+    if not eval_tasks:
+        eval_tasks = [
+            "anli_r1",
+            "anli_r2",
+            "anli_r3",
+            "hellaswag",
+            "lambada",
+            "piqa",
+            "winogrande",
+            "wsc",
+            "mathqa",
+        ]
+
+    # Create the adapter
+    adapter = NoeXEvalHarnessAdapter(model, tokenizer, 50_432, batch_size, device)
+
+    # Run
+    return adapter.run_eval(name, eval_tasks)

labml_nn/neox/evaluation/half_precision.py

@@ -0,0 +1,48 @@
+"""
+---
+title: Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+summary: >
+     Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+---
+
+#  Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+
+This code evaluate [GPT-NeoX](../index.html) using, on a suite of tasks.
+"""
+import argparse
+
+import torch
+from torch import nn
+
+from labml_nn.neox.evaluation import run_eval_harness
+from labml_nn.neox.model import LayerGenerator
+
+
+def main():
+    # Argument parser
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--flash", action='store_true', help="whether to use Flash Attention")
+
+    opt = parser.parse_args()
+
+    # Device
+    device = torch.device('cuda:0')
+    # Load layers
+    layers = list(LayerGenerator(is_clone_layers=True,
+                                 filter_layers=None,
+                                 dtype=torch.float16,
+                                 device=device,
+                                 is_flash_attention=opt.flash,
+                                 ).load())
+
+    # Create `nn.Sequential` model
+    model = nn.Sequential(*layers)
+
+    # Run [evaluation harness](index.html)
+    print(run_eval_harness(model, 'half_precision', ['lambada'], device))
+
+
+#
+if __name__ == '__main__':
+    main()

labml_nn/neox/evaluation/llm_int8.py

@@ -0,0 +1,54 @@
+"""
+---
+title: Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+summary: >
+     Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+---
+
+#  Evaluate GPT-NeoX using LLM.int8() quantization on test suite
+
+This code evaluate [GPT-NeoX](../index.html) using [LLM.int8() quantization](../utils/llm_int8.html),
+on a suite of tasks.
+"""
+
+import torch
+from torch import nn
+
+from labml import monit
+from labml_nn.neox.evaluation import run_eval_harness
+from labml_nn.neox.model import LayerGenerator
+
+
+def main():
+    # Device
+    device = torch.device('cuda:0')
+
+    # Load layers in float16 into CPU. We convert the layers to int8 later, because doing that
+    # on the fly after loading layers to GPU causes CUDA memory fragmentation
+    # (about 3GB memory can get lost due to fragmentation).
+    layer_generator = LayerGenerator(is_clone_layers=True,
+                                     dtype=torch.float16,
+                                     device=torch.device('cpu'),
+                                     )
+    # Load layers
+    layers = list(layer_generator.load())
+
+    # This reduces CUDA memory fragmentation
+    for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
+        layer_generator.post_load_prepare(layer,
+                                          device=device,
+                                          is_llm_int8=True,
+                                          llm_int8_threshold=6.0,
+                                          )
+        layer.to(device)
+
+    # Create `nn.Sequential` model
+    model = nn.Sequential(*layers)
+
+    # Run [evaluation harness](index.html)
+    print(run_eval_harness(model, 'half_precision', [], device))
+
+
+#
+if __name__ == '__main__':
+    main()

labml_nn/neox/model.py

@@ -0,0 +1,717 @@
+"""
+---
+title: GPT-NeoX Model Definition
+summary: >
+    This is the model definition of GPT-NeoX.
+---
+
+# GPT-NeoX Model
+
+Here is the code for layers of GPT-NeoX model and the code to load
+20B checkpoint.
+
+The method `load_state` in the layers load the checkpoints of that layer.
+The checkpoint loading helpers are on [`checkpoint.py`](checkpoint.html)
+"""
+import copy
+import math
+from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
+
+import torch
+from torch import nn
+from torch.cuda.amp import autocast
+
+from labml import monit, logger
+from labml.logger import Text
+from labml_nn.neox import checkpoint
+from labml_nn.neox.utils.cache import get_cache
+
+
+class NeoXModule(nn.Module):
+    def load_state(self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]):
+        pass
+
+
+class Embedding(NeoXModule):
+    """
+    ## Embedding layer
+
+    This is a standard embeddings layer with code to load the checkpoint.
+    """
+
+    def __init__(self, n_vocab: int = 50_432, n_hidden: int = 6_144):
+        """
+        :param n_vocab: is the size of the vocabulary
+        :param n_hidden: is the size of the embeddings
+        """
+        super().__init__()
+
+        self.emb = nn.Embedding(n_vocab, n_hidden)
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: are the token ids of shape `[batch_size, seq_len]`
+        """
+        return self.emb(x)
+
+    def load_state(self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]):
+        """
+        Code to load the checkpoint
+        """
+        with monit.section('Load embedding layer'):
+            checkpoint.merge_params_dim_0(self.emb.weight, 'word_embeddings.weight', p1, p2)
+
+
+class RoPE(nn.Module):
+    """
+    ## Rotary Positional Embeddings
+
+    GPT-NeoX uses [rotary positional embeddings (RoPE)](https://arxiv.org/abs/2104.09864).
+
+    WE have annotated implementation of RoPE [here](https://nn.labml.ai/transformers/rope/index.html)
+    with more notes the theory.
+    """
+
+    def __init__(self, d_rope: int, base: float = 10_000.):
+        """
+        :param d_rope: is the number of features for RoPE embeddings
+        :param base: is the base for $\theta_i = 10000^{\frac{2(i-1)}{d}}$, which defaults to $10000$
+        """
+        super().__init__()
+
+        # To store $\theta_i$ for the features
+        self.theta = None
+        # Cache $\cos m\theta_i$ and $\sin m\theta_i$
+        self.cos_cached = None
+        self.sin_cached = None
+
+        # Base for $\theta_i = 10000^{\frac{2(i-1)}{d}}$
+        self.base = base
+        # Number of features for RoPE
+        self.d_rope = d_rope
+
+    @staticmethod
+    def rotate_half(x: torch.Tensor):
+        """
+        ### Rotate the features
+
+        $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., -x^{(\frac{d}{2})}]$
+        """
+        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def forward(self, x: torch.Tensor, offset: int = 0):
+        """
+        :param x: has shape `[..., seq, n_heads, d_k]`
+        :param offset: is the starting position of `x`. This is $\gt 0$ when we have
+        cached the keys and queries of previous positions
+        """
+
+        # Get the actual sequence length
+        seq_len = x.shape[-3] + offset
+
+        # Initialize $\theta$
+        if self.theta is None:
+            #  $\theta_i = 10000^{\frac{2(i-1)}{d}}$
+            theta = 1.0 / (self.base ** (torch.arange(0, self.d_rope, 2).float() / self.d_rope))
+            self.theta = theta.to(x.device).to(x.dtype)
+
+        # Initialize $\cos m\theta_i$ and $\sin m\theta_i$ cache
+        if (
+                self.cos_cached is None or
+                seq_len > self.cos_cached.shape[1] or
+                self.cos_cached.device != x.device or
+                self.cos_cached.dtype != x.dtype
+        ):
+            # Get position indexes $m$
+            seq_idx = torch.arange(seq_len, device=x.device).type_as(self.theta)
+            # $m \theta_i$
+            idx_theta = torch.einsum("s,d->sd", seq_idx, self.theta)
+            # Concatenate so that for row $m$ we have
+            #
+            # $$[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$$
+            idx_theta2 = torch.cat((idx_theta, idx_theta), dim=-1).to(x.device)
+
+            # Calculate $\cos m\theta_i$ and $\sin m\theta_i$ in fp32
+            with autocast(enabled=False):
+                idx_theta2 = idx_theta2.float()
+                # Add head dimension
+                self.cos_cached = idx_theta2.cos()[:, None, :]
+                self.sin_cached = idx_theta2.sin()[:, None, :]
+
+            # Cache them
+            self.cos_cached = self.cos_cached.to(x.dtype)
+            self.sin_cached = self.sin_cached.to(x.dtype)
+
+        # Split the features. We apply RoPE to only `d_rope` features
+        x_rope, x_pass = x[..., :self.d_rope], x[..., self.d_rope:]
+
+        # Get the sin and cos values from the cache
+        cos, sin = self.cos_cached[offset: seq_len], self.sin_cached[offset: seq_len]
+
+        # RoPE embeddings
+        #
+        # \begin{align}
+        # \begin{pmatrix}
+        # x^{(i)}_m \cos m \theta_i - x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
+        # x^{(i + \frac{d}{2})}_m \cos m\theta_i + x^{(i)}_m \sin m \theta_i \\
+        # \end{pmatrix} \\
+        # \end{align}
+        #
+        # for $i \in {1, 2, ..., \frac{d}{2}}$
+        x_rope = (x_rope * cos) + (self.rotate_half(x_rope) * sin)
+
+        # Concatenate with features that didn't get RoPE embeddings
+        return torch.cat((x_rope, x_pass), dim=-1)
+
+
+class AttentionLayer(nn.Module):
+    """
+    ## Attention layer
+    """
+
+    def __init__(self, n_hidden: int = 6_144, n_heads: int = 64, rope_percentage: float = 0.25,
+                 mask_fill: float = -10_000.0, *, is_flash_attention: bool = False):
+        """
+        :param n_hidden: the number of features in embeddings
+        :param n_heads: the number of attention heads
+        :param rope_percentage: percentage of features to add RoPE embeddings
+        :param mask_fill: masking fill value for attention matrix
+        :param is_flash_attention: specifies whether to use
+            [FlashAttention](https://github.com/HazyResearch/flash-attention)
+        """
+        super().__init__()
+
+        self.n_heads = n_heads
+        self.mask_fill = mask_fill
+
+        # Linear layer for query, key and value
+        self.qkv_lin = nn.Linear(n_hidden, n_hidden * 3)
+        # Final linear layer
+        self.output = nn.Linear(n_hidden, n_hidden)
+
+        # Number of features per head
+        d_k = n_hidden // n_heads
+        # RoPE embedding module
+        self.rope = RoPE(int(d_k * rope_percentage))
+
+        # Attention scaling factor
+        self.scale = 1 / math.sqrt(d_k)
+
+        # To cache causal mask
+        self.causal_mask = None
+
+        # Attention softmax module
+        self.softmax = nn.Softmax(dim=-2)
+
+        # [FlashAttention](https://github.com/HazyResearch/flash-attention)
+        if is_flash_attention:
+            try:
+                from flash_attn.flash_attention import FlashAttention
+                self.flash_attention = FlashAttention()
+            except ImportError:
+                logger.log('Install flash attention github.com/HazyResearch/flash-attention. '
+                           'Falling back to normal attention', Text.warning)
+                self.flash_attention = None
+        else:
+            self.flash_attention = None
+
+    def _get_mask(self, attn: torch.Tensor):
+        """
+        #### Calculate the causal mask
+
+        * `attn` has shape [batch_size, query_seq_len, key_seq_len, n_heads]
+        """
+
+        # Query and key lengths
+        nq, nk = attn.shape[1:3]
+
+        # Create mask
+        if (
+                self.causal_mask is None or
+                self.causal_mask.shape[0] != nq or
+                self.causal_mask.shape[1] != nk or
+                self.causal_mask.device != attn.device
+        ):
+            self.causal_mask = torch.triu(attn.new_ones([nq, nk], dtype=torch.bool), 1 + nk - nq)
+
+        # Return from cache
+        return self.causal_mask[None, :, :, None]
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: has shape `[batch_size, seq_len, n_hidden]`
+        """
+        # Get query, key and value embeddings (all concatenated).
+        # The last dimension size will change from n_hidden -> `3 x n_hidden`
+        qkv = self.qkv_lin(x)
+
+        # Split into heads by changing the shape to `[batch_size, seq_len, n_heads, 3 * d_k]`
+        qkv = qkv.view(*qkv.shape[:-1], self.n_heads, -1)
+        # Split into query, key and value each of shape `[batch_size, seq_len, n_heads, 3 * d_k]`
+        q, k, v = torch.split(qkv, qkv.shape[-1] // 3, dim=-1)
+
+        # If we are caching the states of previous tokens
+        if get_cache().get('use_cache', False):
+            # Get the state id's. We use to retrieve previous states and store the next states
+            prev_state_id, next_state_id = get_cache().get('state_ids')
+            # If there's cache
+            if prev_state_id is not None:
+                # Get the past keys and values. These will have shape `[batch_size, prev_seq_len, n_heads, d_k]`
+                k_past, v_past = get_cache().pop(f'attn_kv_{prev_state_id}')
+                # Offset of the current embeddings
+                offset = k_past.shape[1]
+
+                # Add RoPE embeddings
+                q = self.rope(q, offset=offset)
+                k = self.rope(k, offset=offset)
+
+                # Concatenate the past
+                k = torch.cat([k_past, k], dim=1)
+                v = torch.cat([v_past, v], dim=1)
+            else:
+                # Add RoPE embeddings
+                q = self.rope(q)
+                k = self.rope(k)
+
+            # Save the current state
+            get_cache().push(f'attn_kv_{next_state_id}', (k, v))
+        else:
+            # No cache - simply add RoPE embeddings
+            q = self.rope(q)
+            k = self.rope(k)
+
+        # Use flash attention
+        if self.flash_attention is not None and q.shape[1] == k.shape[1] and q.shape[-1] <= 128:
+            output = self.compute_flash_attention(q, k, v)
+        # Otherwise, use normal attention
+        else:
+            output = self.compute_attention(q, k, v)
+
+        # Reshape from `[batch_size, seq_len, n_heads, d_k] to `[batch_size, seq_len, n_hidden]`
+        output = output.reshape(*x.shape)
+
+        # Final linear layer
+        return self.output(output)
+
+    def compute_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+        # Stack them into shape `[batch_size, seq_len, 3, n_heads, d_k]`
+        qkv = torch.stack((q, k, v), dim=2)
+        d_k = qkv.shape[-1]
+        if d_k <= 32:
+            pad = 32 - d_k
+        elif d_k <= 64:
+            pad = 64 - d_k
+        elif d_k <= 128:
+            pad = 128 - d_k
+        else:
+            raise ValueError(f'Head size {d_k} too large for flash attention')
+
+        if pad > 0:
+            qkv = torch.cat((qkv, qkv.new_zeros(*qkv.shape[:-1], pad)), dim=-1)
+
+        output, _ = self.flash_attention(qkv, causal=True)
+        # The output is of shape `[batch_size, seq_len, n_heads, d_k + padding]`
+        output = output[:, :, :, :d_k]
+
+        return output
+
+    def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+        # Disable auto-casting to fp16 for attention computation
+        with autocast(enabled=False):
+            if q.dtype == torch.float16:
+                # Convert to fp32 if the current dtype is fp16
+                attn = torch.einsum('bihk,bjhk->bijh', q.float(), k.float())
+            else:
+                # Do not cast for bfloat
+                attn = torch.einsum('bihk,bjhk->bijh', q, k)
+
+            # Scale attention
+            attn = attn * self.scale
+
+            # Get causal mask
+            mask = self._get_mask(attn)
+            # Apply mask
+            attn.masked_fill_(mask, self.mask_fill)
+
+            # Attention softmax
+            attn = self.softmax(attn)
+
+        # Get attention weighted values
+        output = torch.einsum('bijh,bjhk->bihk', attn.to(v.dtype), v)
+
+        return output
+
+
+class FFNLayer(nn.Module):
+    """
+    ## Feedforward Network
+    """
+
+    def __init__(self, n_hidden: int = 6_144, d_ff: int = 0):
+        """
+        :param n_hidden: is the embedding size
+        """
+        super().__init__()
+
+        if not d_ff:
+            d_ff = n_hidden * 4
+
+        # Expansion linear layer
+        self.dense_h_h4 = nn.Linear(n_hidden, d_ff)
+        # GELU activation
+        self.activation = nn.GELU()
+        # Contraction linear layer
+        self.dense_h4_h = nn.Linear(d_ff, n_hidden)
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: has shape `[batch_size, seq_len, n_hidden]`
+        """
+        x = self.dense_h_h4(x)
+        x = self.activation(x)
+        x = self.dense_h4_h(x)
+
+        return x
+
+
+class TransformerLayer(NeoXModule):
+    """
+    ## Transformer Layer
+    """
+
+    def __init__(self, n_hidden: int = 6_144, n_heads: int = 64, *, is_flash_attention: bool = False):
+        """
+        :param n_hidden: is the embedding size
+        :param n_heads: is the number of heads
+        :param is_flash_attention: specifies whether to use
+            [FlashAttention](https://github.com/HazyResearch/flash-attention)
+
+        *Out implementation doesn't include dropout*.
+        """
+        super().__init__()
+
+        # Layer normalization before attention
+        self.pre_ln_attn = nn.LayerNorm(n_hidden)
+        # Layer normalization before FFN
+        self.pre_ln_ffn = nn.LayerNorm(n_hidden)
+
+        # Attention layer
+        self.attention = AttentionLayer(n_hidden, n_heads, is_flash_attention=is_flash_attention)
+        # FFN layer
+        self.ffn = FFNLayer(n_hidden)
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: are the embeddings of shape `[batch_size, seq_len, n_hidden]`
+        """
+
+        # Residual connection
+        residual = x
+        # NeoX runs attention and feedforward network in parallel
+        attn = self.attention(self.pre_ln_attn(x))
+        ffn = self.ffn(self.pre_ln_ffn(x))
+        # Add them and the residual connection
+        return attn + ffn + residual
+
+    def load_state(self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]):
+        """
+        Code to load the checkpoint
+        """
+        with monit.section('Load transformer layer'):
+            # Attention output transform
+            checkpoint.merge_params_sum(self.attention.output.bias, 'attention.dense.bias', p1, p2)
+            checkpoint.merge_params_dim_1(self.attention.output.weight, 'attention.dense.weight', p1, p2)
+
+            # Attention query, key and value transform
+            checkpoint.merge_params_dim_0(self.attention.qkv_lin.bias, 'attention.query_key_value.bias', p1, p2)
+            checkpoint.merge_params_dim_0(self.attention.qkv_lin.weight, 'attention.query_key_value.weight', p1, p2)
+
+            # Layer norm before attention
+            checkpoint.merge_params_duplicate(self.pre_ln_attn.bias, 'input_layernorm.bias', p1, p2)
+            checkpoint.merge_params_duplicate(self.pre_ln_attn.weight, 'input_layernorm.weight', p1, p2)
+
+            # FFN second transform
+            checkpoint.merge_params_dim_0(self.ffn.dense_h_h4.bias, 'mlp.dense_h_to_4h.bias', p1, p2)
+            checkpoint.merge_params_dim_0(self.ffn.dense_h_h4.weight, 'mlp.dense_h_to_4h.weight', p1, p2)
+
+            # FFN first transform
+            checkpoint.merge_params_sum(self.ffn.dense_h4_h.bias, 'mlp.dense_4h_to_h.bias', p1, p2)
+            checkpoint.merge_params_dim_1(self.ffn.dense_h4_h.weight, 'mlp.dense_4h_to_h.weight', p1, p2)
+
+            # Layer norm before FFN
+            checkpoint.merge_params_duplicate(self.pre_ln_ffn.bias, 'post_attention_layernorm.bias', p1, p2)
+            checkpoint.merge_params_duplicate(self.pre_ln_ffn.weight, 'post_attention_layernorm.weight', p1, p2)
+
+
+class FinalNorm(NeoXModule):
+    """
+    ## Final normalization layer
+    """
+
+    def __init__(self, n_hidden: int = 6_144):
+        """
+        :param n_hidden: is the embedding size
+        """
+        super().__init__()
+
+        self.ln = nn.LayerNorm(n_hidden)
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: are the embeddings of shape `[batch_size, seq_len, n_hidden]`
+        """
+        return self.ln(x)
+
+    def load_state(self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]):
+        """
+        Code to load the checkpoint
+        """
+        with monit.section('Load final normalization layer'):
+            checkpoint.merge_params_duplicate(self.ln.bias, 'norm.bias', p1, p2)
+            checkpoint.merge_params_duplicate(self.ln.weight, 'norm.weight', p1, p2)
+
+
+class ReadoutLayer(NeoXModule):
+    """
+    Readout layer
+    """
+
+    def __init__(self, n_hidden: int = 6_144, n_vocab: int = 50_432):
+        """
+        :param n_hidden: is the embedding size
+        :param n_vocab: is the size of the vocabulary
+        """
+        super().__init__()
+
+        self.linear = nn.Linear(n_hidden, n_vocab, bias=False)
+
+    def forward(self, x: torch.Tensor):
+        """
+        :param x: are the embeddings of shape `[batch_size, seq_len, n_hidden]`
+        """
+        return self.linear(x)
+
+    def load_state(self, p1: Dict[str, torch.Tensor], p2: Dict[str, torch.Tensor]):
+        """
+        Code to load the checkpoint
+        """
+        with monit.section('Load final linear layer'):
+            checkpoint.merge_params_dim_0(self.linear.weight, 'final_linear.weight', p1, p2)
+
+
+class LayerGenerator:
+    pre_created_layers: Dict[Any, Optional[NeoXModule]]
+
+    def __init__(self, *, n_vocab: int = 50_432, n_hidden: int = 6_144,
+                 n_layers: int = 44, n_heads: int = 64,
+                 filter_layers: Optional[Set] = None,
+                 is_clone_layers: bool = True,
+                 dtype: torch.dtype = torch.float,
+                 device: torch.device = torch.device('cpu'),
+                 is_llm_int8: bool = False,
+                 llm_int8_threshold: float = 6.0,
+                 is_flash_attention: bool = False
+                 ):
+        """
+        ### Generator to create layers
+
+        The layers are generated in the same order as checkpoints.
+
+        It gives `None` when a layer is not available; we use the layer indices as NeoX and there are two
+        transformation layers we don't need in our implementation.
+
+        :param n_vocab: is the number of tokens in the vocabulary
+        :param n_hidden: is the number of features in the embeddings
+        :param n_layers: is the number of transformer layers
+        :param n_heads: is the number of attention heads
+        :param filter_layers: are the set of layers to be used. All layers will be used if None.
+            This is used to test smaller versions of the model with fewer layers
+        :param is_clone_layers: specifies whether to clone the transformer layers (a bit faster)
+        :param dtype: is the data type of the model
+        :param device: is the device of the model
+        :param is_llm_int8: specifies whether to use int8 quantization
+        :param llm_int8_threshold: is the threshold $\alpha$ used to separate outlier features
+        :param is_flash_attention: specifies whether to use
+            [FlashAttention](https://github.com/HazyResearch/flash-attention)
+        """
+        if filter_layers is None:
+            filter_layers = set(range(n_layers + 3))
+
+        self.n_vocab = n_vocab
+        self.n_hidden = n_hidden
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.filter_layers = filter_layers
+        self.is_clone_layers = is_clone_layers
+        self.dtype = dtype
+        self.device = device
+        self.is_llm_int8 = is_llm_int8
+        self.llm_int8_threshold = llm_int8_threshold
+        self.is_flash_attention = is_flash_attention
+
+        self.pre_created_layers = dict(
+            transformer_layer=None,
+        )
+
+    def _prepare_layer(self, layer: NeoXModule):
+        """
+        #### Prepares the layer for usage
+
+        We move the layer to the device and convert it to the correct data type
+
+        :param layer: is the layer to prepare
+        :return: the prepared layer
+        """
+        return layer.to(self.device, self.dtype)
+
+    @torch.no_grad()
+    def post_load_prepare(self, layer: NeoXModule, *,
+                          is_llm_int8: bool = None,
+                          device: torch.device = None,
+                          llm_int8_threshold: float = None,
+                          ):
+        """
+        <a id="post_load_prepare"></a>
+
+        ### Layer transformations after loading the checkpoint
+
+        This function implements layer transformations after loading the checkpoint.
+
+        Currently, it only applies the int8 quantization.
+
+        :param layer: is the layer to prepare
+        :param is_llm_int8: specifies whether to use int8 quantization
+        :param device: is the device of the model
+        :param llm_int8_threshold: is the threshold $\alpha$ used to separate outlier features
+        :return: the prepared layer
+        """
+
+        # Get default values if not specified
+        if is_llm_int8 is None:
+            is_llm_int8 = self.is_llm_int8
+        if device is None:
+            device = self.device
+        if llm_int8_threshold is None:
+            llm_int8_threshold = self.llm_int8_threshold
+
+        # Skip if not using int8 quantization
+        if not is_llm_int8:
+            return layer
+
+        # Only convert the linear layers in the transformer layers
+        if not isinstance(layer, TransformerLayer):
+            return layer
+
+        # Use `make_llm_int8_linear` defined in [utilities](./utils/llm_int8.html).
+        from labml_nn.neox.utils.llm_int8 import make_llm_int8_linear
+
+        # Convert the linear layers
+        with monit.section('Convert to int8'):
+            layer.attention.output = make_llm_int8_linear(layer.attention.output,
+                                                          device=device,
+                                                          threshold=llm_int8_threshold)
+            layer.attention.qkv_lin = make_llm_int8_linear(layer.attention.qkv_lin,
+                                                           device=device,
+                                                           threshold=llm_int8_threshold)
+            layer.ffn.dense_h_h4 = make_llm_int8_linear(layer.ffn.dense_h_h4,
+                                                        device=device,
+                                                        threshold=llm_int8_threshold)
+            layer.ffn.dense_h4_h = make_llm_int8_linear(layer.ffn.dense_h4_h,
+                                                        device=device,
+                                                        threshold=llm_int8_threshold)
+        #
+        return layer
+
+    def _create_and_cache_layer(self, name: str, creator: Callable[[], NeoXModule]):
+        """
+        #### Creates and caches a layer
+
+        Copying cached layers is faster than initializing new layers because it takes time to
+        initialize parameters.
+
+        :param name: is the name of the layer
+        :param creator: is the function to create the layer
+        :return: the created layer or a copy of the cached layer
+        """
+
+        if not self.is_clone_layers:
+            return self._prepare_layer(creator())
+
+        if self.pre_created_layers[name] is None:
+            self.pre_created_layers[name] = self._prepare_layer(creator())
+
+        layer = copy.deepcopy(self.pre_created_layers[name])
+        return layer
+
+    def _create_transformer_layer(self):
+        return self._create_and_cache_layer(
+            'transformer_layer',
+            lambda: TransformerLayer(self.n_hidden, self.n_heads, is_flash_attention=self.is_flash_attention)
+        )
+
+    def _create_embedding_layer(self):
+        return Embedding(self.n_vocab, self.n_hidden)
+
+    def _create_final_norm_layer(self):
+        return FinalNorm(self.n_hidden)
+
+    def _create_readout_layer(self):
+        return ReadoutLayer(self.n_hidden, self.n_vocab)
+
+    @torch.no_grad()
+    def get_layers(self) -> Generator[Tuple[NeoXModule, Tuple[str, str]], None, None]:
+        """
+        ### Generator to get layers
+        """
+        # Embedding layer
+        if 0 in self.filter_layers:
+            with monit.section('Embedding layer'):
+                layer = self._prepare_layer(self._create_embedding_layer())
+            yield layer, ('layer_00-model_00-model_states.pt', 'layer_00-model_01-model_states.pt')
+
+        # Transformer layers
+        for i in range(self.n_layers):
+            # Transformer layer
+            if i + 1 in self.filter_layers:
+                with monit.section(f'Transformer Layer {i}'):
+                    yield self._create_transformer_layer(), \
+                          (f'layer_{i + 2 :02d}-model_00-model_states.pt',
+                           f'layer_{i + 2 :02d}-model_01-model_states.pt')
+
+        # Final normalization layer
+        if self.n_layers + 1 in self.filter_layers:
+            with monit.section('Final norm layer'):
+                layer = self._prepare_layer(self._create_final_norm_layer())
+            yield layer, ('layer_47-model_00-model_states.pt', 'layer_47-model_01-model_states.pt')
+
+        # Readout layer
+        if self.n_layers + 2 in self.filter_layers:
+            with monit.section('Readout layer'):
+                layer = self._prepare_layer(self._create_readout_layer())
+            yield layer, ('layer_48-model_00-model_states.pt', 'layer_48-model_01-model_states.pt')
+
+        for k in self.pre_created_layers.keys():
+            self.pre_created_layers[k] = None
+
+    @property
+    def total_layers(self):
+        """
+        ### Returns the total number of layers
+        """
+        return self.n_layers + 3
+
+    @torch.no_grad()
+    def load(self) -> Generator[NeoXModule, None, None]:
+        """
+        ### Generator to load layers
+        """
+        with monit.section("Layers"):
+            for i, (layer, files) in enumerate(self.get_layers()):
+                if files is not None:
+                    layer.load_state(*checkpoint.load_checkpoint_files(files))
+
+                layer = self.post_load_prepare(layer)
+
+                monit.progress(min(0.99, (i + 1) / self.total_layers))
+                yield layer

labml_nn/neox/readme.md

@@ -0,0 +1,53 @@
+# GPT-NeoX
+
+This is a simple implementation of [Eleuther GPT-NeoX](https://arxiv.org/abs/2204.06745) for inference and fine-tuning.
+
+
+* [Model definition](https://nn.labml.ai/neox/model.html)
+* [Tokenizer](https://nn.labml.ai/neox/tokenizer.html)
+* [Checkpoint downloading and loading helpers](https://nn.labml.ai/neox/checkpoint.html)
+* [Utilities](https://nn.labml.ai/neox/utils/index.html)
+
+### [Samples](https://nn.labml.ai/neox/samples/__init__.py)
+
+* [Generating text](https://nn.labml.ai/neox/samples/generate.html)
+* [Fine tuning the biases with pipeline-parallel](https://nn.labml.ai/neox/samples/finetune.html)
+
+### [Evaluation](https://nn.labml.ai/neox/evaluation/__init__.py)
+
+* [Evaluating half precision model on a single GPU](https://nn.labml.ai/neox/evaluation/half_precision.html)
+
+### Evaluation Results
+
+| Task       | Metric          | NeoX Impl (2 GPU) | This repo (1 GPU) | LLM.Int8 |
+|------------|-----------------|-------------------|-------------------|----------|
+| anli_r1    | acc             | 0.3270            | 0.3360            | 0.3440   |
+|            | acc_stderr      | 0.0148            | 0.0149            | 0.0150   |
+| anli_r2    | acc             | 0.3410            | 0.3350            | 0.3540   |
+|            | acc_stderr      | 0.0150            | 0.0149            | 0.0151   |
+| anli_r3    | acc             | 0.3567            | 0.3525            | 0.3567   |
+|            | acc_stderr      | 0.0138            | 0.0149            | 0.0138   |
+| hellaswag  | acc             | 0.5351            | 0.5353            | 0.5348   |
+|            | acc_stderr      | 0.0050            | 0.0050            | 0.0050   |
+|            | acc_norm        | 0.7140            | 0.7145            | 0.7132   |
+|            | acc_norm_stderr | 0.0045            | 0.0045            | 0.0045   |
+| lambada    | acc             | 0.7211            | 0.7204            | 0.7155   |
+|            | acc_stderr      | 0.0062            | 0.0063            | 0.0063   |
+|            | ppl             | 3.6760            | 3.6375            | 3.7245   |
+|            | ppl_stderr      | 0.0760            | 0.0747            | 0.0768   |
+| piqa       | acc             | 0.7748            | 0.7758            | 0.7769   |
+|            | acc_stderr      | 0.0097            | 0.0097            | 0.0097   |
+|            | acc_norm        | 0.7786            | 0.7845            | 0.7829   |
+|            | acc_norm_stderr | 0.0097            | 0.0096            | 0.0096   |
+| winogrande | acc             | 0.6598            | 0.6582            | 0.6606   |
+|            | acc_stderr      | 0.0133            | 0.0133            | 0.0133   |
+| wsc        | acc             | 0.5096            | 0.5000            | 0.5288   |
+|            | acc_stderr      | 0.0493            | 0.0493            | 0.0492   |
+| mathqa     | acc             |                   |                   | 0.2720   |
+|            | acc_stderr      |                   |                   | 0.0081   |
+|            | acc_norm        |                   |                   | 0.2727   |
+|            | acc_norm_stderr |                   |                   | 0.0082   |
+
+**Official [Eleuther](https://www.eleuther.ai)
+GPT-NoeX is source code is available at [eleutherai/gpt-neox](https://github.com/eleutherai/gpt-neox).**
+

labml_nn/neox/samples/__init__.py

@@ -0,0 +1,12 @@
+"""
+---
+title: Samples
+summary: >
+    Samples for inference and fine-tuning
+---
+
+# Samples
+
+* [Generating text](generate.html)
+* [Fine tuning the biases with pipeline-parallel training](finetune.html)
+"""

labml_nn/neox/samples/finetune.py

@@ -0,0 +1,128 @@
+"""
+---
+title: Fine Tune GPT-NeoX
+summary: >
+    Fine tune GPT-NeoX biases with Fairscale pipeline parallel module
+---
+
+# Fine Tune GPT-NeoX
+
+This shows how to fine tune GPT-NeoX with pipeline parallelism.
+"""
+
+import fairscale
+import torch
+import torch.nn as nn
+import torch.utils.data
+import torch.utils.data
+import typing
+from torch.utils.data import DataLoader, RandomSampler
+
+from labml import experiment, monit, tracker, lab
+from labml.configs import option
+from labml.logger import inspect
+from labml_nn.neox.utils.text_dataset import get_training_data
+from labml_nn.neox.utils.finetune import FineTuneBiases
+from labml_nn.neox.model import LayerGenerator, NeoXModule
+from labml_nn.neox.utils import balance_layers_simple
+from labml_nn.neox.utils.trainer import PipelineParallelTrainerConf
+
+
+@option(PipelineParallelTrainerConf.layers, 'PipelineBiases')
+def neox_layers(c: PipelineParallelTrainerConf):
+    """
+    ### Load GPT-NeoX layers
+    """
+    return list(LayerGenerator(is_clone_layers=c.is_clone_layers,
+                               filter_layers=c.filter_layers,
+                               dtype=c.dtype,
+                               ).load())
+
+
+@option(PipelineParallelTrainerConf.fine_tuner, 'PipelineBiases')
+def fine_tune_biases(c: PipelineParallelTrainerConf):
+    """
+    ### Create fine tuner for biases
+    """
+
+    fine_tuner = FineTuneBiases(typing.cast(typing.List[NeoXModule], c.layers))
+    # Mark biases as trainable
+    fine_tuner.set_trainable_params()
+
+    #
+    return fine_tuner
+
+
+@option(PipelineParallelTrainerConf.model, 'PipelineBiases')
+def pipe_model(c: PipelineParallelTrainerConf):
+    """
+    ### Create pipeline parallel model
+    """
+
+    if c.is_checkpointing:
+        raise NotImplementedError()
+    else:
+        layers = c.layers
+
+    # Make sure the finetuner is initialized
+    _ = c.fine_tuner
+
+    # Create the Pipe module
+    with monit.section('Pipe'):
+        # Get the layer distribution across GPUs
+        balance = balance_layers_simple(len(layers), c.n_gpus)
+        inspect(balance=balance)
+        # Devices for each GPU
+        devices = [torch.device(f'cuda:{i}') for i in range(c.n_gpus)]
+        # Create Fairscale Pipe module
+        pipe_model = fairscale.nn.Pipe(nn.Sequential(*layers),
+                                       balance=balance,
+                                       devices=devices,
+                                       chunks=c.chunks)
+
+    #
+    return pipe_model
+
+
+@option(PipelineParallelTrainerConf.train_loader)
+def tiny_shakespeare(c: PipelineParallelTrainerConf):
+    """
+    #### Tiny Shakespeare dataset
+    """
+    dataset = get_training_data(c.max_seq_len)
+
+    return DataLoader(dataset,
+                      batch_size=c.batch_size,
+                      sampler=RandomSampler(dataset, replacement=True))
+
+
+def main():
+    # Create experiment
+    experiment.create(name='pipe_neox_biases',
+                      writers={'screen', 'web_api'})
+
+    # Initialize configs
+    conf = PipelineParallelTrainerConf()
+    experiment.configs(conf, {
+        'learning_rate': 3e-4,
+        'is_checkpointing': False,
+        'max_seq_len': 128,
+        'batch_size': 64,
+        'chunks': 8,
+    })
+
+    # Start the experiment
+    with experiment.start():
+        # Initialize the model. Do this before the loop for cleaner logs.
+        _ = conf.model
+
+        # Train
+        for epoch in monit.loop(conf.epochs):
+            conf.train_epoch()
+            tracker.new_line()
+            torch.save(conf.fine_tuner.state_dict(), str(lab.get_data_path() / 'fine_tune.pt'))
+
+
+#
+if __name__ == '__main__':
+    main()

labml_nn/neox/samples/generate.py

@@ -0,0 +1,102 @@
+"""
+---
+title: Generate Text with GPT-NeoX
+summary: >
+     Generate Text with GPT-NeoX
+---
+
+#  Generate Text with GPT-NeoX
+
+This shows how to generate text from GPT-NeoX with a single GPU.
+
+This needs a GPU with more than 45GB memory.
+"""
+
+# Imports
+from typing import List
+
+import torch
+from torch import nn
+
+from labml import monit
+from labml_nn.neox.model import LayerGenerator
+from labml_nn.neox.utils import get_tokens, print_tokens
+from labml_nn.neox.utils.cache import get_cache
+
+# List of layers to load. This is used for testing.
+# You can assign a subset of layers like `{0, 1}` so that it only loads
+# the first to transformer layers.
+LAYERS = None
+
+# Prompt to complete
+PROMPT = 'Einstein was born in the German Empire, but moved to Switzerland in 1895, forsaking his German'
+
+
+def infer(model: nn.Module, ids: List[int], device: torch.device):
+    """
+    ### Predict the next token
+
+    :param model: is the model
+    :param ids: are the input token ids
+    :param device: is the device of the model
+    """
+
+    with torch.no_grad():
+        # Get the tokens
+        x = torch.tensor(ids)[None, :].to(device)
+        # Eval model
+        x = model(x)
+
+    # Return predicted token
+    return x[0].max(dim=-1)[1].tolist()
+
+
+def generate():
+    """
+    ## Generate text
+    """
+
+    # Setup [cache](../utils/cache.html) to cache intermediate key/value pairs for faster generation
+    cache = get_cache()
+    cache.set('use_cache', True)
+
+    # Device
+    device = torch.device('cuda:0')
+
+    # Load layers
+    layers = list(LayerGenerator(is_clone_layers=True,
+                                 filter_layers=LAYERS,
+                                 dtype=torch.float16,
+                                 device=device,
+                                 ).load())
+
+    model = nn.Sequential(*layers)
+
+    # Get token ids
+    ids = get_tokens(PROMPT)
+
+    # Run the model
+    cache.set('state_ids', (None, 1))
+    with monit.section('Infer'):
+        next_token = infer(model, ids, device)[-1]
+
+    # Append the predicted token
+    ids += [next_token]
+
+    # Predict 100 tokens
+    for i in range(1, 100):
+        # Set the state to use cached activations
+        cache.set('state_ids', (i, i + 1))
+        # Get next token. Note that we only feed the last token to the model because
+        # we cache the key/value pairs of previous tokens.
+        with monit.section('Infer'):
+            next_token = infer(model, [next_token], device)[-1]
+        # Append the predicted token
+        ids += [next_token]
+        # Print
+        print_tokens(ids, [ids])
+
+
+#
+if __name__ == '__main__':
+    generate()

labml_nn/neox/samples/llm_int8.py

@@ -0,0 +1,91 @@
+"""
+---
+title: Generate Text with GPT-NeoX using LLM.int8() quantization
+summary: >
+     Generate Text with GPT-NeoX using LLM.int8() quantization
+---
+
+#  Generate Text with GPT-NeoX using LLM.int8() quantization
+
+This shows how to generate text from GPT-NeoX using [LLM.int8() quantization](../utils/llm_int8.html).
+
+This needs a GPU with 24GB memory.
+"""
+
+import torch
+from torch import nn
+
+from labml import monit
+from labml_nn.neox.model import LayerGenerator
+from labml_nn.neox.samples.generate import PROMPT, infer
+from labml_nn.neox.utils import get_tokens, print_tokens
+from labml_nn.neox.utils.cache import get_cache
+
+
+def generate():
+    """
+    ## Generate text
+    """
+
+    # Setup [cache](../utils/cache.html) to cache intermediate key/value pairs for faster generation
+    cache = get_cache()
+    cache.set('use_cache', True)
+
+    # Device
+    device = torch.device('cuda:0')
+
+    # Load layers in float16 into CPU. We convert the layers to int8 later, because doing that
+    # on the fly after loading layers to GPU causes CUDA memory fragmentation
+    # (about 3GB memory can get lost due to fragmentation).
+    layer_generator = LayerGenerator(is_clone_layers=True,
+                                     dtype=torch.float16,
+                                     device=torch.device('cpu'),
+                                     is_llm_int8=False,
+                                     )
+    layers = list(layer_generator.load())
+
+    # This reduces CUDA memory fragmentation
+    for layer in monit.iterate('Convert to int8', layers, is_children_silent=True):
+        layer_generator.post_load_prepare(layer,
+                                          device=device,
+                                          is_llm_int8=True,
+                                          llm_int8_threshold=6.0,
+                                          )
+        layer.to(device)
+
+    # Create `nn.Sequential` model
+    model = nn.Sequential(*layers)
+
+    # Clear cache and print memory summary for debugging
+    torch.cuda.empty_cache()
+    print(torch.cuda.memory_summary())
+
+    # Get token ids
+    ids = get_tokens(PROMPT)
+
+    # Run the model.
+    # We use the [`infer`](generate.html) function defined in [`generate.py`](generate.html)
+    cache.set('state_ids', (None, 1))
+    with monit.section('Infer'):
+        next_token = infer(model, ids, device)[-1]
+
+    # Append the predicted token
+    ids += [next_token]
+
+    # Predict 100 tokens
+    for i in range(1, 100):
+        # Set the state to use cached activations
+        cache.set('state_ids', (i, i + 1))
+        # Get next token. Note that we only feed the last token to the model because
+        # we cache the key/value pairs of previous tokens.
+        with monit.section('Infer'):
+            next_token = infer(model, [next_token], device)[-1]
+        # Append the predicted token
+        ids += [next_token]
+        # Print
+        print_tokens(ids, [ids])
+
+
+#
+if __name__ == '__main__':
+    generate()

labml_nn/neox/tokenizer.py

@@ -0,0 +1,28 @@
+"""
+---
+title: GPT-NeoX Tokenizer
+summary: >
+    Loads the GPT-NeoX tokenizer
+---
+
+# GPT-NeoX Tokenizer
+
+This initializes a Hugging Face tokenizer from the downloaded vocabulary.
+"""
+
+from tokenizers import Tokenizer
+
+from labml import lab, monit
+
+
+@monit.func('Load NeoX Tokenizer')
+def get_tokenizer() -> Tokenizer:
+    """
+    ### Load NeoX Tokenizer
+
+    :return: the tokenizer
+    """
+    vocab_file = lab.get_data_path() / 'neox' / 'slim_weights' / '20B_tokenizer.json'
+    tokenizer = Tokenizer.from_file(str(vocab_file))
+
+    return tokenizer

labml_nn/neox/utils/__init__.py

@@ -0,0 +1,134 @@
+"""
+---
+title: Utilities and Helpers
+summary: >
+    Utilities and helper functions
+---
+
+# Utilities and Helpers
+
+* [Cache for intermediate activations (for faster inference)](cache.html)
+* [Tools for finetuning](finetune.html)
+* [Trainer](trainer.html)
+* [Text dataset](text_dataset.html)
+"""
+import typing
+from typing import List, Optional
+
+import torch
+
+from labml import logger
+from labml.logger import Text
+from labml_nn.neox.tokenizer import get_tokenizer
+
+if typing.TYPE_CHECKING:
+    from tokenizers import Tokenizer
+
+# Tokenizer singleton
+_TOKENIZER: Optional['Tokenizer'] = None
+
+
+def get_tokens(text: str) -> List[int]:
+    """
+    ### Get token ids
+
+    :param text: is the text to tokenize
+    :return: the token ids
+    """
+    global _TOKENIZER
+    if _TOKENIZER is None:
+        _TOKENIZER = get_tokenizer()
+    return _TOKENIZER.encode_batch([text])[0].ids
+
+
+def print_token_outputs(ids: List[int], *xs: torch.Tensor):
+    """
+    ### Print tokens from model outputs
+
+    Pretty prints target tokens along side outputs from the model(s).
+
+    :param ids: are the target token ids
+    :param xs: are the model(s) outputs
+    """
+    ids = ids + [-1]
+    xs = [[-1] + x[0].max(dim=-1)[1].tolist() for x in xs]
+
+    print_tokens(ids, xs)
+
+
+def print_tokens(target: List[int], others: List[List[int]]):
+    """
+    ### Print tokens
+
+    Pretty prints tokens for comparison
+
+    :param target: are the target token ids
+    :param others: are the sampled outputs from the model(s)
+    """
+
+    # Load tokenizer
+    global _TOKENIZER
+    if _TOKENIZER is None:
+        _TOKENIZER = get_tokenizer()
+
+    # Convert the tokens to list of strings
+    text = []
+    for i in range(len(target)):
+        tokens = [_TOKENIZER.decode([target[i]]) if target[i] != -1 else '---']
+        for j in range(len(others)):
+            tokens.append(_TOKENIZER.decode([others[j][i]]) if others[j][i] != -1 else '---')
+
+        text.append(tokens)
+
+    # Stats
+    correct = [0 for _ in others]
+    total = 0
+
+    # Iterate through tokens
+    for i in range(len(target)):
+        parts = [(f'{i}: ', Text.meta)]
+        parts += [('"', Text.subtle), (text[i][0], Text.subtle), ('"', Text.subtle), '\t']
+
+        # Empty target
+        if target[i] == -1:
+            for j in range(len(others)):
+                parts += [('"', Text.subtle), (text[i][j + 1], Text.subtle), ('"', Text.subtle), '\t']
+
+            logger.log(parts)
+            continue
+
+        # Number of tokens
+        total += 1
+
+        # Other outputs
+        for j in range(len(others)):
+            correct[j] += 1 if others[j][i] == target[i] else 0
+
+            parts += [('"', Text.subtle),
+                      (text[i][j + 1], Text.success if others[j][i] == target[i] else Text.danger),
+                      ('"', Text.subtle), '\t']
+
+        logger.log(parts)
+
+    # Stats
+    parts = [(f'{total}', Text.highlight), '\t']
+    for j in range(len(others)):
+        parts += [(f'{correct[j]}', Text.value), '\t']
+    logger.log(parts)
+
+
+def balance_layers_simple(n_layers: int, n_chunks: int):
+    """
+    ### Balance layers
+
+    Split the `n_layers` into `n_chunks`. This is used for pipeline parallel training.
+
+    :param n_layers: is the number of layers
+    :param n_chunks: is the number of chunks
+    :return: returns a list with the number of layers for each chunk
+    """
+    balance = []
+    for i in range(n_chunks):
+        balance.append((n_layers - sum(balance)) // (n_chunks - i))
+
+    return list(reversed(balance))

labml_nn/neox/utils/cache.py

@@ -0,0 +1,118 @@
+"""
+---
+title: Cache for Intermediate Activations
+summary: >
+    Cache for intermediate activations for faster inference.
+---
+
+# Cache for Intermediate Activations
+
+During inference the model outputs token by token.
+We use this simple cache to store key's and value's attention layers,
+so that we don't have to recompute them for previous tokens.
+"""
+
+from typing import Any
+
+
+class Cache:
+    """
+    ## Cache
+
+    This maintains a key-value cache and queues push values and pop them in the same order.
+    The queues are useful since we have multiple attention layers.
+    """
+
+    def __init__(self):
+        self._cache = {}
+
+    def clear_all(self):
+        """
+        ### Clear cache
+        """
+        self._cache = {}
+
+    def push(self, name: str, value: Any):
+        """
+        ### Push a value to a queue
+
+        :param name: is the name of the queue
+        :param value: is the value to be pushed
+        """
+
+        # Create an empty queue if it's not present
+        if name not in self._cache:
+            self._cache[name] = []
+
+        # Push to the queue
+        self._cache[name].append(value)
+
+    def q_size(self, name):
+        """
+        ### Return the size of the queue
+
+        :param name: is the name of the queue
+        :return: size of the queue if exists else None
+        """
+
+        if name not in self._cache:
+            return None
+
+        if type(self._cache[name]) != list:
+            return None
+
+        return len(self._cache[name])
+
+    def pop(self, name: str):
+        """
+        ### Pop from a queue
+
+        :param name: is the name of the queue
+        :return: the value
+        """
+        return self._cache[name].pop(0)
+
+    def set(self, key: str, value: Any):
+        """
+        ### Cache a value
+
+        :param key: is the name of the value to be cached
+        :param value: is the value
+        """
+        self._cache[key] = value
+
+    def get(self, key: str, default: Any = None):
+        """
+        ### Retrieve a value from cache
+
+        :param key: is the name used when caching
+        :param default: is the default value if the cache is empty
+        :return: the cached value
+        """
+        return self._cache.get(key, default)
+
+    def clear(self, key: str):
+        """
+        ### Clear a cache value
+
+        :param key: is the name used when caching
+        """
+        del self._cache[key]
+
+
+# Singleton for cache
+_INSTANCE = None
+
+
+def get_cache() -> Cache:
+    """
+    ### Get the cache instance
+
+    :return: the cache instance
+    """
+    global _INSTANCE
+
+    if _INSTANCE is None:
+        _INSTANCE = Cache()
+
+    return _INSTANCE

labml_nn/neox/utils/finetune.py

@@ -0,0 +1,55 @@
+from typing import List, Dict
+
+import torch
+from torch import nn
+
+from labml_nn.neox.model import TransformerLayer, NeoXModule
+
+
+class FineTuner:
+    def __init__(self, layers: List[NeoXModule]):
+        self.layers = layers
+
+    def get_trainable_params(self) -> Dict[str, nn.Parameter]:
+        params = {}
+        for i, layer in enumerate(self.layers):
+            params.update(self.get_layer_trainable_params(layer, prefix=f'layer_{i :02d}'))
+
+        return params
+
+    def get_layer_trainable_params(self, layer: NeoXModule, prefix: str) -> Dict[str, nn.Parameter]:
+        raise NotImplementedError
+
+    def set_trainable_params(self):
+        for layer in self.layers:
+            # Set `requires_grad` to `False` for the entire layer.
+            layer.requires_grad_(False)
+            #
+            for p in self.get_trainable_params().values():
+                p.requires_grad_(True)
+
+    def state_dict(self):
+        return {n: p.data.cpu() for n, p in self.get_trainable_params().items()}
+
+    def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
+        params = self.get_trainable_params()
+        for n, p in params.items():
+            p.data[:] = state_dict[n].to(p.data.device)
+
+        for n in state_dict.keys():
+            assert n in params, n
+
+
+class FineTuneBiases(FineTuner):
+    def get_layer_trainable_params(self, layer: NeoXModule, prefix: str) -> Dict[str, nn.Parameter]:
+        params = {}
+
+        if isinstance(layer, TransformerLayer):
+            # No need to train the mlp bias because we are adding it with attention output
+            params[f'{prefix}.attention.output.bias'] = layer.attention.output.bias
+            params[f'{prefix}.attention.qkv_lin.bias'] = layer.attention.qkv_lin.bias
+            params[f'{prefix}.ffn.dense_h_h4.bias'] = layer.ffn.dense_h_h4.bias
+        else:
+            pass
+
+        return params

labml_nn/neox/utils/llm_int8.py

@@ -0,0 +1,75 @@
+"""
+---
+title: LLM.int8() on GPT-NeoX
+summary: >
+    Transform nn.Linear layers to 8-bit integer layers.
+---
+
+# LLM.int() on GPT-NeoX
+
+This implements a utility function to transform a `nn.Linear` layer to LLM.int8() linear layer.
+
+[LLM.int8() paper](https://arxiv.org/abs/eb2bcaee1d0011edaa66a71c10a887e7)
+ shows you can use int8 quantization while handling outliers to
+reduce memory footprint without performance degradation in large language models.
+They convert weights and inputs to scaled 8-bit integers and does matrix multiplication
+producing int32 results which is then converted back to float16 and rescaled.
+They show that in large langauge models, some features can give extreme values (outliers)
+that dominate the model's output.
+These features get clamped in 8-bit integer space which causes the model performance to degrade.
+As a solution they pick these outliers (greater than a specified threshold)
+and compute their multiplications separately in float16 space.
+Since the percentage of outliers is around 0.01% this doesn't increase memory usage,
+and prevents the model from degrading performance.
+
+The code to transform GPT-NoeX layers is defined in [model.py](../model.html#post_load_prepare).
+
+Here are example uses of GPT-NeoX with int8 quantization.
+
+* [Generate Text](../samples/llm_int8.html)
+* [Run Evaluation Tests](../evaluation/llm_int8.html)
+"""
+
+# Import [`bitsandbytes`](https://github.com/timdettmers/bitsandbytes) package
+try:
+    from bitsandbytes.nn import Linear8bitLt, Int8Params
+except ImportError:
+    raise ImportError('''Please install `bitsandbytes` with `pip install bitsandbytes -U`''')
+
+import torch
+from torch import nn
+
+
+def make_llm_int8_linear(linear_module: nn.Linear, device: torch.device, threshold: float = 6.0):
+    """
+    ## Transform a `nn.Linear` layer to LLM.int8() linear layer
+
+    :param linear_module: is the `nn.Linear` layer to transform
+    :param device: is the device of the model
+    :param threshold: is the threshold $\alpha$ to use for outlier detection
+    """
+
+    #
+    assert isinstance(linear_module, nn.Linear)
+
+    # Create an empty Linear8bitLt module
+    int8_lin = Linear8bitLt(
+        linear_module.in_features,
+        linear_module.out_features,
+        linear_module.bias is not None,
+        has_fp16_weights=False,
+        threshold=threshold,
+    )
+
+    # Quantize the weights
+    int8_lin._parameters['weight'] = Int8Params(linear_module.weight.data.cpu(),
+                                                requires_grad=False,
+                                                has_fp16_weights=False).to(device)
+
+    # Set the bias in float16 space
+    if linear_module.bias is not None:
+        int8_lin._parameters['bias'] = nn.Parameter(linear_module.bias.data,
+                                                    requires_grad=False)
+
+    #
+    return int8_lin

labml_nn/neox/utils/text_dataset.py

@@ -0,0 +1,132 @@
+"""
+---
+title: Text Dataset for GPT-NeoX
+summary: >
+    Loads text datasets to fine-tune GPT-NeoX
+---
+
+# Text Dataset for GPT-NeoX
+"""
+from pathlib import PurePath, Path
+from typing import Optional, List
+
+import torch
+import torch.utils.data
+from labml import lab
+from labml import monit
+from labml.logger import inspect
+from labml.utils.download import download_file
+
+from labml_nn.neox.tokenizer import get_tokenizer
+
+
+def load_text(path: PurePath, url: Optional[str] = None, *, filter_subset: Optional[int] = None):
+    """
+    ### Load text file
+
+    :param path: is the location of the text file
+    :param url: is the URL to download the file from
+    :param filter_subset: is the number of characters to filter.
+     Use this during testing when trying large datasets
+    :return: the text content
+    """
+
+    path = Path(path)
+
+    # Download if it doesn't exist
+    if not path.exists():
+        if not url:
+            raise FileNotFoundError(str(path))
+        else:
+            download_file(url, path)
+
+    with monit.section("Load data"):
+        # Load data
+        with open(str(path), 'r') as f:
+            text = f.read()
+        # Filter
+        if filter_subset:
+            text = text[:filter_subset]
+
+    #
+    return text
+
+
+class NeoXDataset(torch.utils.data.Dataset):
+    """
+    ## Dataset for fine-tuning GPT-NeoX
+
+    This is not optimized to very large datasets.
+    """
+
+    def __init__(self, tokens: List[int], seq_len: int):
+        """
+        :param tokens: is the list of token ids
+        :param seq_len: is the sequence length of a single training sample
+        """
+
+        self.seq_len = seq_len
+        # Number of samples
+        n_samples = len(tokens) // seq_len
+        self.n_samples = n_samples
+        # Truncate
+        tokens = tokens[:n_samples * seq_len + 1]
+        # Create a PyTorch tensor
+        self.tokens = torch.tensor(tokens)
+
+    def __len__(self):
+        return self.n_samples
+
+    def __getitem__(self, idx: int):
+        """
+        ### Get a sample
+
+        :param idx: is the index of the sample
+        :return: the input and the target
+        """
+        offset = idx * self.seq_len
+        return self.tokens[offset:offset + self.seq_len], self.tokens[offset + 1:offset + 1 + self.seq_len]
+
+
+DATASETS = {
+    'tiny_shakespeare': {
+        'file': 'tiny_shakespeare.txt',
+        'url': 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
+    }
+}
+
+
+def get_training_data(seq_len: int = 32, dataset_name: str = 'tiny_shakespeare', truncate: int = -1):
+    """
+    ### Load Dataset
+
+    :param seq_len: is the sequence length of a single training sample
+    :param dataset_name: is the name of the dataset
+    :return: the dataset
+    """
+
+    ds = DATASETS[dataset_name]
+    # Load the content
+    text = load_text(lab.get_data_path() / ds['file'], ds['url'])
+    # Tokenize
+    tokenizer = get_tokenizer()
+    tokens = tokenizer.encode_batch([text])[0]
+
+    if truncate > 0:
+        token_ids = tokens.ids[:truncate * seq_len]
+    else:
+        token_ids = tokens.ids
+
+    #
+    return NeoXDataset(token_ids, seq_len)
+
+
+def _test():
+    dataset = get_training_data()
+
+    inspect(tokens=len(dataset.tokens))
+
+
+#
+if __name__ == '__main__':
+    _test()

labml_nn/neox/utils/trainer.py

@@ -0,0 +1,182 @@
+from typing import Optional, Set, List
+
+import torch.nn as nn
+import torch.optim
+import torch.utils.data
+from torch.cuda import amp
+from torch.cuda.amp import GradScaler
+
+from labml import monit, tracker
+from labml.configs import BaseConfigs, option
+from labml_nn.neox.utils.finetune import FineTuner
+
+
+def get_trainable_params(model: nn.Module):
+    """
+    ### Get trainable parameters
+
+    :param model: is the model to train
+    :return: a list of parameters for training
+    """
+
+    # Get all parameters
+    params = list(model.parameters())
+    # Filter parameters that require gradients
+    trainable_params = [p for p in params if p.requires_grad]
+
+    #
+    return trainable_params
+
+
+class TrainerConf(BaseConfigs):
+    model: nn.Module
+    layers: List[nn.Module]
+    optimizer: torch.optim.Optimizer = 'Adam'
+    train_loader: torch.utils.data.DataLoader
+    valid_loader: Optional[torch.utils.data.DataLoader] = None,
+    device: torch.device = torch.device('cuda:0')
+    scaler: Optional[GradScaler] = 'Default'
+    is_amp: bool = True
+    dtype: torch.dtype = torch.float16
+
+    is_clone_layers: bool = True
+
+    loss_func: nn.Module = nn.CrossEntropyLoss()
+    checkpoints_per_epoch: int = 0
+    samples_per_epoch: int = 0
+
+    grad_norm: Optional[float] = 1.0
+    learning_rate: float = 3e-4
+    max_seq_len: int = 1024
+    batch_size: int = 64
+    epochs: int = 16
+
+    n_gpus: int = torch.cuda.device_count()
+
+    filter_layers: Optional[Set] = None
+
+    def get_loss(self, sample, dataset_split: str):
+        """
+        :param dataset_split: train/valid
+        :param sample: is the sample
+        :return: the loss, output and the target
+        """
+        data, target = sample
+
+        # Forward pass
+        with monit.section('Forward pass'):
+            output = self.model(data.to(self.device))
+        # Move targets to the same device as output
+        target = target.to(output.device)
+        # Calculate loss
+        loss = self.loss_func(output.view(target.numel(), -1), target.view(-1))
+
+        return loss, output, target
+
+    def train(self):
+        for epoch in monit.loop(self.epochs):
+            self.train_epoch()
+            tracker.new_line()
+
+    def sample(self, idx):
+        pass
+
+    def save_checkpoint(self, idx):
+        pass
+
+    def get_iterators(self):
+        # Iterate through the batches
+        iterators = [('train', self.train_loader)]
+        if self.valid_loader is not None:
+            iterators.append(('valid', self.valid_loader))
+
+        if self.samples_per_epoch > 0:
+            iterators.append((self.sample, [i for i in range(self.samples_per_epoch)]))
+
+        if self.checkpoints_per_epoch > 0:
+            iterators.append((self.save_checkpoint, [i for i in range(self.checkpoints_per_epoch)]))
+
+        return iterators
+
+    def train_epoch(self):
+        # Set model for train
+        self.model.train()
+
+        iterators = self.get_iterators()
+        for split_name, sample in monit.mix(1024, *iterators):
+            if split_name == 'train':
+                # Set gradients to zero
+                self.optimizer.zero_grad()
+                tracker.add_global_step()
+
+            with torch.set_grad_enabled(split_name == 'train'):
+                if self.is_amp:
+                    # Forward pass
+                    with amp.autocast():
+                        loss, output, target = self.get_loss(sample, split_name)
+                else:
+                    loss, output, target = self.get_loss(sample, split_name)
+
+                # Get predictions
+                pred = output.argmax(dim=-1)
+                # Calculate accuracy
+                accuracy = pred.eq(target).sum().item() / (target != -100).sum()
+
+                tracker.add({f'loss.{split_name}': loss, f'acc.{split_name}': accuracy * 100})
+
+            if split_name == 'train':
+                if self.scaler is not None:
+                    # Backward pass
+                    loss = self.scaler.scale(loss)
+                    # tracker.add({'loss.scaled': loss})
+
+                with monit.section('Backward pass'):
+                    loss.backward()
+
+                # Optimize
+                with monit.section('Optimize'):
+                    if self.scaler is None:
+                        self.optimizer.step()
+                    else:
+                        self.scaler.unscale_(self.optimizer)
+                        if self.grad_norm is not None:
+                            torch.nn.utils.clip_grad_norm_(get_trainable_params(self.model), self.grad_norm)
+                        self.scaler.step(self.optimizer)
+                        self.scaler.update()
+
+            tracker.save()
+
+
+@option(TrainerConf.optimizer, 'Adam')
+def adam_optimizer(c: TrainerConf):
+    if c.dtype == torch.float32:
+        return torch.optim.Adam(get_trainable_params(c.model), lr=c.learning_rate)
+    elif c.dtype == torch.float16:
+        from labml_nn.optimizers.adam_fp16 import AdamFP16
+        return AdamFP16(get_trainable_params(c.model), lr=c.learning_rate)
+    else:
+        raise NotImplementedError()
+
+
+@option(TrainerConf.optimizer, 'SGD')
+def sgd_optimizer(c: TrainerConf):
+    return torch.optim.SGD(get_trainable_params(c.model), lr=c.learning_rate)
+
+
+@option(TrainerConf.scaler, 'Default')
+def grad_scaler(c: TrainerConf):
+    if not c.is_amp:
+        return None
+
+    if c.dtype == torch.float16:
+        from labml_nn.optimizers.adam_fp16 import GradScalerFP16
+        return GradScalerFP16()
+    else:
+        return GradScaler()
+
+
+class PipelineParallelTrainerConf(TrainerConf):
+    is_checkpointing: bool = False
+    chunks: int
+
+    fine_tuner: FineTuner