model

app.py

@@ -0,0 +1,36 @@
+import gradio as gr
+from pathlib import Path
+import torch
+from tsai_gpt.generate_for_app import generate_for_app
+
+pythia_model = "checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth"
+def generate_text(prompt):
+    generated_text = generate_for_app(prompt, num_samples=1, max_new_tokens=200, temperature=0.9, checkpoint_dir=Path("checkpoints/meta-llama/Llama-2-7b-chat-hf/"))
+   
+
+    return generated_text
+
+
+#gr.Interface(fn=generate_nanogpt_text, inputs=gr.Button(value="Generate text!"), outputs='text').launch(share=True)
+
+
+with  gr.Blocks() as demo:
+    gr.Markdown(
+    """
+    # Example of text generation with our pythia 160M model based on the RedPajama sample data:
+
+
+   The model checkpoint is the 'checkpoints/meta-llama/Llama-2-7b-chat-hf' dir. The hyper params used are the exact same emitted by the training main.ipynb notebook. The loss is less than 3.5; we can see syntactically correct but semantically meaningless sentences.
+
+   Keep in mind the output is limited to 250 tokens so the inference runs within reasonable time (10s) on CPU. (Huggingface free tier)
+
+   GPU inference can output much much longer sequences.
+   Click on the "Generate text" button to see the generated text. 
+    """)
+
+    generate_button = gr.Button("Generate text!")
+    input=gr.Textbox(label="Enter your prompt here")
+    output = gr.Textbox(label="Text generated by Pythia 160M trained model")
+    generate_button.click(fn=generate_text, inputs=input, outputs=output, api_name='text generation sample')
+
+demo.launch()

checkpoints/meta-llama/Llama-2-7b-chat-hf/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 2,
+  "max_length": 4096,
+  "pad_token_id": 0,
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.32.0.dev0"
+}

checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_config.json

@@ -0,0 +1 @@
+{"name": "Llama-2-7b-chat-hf", "hf_config": {"org": "meta-llama", "name": "Llama-2-7b-chat-hf"}, "block_size": 2048, "vocab_size": 50254, "padding_multiple": 128, "padded_vocab_size": 50304, "n_layer": 12, "n_head": 12, "n_embd": 768, "rotary_percentage": 0.25, "parallel_residual": true, "bias": true, "lm_head_bias": false, "n_query_groups": 12, "shared_attention_norm": false, "_norm_class": "LayerNorm", "norm_eps": 1e-05, "_mlp_class": "GptNeoxMLP", "gelu_approximate": "none", "intermediate_size": 3072, "rope_condense_ratio": 1, "rope_base": 10000}

checkpoints/meta-llama/Llama-2-7b-chat-hf/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

checkpoints/meta-llama/Llama-2-7b-chat-hf/tokenizer_config.json

@@ -0,0 +1,36 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": null,
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

main.ipynb

@@ -0,0 +1,661 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "torch.cuda.is_available()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "import math\n",
+    "import sys\n",
+    "import time\n",
+    "from pathlib import Path\n",
+    "from typing import Optional, Tuple, Union\n",
+    "\n",
+    "import lightning as L\n",
+    "import torch\n",
+    "from lightning.fabric.loggers import CSVLogger\n",
+    "from lightning.fabric.strategies import FSDPStrategy\n",
+    "from torch.utils.data import DataLoader\n",
+    "\n",
+    "# # support running without installing as a package\n",
+    "# wd = Path(__file__).parent.parent.resolve()\n",
+    "# sys.path.append(str(wd))\n",
+    "\n",
+    "from tsai_gpt.model import GPT, Block, Config\n",
+    "from tsai_gpt.packed_dataset import CombinedDataset, PackedDataset\n",
+    "from tsai_gpt.speed_monitor import SpeedMonitorBase, estimate_flops, measure_flops\n",
+    "from tsai_gpt.speed_monitor import SpeedMonitorFabric as SpeedMonitor\n",
+    "from tsai_gpt.utils import chunked_cross_entropy, get_default_supported_precision, num_parameters, load_checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_name = \"pythia-160m\"\n",
+    "name = \"redpajama\"\n",
+    "out_dir = Path(\"out\") / name\n",
+    "save_interval = 1000\n",
+    "eval_interval = 1000\n",
+    "eval_iters = 100\n",
+    "log_interval = 100"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Hyperparameters\n",
+    "learning_rate = 6e-3\n",
+    "batch_size = 32\n",
+    "micro_batch_size = 8\n",
+    "gradient_accumulation_steps = batch_size // micro_batch_size\n",
+    "assert gradient_accumulation_steps > 0\n",
+    "#max_iters = 600000  # num_epochs * (epoch_size // micro_batch_size) // devices\n",
+    "max_iters = 15000\n",
+    "weight_decay = 1e-1\n",
+    "beta1 = 0.9\n",
+    "beta2 = 0.95\n",
+    "grad_clip = 1.0\n",
+    "decay_lr = True\n",
+    "warmup_iters = 2000\n",
+    "lr_decay_iters = max_iters\n",
+    "min_lr = 6e-6"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Data proportions from https://arxiv.org/pdf/2302.13971.pdf Table 1\n",
+    "data_config = [\n",
+    "    (\"arxiv\", 2.5),\n",
+    "    (\"book\", 4.5),\n",
+    "    (\"c4\", 15.0),\n",
+    "    (\"cc\", 67.0),\n",
+    "    (\"github\", 4.5),\n",
+    "    (\"stackexchange\", 2.0),\n",
+    "    (\"wikipedia\", 4.5),\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith(\"_\")}\n",
+    "logger = CSVLogger(\"out\", name, flush_logs_every_n_steps=log_interval)\n",
+    "\n",
+    "\n",
+    "def setup(\n",
+    "    devices: int = 4,\n",
+    "    train_data_dir: Path = Path(\"data/redpajama_sample\"),\n",
+    "    val_data_dir: Optional[Path] = None,\n",
+    "    precision: Optional[str] = None,\n",
+    "    resume: Union[bool, Path] = False,\n",
+    ") -> None:\n",
+    "    precision = precision or get_default_supported_precision(training=True)\n",
+    "\n",
+    "    if devices > 1:\n",
+    "        strategy = FSDPStrategy(\n",
+    "            auto_wrap_policy={Block},\n",
+    "            activation_checkpointing_policy={Block},\n",
+    "            state_dict_type=\"full\",\n",
+    "            limit_all_gathers=True,\n",
+    "            cpu_offload=False,\n",
+    "        )\n",
+    "    else:\n",
+    "        strategy = \"auto\"\n",
+    "\n",
+    "    fabric = L.Fabric(devices=devices, strategy=strategy, precision=precision, loggers=logger)\n",
+    "    fabric.print(hparams)\n",
+    "    fabric.launch(main, train_data_dir, val_data_dir, resume)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_copy = None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def main(fabric: L.Fabric, train_data_dir: Path, val_data_dir: Path, resume: Union[bool, Path]) -> None:\n",
+    "    global model_copy\n",
+    "    speed_monitor = SpeedMonitor(fabric, window_size=50, time_unit=\"seconds\")\n",
+    "\n",
+    "    if fabric.global_rank == 0:\n",
+    "        out_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "    config = Config.from_name(model_name)\n",
+    "\n",
+    "    train_dataloader, val_dataloader = create_dataloaders(\n",
+    "        batch_size=micro_batch_size,\n",
+    "        block_size=config.block_size,\n",
+    "        fabric=fabric,\n",
+    "        train_data_dir=train_data_dir,\n",
+    "        val_data_dir=val_data_dir,\n",
+    "        seed=(1337 + fabric.global_rank),\n",
+    "    )\n",
+    "    if val_dataloader is None:\n",
+    "        train_dataloader = fabric.setup_dataloaders(train_dataloader)\n",
+    "    else:\n",
+    "        train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)\n",
+    "\n",
+    "    fabric.seed_everything(1337)  # same seed for every process to init model (FSDP)\n",
+    "\n",
+    "    fabric.print(f\"Loading model with {config.__dict__}\")\n",
+    "    t0 = time.perf_counter()\n",
+    "    import torch\n",
+    "    import torch.nn as nn\n",
+    "    def _init_weights(module: nn.Module) -> None:\n",
+    "            \"\"\"Meant to be used with `gpt.apply(gpt._init_weights)`.\"\"\"\n",
+    "            if isinstance(module, nn.Linear):\n",
+    "                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
+    "                if module.bias is not None:\n",
+    "                    torch.nn.init.zeros_(module.bias)\n",
+    "            elif isinstance(module, nn.Embedding):\n",
+    "                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
+    "            \n",
+    "    with fabric.init_module(empty_init=True):\n",
+    "        model = GPT(config)\n",
+    "        model.apply(_init_weights)\n",
+    "    model.apply(_init_weights)\n",
+    "\n",
+    "    \n",
+    "    # checkpoint_path = Path(\"out/redpajama/iter-000999-ckpt.pth\")\n",
+    "\n",
+    "    # load_checkpoint(fabric, model, checkpoint_path)\n",
+    "        \n",
+    "    # print(model.transformer.h[0].mlp.fc.weight)\n",
+    "\n",
+    "    fabric.print(f\"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.\")\n",
+    "    fabric.print(f\"Total parameters {num_parameters(model):,}\")\n",
+    "\n",
+    "    model = fabric.setup(model)\n",
+    "    optimizer = torch.optim.AdamW(\n",
+    "        model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2), foreach=False\n",
+    "    )\n",
+    "\n",
+    "    # model_copy = model\n",
+    "\n",
+    "    optimizer = fabric.setup_optimizers(optimizer)\n",
+    "\n",
+    "    state = {\"model\": model, \"optimizer\": optimizer, \"hparams\": hparams, \"iter_num\": 0, \"step_count\": 0}\n",
+    "\n",
+    "    if resume is True:\n",
+    "        resume = max(out_dir.glob(\"*.pth\"), key=lambda p: int(p.name.split(\"-\")[1]))\n",
+    "    if resume:\n",
+    "        fabric.print(f\"Resuming training from {resume}\")\n",
+    "        fabric.load(resume, state)\n",
+    "\n",
+    "    train_time = time.perf_counter()\n",
+    "    train(fabric, state, train_dataloader, val_dataloader, speed_monitor)\n",
+    "    fabric.print(f\"Training time: {(time.perf_counter()-train_time):.2f}s\")\n",
+    "    if fabric.device.type == \"cuda\":\n",
+    "        fabric.print(f\"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB\")\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train(\n",
+    "    fabric: L.Fabric,\n",
+    "    state: dict,\n",
+    "    train_dataloader: DataLoader,\n",
+    "    val_dataloader: DataLoader,\n",
+    "    speed_monitor: SpeedMonitorBase,\n",
+    ") -> None:\n",
+    "    model = state[\"model\"]\n",
+    "    optimizer = state[\"optimizer\"]\n",
+    "\n",
+    "    if val_dataloader is not None:\n",
+    "        validate(fabric, model, val_dataloader)  # sanity check\n",
+    "\n",
+    "    with torch.device(\"meta\"):\n",
+    "        meta_model = GPT(model.config)\n",
+    "        # \"estimated\" is not as precise as \"measured\". Estimated is optimistic but widely used in the wild.\n",
+    "        # When comparing MFU or FLOP numbers with other projects that use estimated FLOPs,\n",
+    "        # consider passing `SpeedMonitor(flops_per_batch=estimated_flops)` instead\n",
+    "        estimated_flops = estimate_flops(meta_model) * micro_batch_size\n",
+    "        fabric.print(f\"Estimated TFLOPs: {estimated_flops * fabric.world_size / 1e12:.2f}\")\n",
+    "        x = torch.randint(0, 1, (micro_batch_size, model.max_seq_length))\n",
+    "        measured_flops = measure_flops(meta_model, x)\n",
+    "        fabric.print(f\"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}\")\n",
+    "        del meta_model, x\n",
+    "\n",
+    "    total_lengths = 0\n",
+    "    total_t0 = time.perf_counter()\n",
+    "\n",
+    "    for state[\"iter_num\"], train_data in enumerate(train_dataloader, state[\"iter_num\"]):\n",
+    "        if state[\"iter_num\"] >= max_iters:\n",
+    "            checkpoint_path = out_dir / f\"iter-{state['iter_num']:06d}-ckpt.pth\"\n",
+    "            fabric.print(f\"Saving checkpoint to {str(checkpoint_path)!r}\")\n",
+    "            fabric.save(checkpoint_path, state)\n",
+    "            break\n",
+    "\n",
+    "        # determine and set the learning rate for this iteration\n",
+    "        lr = get_lr(state[\"iter_num\"]) if decay_lr else learning_rate\n",
+    "        for param_group in optimizer.param_groups:\n",
+    "            param_group[\"lr\"] = lr\n",
+    "\n",
+    "        iter_t0 = time.perf_counter()\n",
+    "\n",
+    "        input_ids = train_data[:, 0 : model.max_seq_length].contiguous()\n",
+    "        targets = train_data[:, 1 : model.max_seq_length + 1].contiguous()\n",
+    "\n",
+    "        is_accumulating = (state[\"iter_num\"] + 1) % gradient_accumulation_steps != 0\n",
+    "        with fabric.no_backward_sync(model, enabled=is_accumulating):\n",
+    "            logits = model(input_ids)\n",
+    "            loss = chunked_cross_entropy(logits, targets, chunk_size=0)\n",
+    "            fabric.backward(loss / gradient_accumulation_steps)\n",
+    "        \n",
+    "        # return \n",
+    "\n",
+    "        if not is_accumulating:\n",
+    "            fabric.clip_gradients(model, optimizer, max_norm=grad_clip)\n",
+    "            optimizer.step()\n",
+    "            optimizer.zero_grad()\n",
+    "            state[\"step_count\"] += 1\n",
+    "\n",
+    "        t1 = time.perf_counter()\n",
+    "        total_lengths += input_ids.size(1)\n",
+    "        speed_monitor.on_train_batch_end(\n",
+    "            (state[\"iter_num\"] + 1) * micro_batch_size,\n",
+    "            t1 - total_t0,\n",
+    "            # this assumes that device FLOPs are the same and that all devices have the same batch size\n",
+    "            fabric.world_size,\n",
+    "            flops_per_batch=measured_flops,\n",
+    "            lengths=total_lengths,\n",
+    "        )\n",
+    "        if state[\"iter_num\"] % log_interval == 0:\n",
+    "            fabric.print(\n",
+    "                f\"iter {state['iter_num']} step {state['step_count']}: loss {loss.item():.4f}, LR: {lr:.6f}, iter time:\"\n",
+    "                f\" {(t1 - iter_t0) * 1000:.2f}ms{' (optimizer.step)' if not is_accumulating else ''}\"\n",
+    "            )\n",
+    "\n",
+    "        if val_dataloader is not None and not is_accumulating and state[\"step_count\"] % eval_interval == 0:\n",
+    "            t0 = time.perf_counter()\n",
+    "            val_loss = validate(fabric, model, val_dataloader)\n",
+    "            t1 = time.perf_counter() - t0\n",
+    "            speed_monitor.eval_end(t1)\n",
+    "            fabric.print(f\"step {state['iter_num']}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f}ms\")\n",
+    "            fabric.barrier()\n",
+    "        if not is_accumulating and state[\"step_count\"] % save_interval == 0:\n",
+    "            checkpoint_path = out_dir / f\"iter-{state['iter_num']:06d}-ckpt.pth\"\n",
+    "            fabric.print(f\"Saving checkpoint to {str(checkpoint_path)!r}\")\n",
+    "            fabric.save(checkpoint_path, state)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@torch.inference_mode()\n",
+    "def validate(fabric: L.Fabric, model: torch.nn.Module, val_dataloader: DataLoader) -> torch.Tensor:\n",
+    "    fabric.print(\"Validating ...\")\n",
+    "    model.eval()\n",
+    "\n",
+    "    losses = torch.zeros(eval_iters, device=fabric.device)\n",
+    "    for k, val_data in enumerate(val_dataloader):\n",
+    "        input_ids = val_data[:, 0 : model.max_seq_length].contiguous()\n",
+    "        targets = val_data[:, 1 : model.max_seq_length + 1].contiguous()\n",
+    "        logits = model(input_ids)\n",
+    "        losses[k] = chunked_cross_entropy(logits, targets, chunk_size=0)\n",
+    "    out = losses.mean()\n",
+    "\n",
+    "    model.train()\n",
+    "    return out"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_dataloader(\n",
+    "    batch_size: int, block_size: int, data_dir: Path, fabric: L.Fabric, shuffle: bool = True, seed: int = 12345\n",
+    ") -> DataLoader:\n",
+    "    datasets = []\n",
+    "    for prefix, _ in data_config:\n",
+    "        filenames = glob.glob(str(data_dir / f\"{prefix}*\"))\n",
+    "        dataset = PackedDataset(\n",
+    "            filenames,\n",
+    "            n_chunks=4,\n",
+    "            block_size=block_size,\n",
+    "            shuffle=shuffle,\n",
+    "            seed=seed,\n",
+    "            num_processes=fabric.world_size,\n",
+    "            process_rank=fabric.global_rank,\n",
+    "        )\n",
+    "        datasets.append(dataset)\n",
+    "\n",
+    "    if not datasets:\n",
+    "        raise RuntimeError(\n",
+    "            f\"No data found at {data_dir}. Make sure you ran prepare_redpajama.py to create the dataset.\"\n",
+    "        )\n",
+    "\n",
+    "    weights = [weight for _, weight in data_config]\n",
+    "    sum_weights = sum(weights)\n",
+    "    weights = [el / sum_weights for el in weights]\n",
+    "\n",
+    "    combined_dataset = CombinedDataset(datasets=datasets, seed=seed, weights=weights)\n",
+    "\n",
+    "    return DataLoader(combined_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def create_dataloaders(\n",
+    "    batch_size: int,\n",
+    "    block_size: int,\n",
+    "    fabric: L.Fabric,\n",
+    "    train_data_dir: Path = Path(\"data/redpajama_sample\"),\n",
+    "    val_data_dir: Optional[Path] = None,\n",
+    "    seed: int = 12345,\n",
+    ") -> Tuple[DataLoader, DataLoader]:\n",
+    "    # Increase by one because we need the next word as well\n",
+    "    effective_block_size = block_size + 1\n",
+    "    train_dataloader = create_dataloader(\n",
+    "        batch_size=batch_size,\n",
+    "        block_size=effective_block_size,\n",
+    "        fabric=fabric,\n",
+    "        data_dir=train_data_dir,\n",
+    "        shuffle=True,\n",
+    "        seed=seed,\n",
+    "    )\n",
+    "    val_dataloader = (\n",
+    "        create_dataloader(\n",
+    "            batch_size=batch_size,\n",
+    "            block_size=effective_block_size,\n",
+    "            fabric=fabric,\n",
+    "            data_dir=val_data_dir,\n",
+    "            shuffle=False,\n",
+    "            seed=seed,\n",
+    "        )\n",
+    "        if val_data_dir\n",
+    "        else None\n",
+    "    )\n",
+    "    return train_dataloader, val_dataloader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_lr(it: int) -> float:\n",
+    "    # 1) linear warmup for warmup_iters steps\n",
+    "    if it < warmup_iters:\n",
+    "        return learning_rate * it / warmup_iters\n",
+    "    # 2) if it > lr_decay_iters, return min learning rate\n",
+    "    if it > lr_decay_iters:\n",
+    "        return min_lr\n",
+    "    # 3) in between, use cosine decay down to min learning rate\n",
+    "    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)\n",
+    "    assert 0 <= decay_ratio <= 1\n",
+    "    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))  # coeff ranges 0..1\n",
+    "    return min_lr + coeff * (learning_rate - min_lr)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using bfloat16 Automatic Mixed Precision (AMP)\n",
+      "Seed set to 1337\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'model_name': 'pythia-160m', 'name': 'redpajama', 'save_interval': 1000, 'eval_interval': 1000, 'eval_iters': 100, 'log_interval': 100, 'learning_rate': 0.006, 'batch_size': 32, 'micro_batch_size': 8, 'gradient_accumulation_steps': 4, 'max_iters': 15000, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'decay_lr': True, 'warmup_iters': 2000, 'lr_decay_iters': 15000, 'min_lr': 6e-06}\n",
+      "Loading model with {'name': 'pythia-160m', 'hf_config': {'org': 'EleutherAI', 'name': 'pythia-160m-deduped'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.03 seconds.\n",
+      "Total parameters 162,322,944\n",
+      "Estimated TFLOPs: 22.14\n",
+      "Measured TFLOPs: 15.86\n",
+      "iter 0 step 0: loss 10.9790, LR: 0.000000, iter time: 993.87ms\n",
+      "iter 100 step 25: loss 7.3972, LR: 0.000300, iter time: 56.71ms\n",
+      "iter 200 step 50: loss 5.9952, LR: 0.000600, iter time: 50.39ms\n",
+      "iter 300 step 75: loss 5.8594, LR: 0.000900, iter time: 50.64ms\n",
+      "iter 400 step 100: loss 6.0047, LR: 0.001200, iter time: 56.22ms\n",
+      "iter 500 step 125: loss 5.9611, LR: 0.001500, iter time: 55.67ms\n",
+      "iter 600 step 150: loss 5.7425, LR: 0.001800, iter time: 56.04ms\n",
+      "iter 700 step 175: loss 5.6345, LR: 0.002100, iter time: 56.77ms\n",
+      "iter 800 step 200: loss 5.4736, LR: 0.002400, iter time: 56.59ms\n",
+      "iter 900 step 225: loss 5.3942, LR: 0.002700, iter time: 55.94ms\n",
+      "iter 1000 step 250: loss 5.3758, LR: 0.003000, iter time: 55.97ms\n",
+      "iter 1100 step 275: loss 5.4347, LR: 0.003300, iter time: 50.85ms\n",
+      "iter 1200 step 300: loss 5.6140, LR: 0.003600, iter time: 55.73ms\n",
+      "iter 1300 step 325: loss 5.1840, LR: 0.003900, iter time: 50.78ms\n",
+      "iter 1400 step 350: loss 5.5744, LR: 0.004200, iter time: 55.55ms\n",
+      "iter 1500 step 375: loss 4.8744, LR: 0.004500, iter time: 56.58ms\n",
+      "iter 1600 step 400: loss 5.2784, LR: 0.004800, iter time: 56.40ms\n",
+      "iter 1700 step 425: loss 4.6915, LR: 0.005100, iter time: 56.13ms\n",
+      "iter 1800 step 450: loss 5.1188, LR: 0.005400, iter time: 51.62ms\n",
+      "iter 1900 step 475: loss 5.1236, LR: 0.005700, iter time: 54.36ms\n",
+      "iter 2000 step 500: loss 4.8566, LR: 0.006000, iter time: 56.13ms\n",
+      "iter 2100 step 525: loss 4.7733, LR: 0.005999, iter time: 56.35ms\n",
+      "iter 2200 step 550: loss 4.8778, LR: 0.005997, iter time: 50.88ms\n",
+      "iter 2300 step 575: loss 4.4962, LR: 0.005992, iter time: 56.58ms\n",
+      "iter 2400 step 600: loss 4.7976, LR: 0.005986, iter time: 55.87ms\n",
+      "iter 2500 step 625: loss 4.3805, LR: 0.005978, iter time: 57.95ms\n",
+      "iter 2600 step 650: loss 4.5287, LR: 0.005969, iter time: 59.25ms\n",
+      "iter 2700 step 675: loss 4.3517, LR: 0.005957, iter time: 52.44ms\n",
+      "iter 2800 step 700: loss 4.6585, LR: 0.005944, iter time: 51.33ms\n",
+      "iter 2900 step 725: loss 4.4335, LR: 0.005929, iter time: 58.50ms\n",
+      "iter 3000 step 750: loss 4.4874, LR: 0.005913, iter time: 57.96ms\n",
+      "iter 3100 step 775: loss 4.3373, LR: 0.005895, iter time: 57.91ms\n",
+      "iter 3200 step 800: loss 4.1922, LR: 0.005875, iter time: 56.81ms\n",
+      "iter 3300 step 825: loss 4.5304, LR: 0.005853, iter time: 57.85ms\n",
+      "iter 3400 step 850: loss 4.1766, LR: 0.005830, iter time: 58.52ms\n",
+      "iter 3500 step 875: loss 4.2740, LR: 0.005805, iter time: 57.50ms\n",
+      "iter 3600 step 900: loss 4.2820, LR: 0.005779, iter time: 58.88ms\n",
+      "iter 3700 step 925: loss 4.2292, LR: 0.005751, iter time: 57.29ms\n",
+      "iter 3800 step 950: loss 4.4272, LR: 0.005721, iter time: 51.97ms\n",
+      "iter 3900 step 975: loss 4.2242, LR: 0.005690, iter time: 56.50ms\n",
+      "Saving checkpoint to 'out/redpajama/iter-003999-ckpt.pth'\n",
+      "iter 4000 step 1000: loss 4.2313, LR: 0.005657, iter time: 52.22ms\n",
+      "iter 4100 step 1025: loss 3.7856, LR: 0.005622, iter time: 58.16ms\n",
+      "iter 4200 step 1050: loss 3.6729, LR: 0.005586, iter time: 57.64ms\n",
+      "iter 4300 step 1075: loss 3.3938, LR: 0.005549, iter time: 57.79ms\n",
+      "iter 4400 step 1100: loss 4.0303, LR: 0.005510, iter time: 59.47ms\n",
+      "iter 4500 step 1125: loss 4.0657, LR: 0.005469, iter time: 57.41ms\n",
+      "iter 4600 step 1150: loss 4.1320, LR: 0.005428, iter time: 52.20ms\n",
+      "iter 4700 step 1175: loss 3.7960, LR: 0.005384, iter time: 58.86ms\n",
+      "iter 4800 step 1200: loss 4.2554, LR: 0.005340, iter time: 58.40ms\n",
+      "iter 4900 step 1225: loss 3.8772, LR: 0.005294, iter time: 56.95ms\n",
+      "iter 5000 step 1250: loss 3.6817, LR: 0.005246, iter time: 57.29ms\n",
+      "iter 5100 step 1275: loss 3.9250, LR: 0.005198, iter time: 57.37ms\n",
+      "iter 5200 step 1300: loss 4.3746, LR: 0.005148, iter time: 57.65ms\n",
+      "iter 5300 step 1325: loss 4.0608, LR: 0.005096, iter time: 56.79ms\n",
+      "iter 5400 step 1350: loss 4.0989, LR: 0.005044, iter time: 56.71ms\n",
+      "iter 5500 step 1375: loss 3.7373, LR: 0.004990, iter time: 56.18ms\n",
+      "iter 5600 step 1400: loss 3.7889, LR: 0.004936, iter time: 49.99ms\n",
+      "iter 5700 step 1425: loss 4.0925, LR: 0.004880, iter time: 54.77ms\n",
+      "iter 5800 step 1450: loss 3.8588, LR: 0.004823, iter time: 55.27ms\n",
+      "iter 5900 step 1475: loss 4.1029, LR: 0.004765, iter time: 50.15ms\n",
+      "iter 6000 step 1500: loss 3.7252, LR: 0.004705, iter time: 49.90ms\n",
+      "iter 6100 step 1525: loss 3.4831, LR: 0.004645, iter time: 55.53ms\n",
+      "iter 6200 step 1550: loss 3.9866, LR: 0.004584, iter time: 54.90ms\n",
+      "iter 6300 step 1575: loss 3.7148, LR: 0.004522, iter time: 55.46ms\n",
+      "iter 6400 step 1600: loss 3.5724, LR: 0.004459, iter time: 50.00ms\n",
+      "iter 6500 step 1625: loss 3.7567, LR: 0.004396, iter time: 55.14ms\n",
+      "iter 6600 step 1650: loss 3.6850, LR: 0.004331, iter time: 54.25ms\n",
+      "iter 6700 step 1675: loss 3.8028, LR: 0.004266, iter time: 54.40ms\n",
+      "iter 6800 step 1700: loss 3.8552, LR: 0.004200, iter time: 54.60ms\n",
+      "iter 6900 step 1725: loss 3.8593, LR: 0.004133, iter time: 54.42ms\n",
+      "iter 7000 step 1750: loss 3.9515, LR: 0.004066, iter time: 54.44ms\n",
+      "iter 7100 step 1775: loss 3.6975, LR: 0.003998, iter time: 49.42ms\n",
+      "iter 7200 step 1800: loss 3.8195, LR: 0.003929, iter time: 54.56ms\n",
+      "iter 7300 step 1825: loss 3.7486, LR: 0.003860, iter time: 49.36ms\n",
+      "iter 7400 step 1850: loss 3.8587, LR: 0.003790, iter time: 54.44ms\n",
+      "iter 7500 step 1875: loss 3.6892, LR: 0.003720, iter time: 51.14ms\n",
+      "iter 7600 step 1900: loss 3.1359, LR: 0.003650, iter time: 57.86ms\n",
+      "iter 7700 step 1925: loss 3.4292, LR: 0.003579, iter time: 56.81ms\n",
+      "iter 7800 step 1950: loss 3.4693, LR: 0.003508, iter time: 50.99ms\n",
+      "iter 7900 step 1975: loss 3.2231, LR: 0.003436, iter time: 56.80ms\n",
+      "Saving checkpoint to 'out/redpajama/iter-007999-ckpt.pth'\n",
+      "iter 8000 step 2000: loss 3.8785, LR: 0.003364, iter time: 50.98ms\n",
+      "iter 8100 step 2025: loss 3.3586, LR: 0.003292, iter time: 56.01ms\n",
+      "iter 8200 step 2050: loss 3.5540, LR: 0.003220, iter time: 69.23ms\n",
+      "iter 8300 step 2075: loss 3.5014, LR: 0.003148, iter time: 64.78ms\n",
+      "iter 8400 step 2100: loss 4.1081, LR: 0.003075, iter time: 64.00ms\n",
+      "iter 8500 step 2125: loss 3.4069, LR: 0.003003, iter time: 70.03ms\n",
+      "iter 8600 step 2150: loss 3.4040, LR: 0.002931, iter time: 50.92ms\n",
+      "iter 8700 step 2175: loss 3.6502, LR: 0.002858, iter time: 56.68ms\n",
+      "iter 8800 step 2200: loss 3.7889, LR: 0.002786, iter time: 51.22ms\n",
+      "iter 8900 step 2225: loss 3.5215, LR: 0.002714, iter time: 56.65ms\n",
+      "iter 9000 step 2250: loss 3.5122, LR: 0.002642, iter time: 56.33ms\n",
+      "iter 9100 step 2275: loss 3.2583, LR: 0.002570, iter time: 55.80ms\n",
+      "iter 9200 step 2300: loss 3.6411, LR: 0.002498, iter time: 50.97ms\n",
+      "iter 9300 step 2325: loss 3.3789, LR: 0.002427, iter time: 56.73ms\n",
+      "iter 9400 step 2350: loss 3.4600, LR: 0.002356, iter time: 56.51ms\n",
+      "iter 9500 step 2375: loss 3.5573, LR: 0.002286, iter time: 56.18ms\n",
+      "iter 9600 step 2400: loss 3.8386, LR: 0.002216, iter time: 56.41ms\n",
+      "iter 9700 step 2425: loss 3.6447, LR: 0.002146, iter time: 50.92ms\n",
+      "iter 9800 step 2450: loss 3.7056, LR: 0.002077, iter time: 56.08ms\n",
+      "iter 9900 step 2475: loss 3.6347, LR: 0.002008, iter time: 56.52ms\n",
+      "iter 10000 step 2500: loss 3.6540, LR: 0.001940, iter time: 56.33ms\n",
+      "iter 10100 step 2525: loss 2.9236, LR: 0.001873, iter time: 56.44ms\n",
+      "iter 10200 step 2550: loss 3.3438, LR: 0.001806, iter time: 56.55ms\n",
+      "iter 10300 step 2575: loss 2.8825, LR: 0.001740, iter time: 56.82ms\n",
+      "iter 10400 step 2600: loss 3.2412, LR: 0.001675, iter time: 56.60ms\n",
+      "iter 10500 step 2625: loss 3.7394, LR: 0.001610, iter time: 51.23ms\n",
+      "iter 10600 step 2650: loss 3.0055, LR: 0.001547, iter time: 56.51ms\n",
+      "iter 10700 step 2675: loss 3.0301, LR: 0.001484, iter time: 55.63ms\n",
+      "iter 10800 step 2700: loss 3.7498, LR: 0.001422, iter time: 56.42ms\n",
+      "iter 10900 step 2725: loss 3.3228, LR: 0.001361, iter time: 56.71ms\n",
+      "iter 11000 step 2750: loss 3.7291, LR: 0.001301, iter time: 56.78ms\n",
+      "iter 11100 step 2775: loss 3.2531, LR: 0.001241, iter time: 50.94ms\n",
+      "iter 11200 step 2800: loss 3.5005, LR: 0.001183, iter time: 51.03ms\n",
+      "iter 11300 step 2825: loss 3.6694, LR: 0.001126, iter time: 55.72ms\n",
+      "iter 11400 step 2850: loss 3.6801, LR: 0.001070, iter time: 56.28ms\n",
+      "iter 11500 step 2875: loss 3.5435, LR: 0.001016, iter time: 56.29ms\n",
+      "iter 11600 step 2900: loss 3.1989, LR: 0.000962, iter time: 56.46ms\n",
+      "iter 11700 step 2925: loss 3.7118, LR: 0.000910, iter time: 56.78ms\n",
+      "iter 11800 step 2950: loss 3.4840, LR: 0.000858, iter time: 56.68ms\n",
+      "iter 11900 step 2975: loss 3.3171, LR: 0.000808, iter time: 58.20ms\n",
+      "Saving checkpoint to 'out/redpajama/iter-011999-ckpt.pth'\n",
+      "iter 12000 step 3000: loss 3.4315, LR: 0.000760, iter time: 133.55ms\n",
+      "iter 12100 step 3025: loss 3.2361, LR: 0.000712, iter time: 50.96ms\n",
+      "iter 12200 step 3050: loss 3.2333, LR: 0.000666, iter time: 56.34ms\n"
+     ]
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m torch\u001b[39m.\u001b[39mset_float32_matmul_precision(\u001b[39m\"\u001b[39m\u001b[39mmedium\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m----> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m setup(\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=2'>3</a>\u001b[0m     devices\u001b[39m=\u001b[39;49m\u001b[39m1\u001b[39;49m,\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=3'>4</a>\u001b[0m     \u001b[39m#train_data_dir=Path(\"data/lit-redpajama-sample\")\u001b[39;49;00m\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=4'>5</a>\u001b[0m     train_data_dir\u001b[39m=\u001b[39;49mPath(\u001b[39m\"\u001b[39;49m\u001b[39m/home/raghu/work/data/redpajama/data/lit-redpajama-sample\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m      <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m )\n",
+      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m2\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a>\u001b[0m fabric \u001b[39m=\u001b[39m L\u001b[39m.\u001b[39mFabric(devices\u001b[39m=\u001b[39mdevices, strategy\u001b[39m=\u001b[39mstrategy, precision\u001b[39m=\u001b[39mprecision, loggers\u001b[39m=\u001b[39mlogger)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m fabric\u001b[39m.\u001b[39mprint(hparams)\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=26'>27</a>\u001b[0m fabric\u001b[39m.\u001b[39;49mlaunch(main, train_data_dir, val_data_dir, resume)\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:834\u001b[0m, in \u001b[0;36mFabric.launch\u001b[0;34m(self, function, *args, **kwargs)\u001b[0m\n\u001b[1;32m    829\u001b[0m \u001b[39melif\u001b[39;00m \u001b[39misinstance\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstrategy\u001b[39m.\u001b[39mlauncher, (_MultiProcessingLauncher, _XLALauncher)):\n\u001b[1;32m    830\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mTypeError\u001b[39;00m(\n\u001b[1;32m    831\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTo use the `\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstrategy)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m` strategy, `.launch()` needs to be called with a function\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    832\u001b[0m         \u001b[39m\"\u001b[39m\u001b[39m that contains the code to launch in processes.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m    833\u001b[0m     )\n\u001b[0;32m--> 834\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_wrap_and_launch(function, \u001b[39mself\u001b[39;49m, \u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:920\u001b[0m, in \u001b[0;36mFabric._wrap_and_launch\u001b[0;34m(self, to_run, *args, **kwargs)\u001b[0m\n\u001b[1;32m    918\u001b[0m \u001b[39mif\u001b[39;00m (launcher \u001b[39m:=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_strategy\u001b[39m.\u001b[39mlauncher) \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    919\u001b[0m     \u001b[39mreturn\u001b[39;00m launcher\u001b[39m.\u001b[39mlaunch(to_run, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)\n\u001b[0;32m--> 920\u001b[0m \u001b[39mreturn\u001b[39;00m to_run(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:925\u001b[0m, in \u001b[0;36mFabric._wrap_with_setup\u001b[0;34m(self, to_run, *args, **kwargs)\u001b[0m\n\u001b[1;32m    923\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_strategy\u001b[39m.\u001b[39msetup_environment()\n\u001b[1;32m    924\u001b[0m \u001b[39mwith\u001b[39;00m _replace_dunder_methods(DataLoader, \u001b[39m\"\u001b[39m\u001b[39mdataset\u001b[39m\u001b[39m\"\u001b[39m), _replace_dunder_methods(BatchSampler):\n\u001b[0;32m--> 925\u001b[0m     \u001b[39mreturn\u001b[39;00m to_run(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m7\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=67'>68</a>\u001b[0m     fabric\u001b[39m.\u001b[39mload(resume, state)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=69'>70</a>\u001b[0m train_time \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mperf_counter()\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=70'>71</a>\u001b[0m train(fabric, state, train_dataloader, val_dataloader, speed_monitor)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=71'>72</a>\u001b[0m fabric\u001b[39m.\u001b[39mprint(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTraining time: \u001b[39m\u001b[39m{\u001b[39;00m(time\u001b[39m.\u001b[39mperf_counter()\u001b[39m-\u001b[39mtrain_time)\u001b[39m:\u001b[39;00m\u001b[39m.2f\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39ms\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=72'>73</a>\u001b[0m \u001b[39mif\u001b[39;00m fabric\u001b[39m.\u001b[39mdevice\u001b[39m.\u001b[39mtype \u001b[39m==\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mcuda\u001b[39m\u001b[39m\"\u001b[39m:\n",
+      "\u001b[1;32m/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb Cell 14\u001b[0m line \u001b[0;36m5\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=51'>52</a>\u001b[0m \u001b[39m# return \u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=53'>54</a>\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m is_accumulating:\n\u001b[0;32m---> <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=54'>55</a>\u001b[0m     fabric\u001b[39m.\u001b[39;49mclip_gradients(model, optimizer, max_norm\u001b[39m=\u001b[39;49mgrad_clip)\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=55'>56</a>\u001b[0m     optimizer\u001b[39m.\u001b[39mstep()\n\u001b[1;32m     <a href='vscode-notebook-cell://wsl%2Bubuntu/home/raghu/work/ERA-V1-assignments/assignment-22/main.ipynb#X16sdnNjb2RlLXJlbW90ZQ%3D%3D?line=56'>57</a>\u001b[0m     optimizer\u001b[39m.\u001b[39mzero_grad()\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/fabric.py:455\u001b[0m, in \u001b[0;36mFabric.clip_gradients\u001b[0;34m(self, module, optimizer, clip_val, max_norm, norm_type, error_if_nonfinite)\u001b[0m\n\u001b[1;32m    453\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mNone\u001b[39;00m\n\u001b[1;32m    454\u001b[0m \u001b[39mif\u001b[39;00m max_norm \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[0;32m--> 455\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mstrategy\u001b[39m.\u001b[39;49mclip_gradients_norm(\n\u001b[1;32m    456\u001b[0m         _unwrap_objects(module),\n\u001b[1;32m    457\u001b[0m         _unwrap_objects(optimizer),\n\u001b[1;32m    458\u001b[0m         max_norm\u001b[39m=\u001b[39;49mmax_norm,\n\u001b[1;32m    459\u001b[0m         norm_type\u001b[39m=\u001b[39;49mnorm_type,\n\u001b[1;32m    460\u001b[0m         error_if_nonfinite\u001b[39m=\u001b[39;49merror_if_nonfinite,\n\u001b[1;32m    461\u001b[0m     )\n\u001b[1;32m    462\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mValueError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39mYou have to specify either `clip_val` or `max_norm` to do gradient clipping!\u001b[39m\u001b[39m\"\u001b[39m)\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/lightning/fabric/strategies/strategy.py:380\u001b[0m, in \u001b[0;36mStrategy.clip_gradients_norm\u001b[0;34m(self, module, optimizer, max_norm, norm_type, error_if_nonfinite)\u001b[0m\n\u001b[1;32m    378\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprecision\u001b[39m.\u001b[39munscale_gradients(optimizer)\n\u001b[1;32m    379\u001b[0m parameters \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mprecision\u001b[39m.\u001b[39mmain_params(optimizer)\n\u001b[0;32m--> 380\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39;49mnn\u001b[39m.\u001b[39;49mutils\u001b[39m.\u001b[39;49mclip_grad_norm_(\n\u001b[1;32m    381\u001b[0m     parameters, max_norm\u001b[39m=\u001b[39;49mmax_norm, norm_type\u001b[39m=\u001b[39;49mnorm_type, error_if_nonfinite\u001b[39m=\u001b[39;49merror_if_nonfinite\n\u001b[1;32m    382\u001b[0m )\n",
+      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/nn/utils/clip_grad.py:63\u001b[0m, in \u001b[0;36mclip_grad_norm_\u001b[0;34m(parameters, max_norm, norm_type, error_if_nonfinite, foreach)\u001b[0m\n\u001b[1;32m     59\u001b[0m             norms\u001b[39m.\u001b[39mextend([torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39mvector_norm(g, norm_type) \u001b[39mfor\u001b[39;00m g \u001b[39min\u001b[39;00m grads])\n\u001b[1;32m     61\u001b[0m     total_norm \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39mvector_norm(torch\u001b[39m.\u001b[39mstack([norm\u001b[39m.\u001b[39mto(first_device) \u001b[39mfor\u001b[39;00m norm \u001b[39min\u001b[39;00m norms]), norm_type)\n\u001b[0;32m---> 63\u001b[0m \u001b[39mif\u001b[39;00m error_if_nonfinite \u001b[39mand\u001b[39;00m torch\u001b[39m.\u001b[39mlogical_or(total_norm\u001b[39m.\u001b[39misnan(), total_norm\u001b[39m.\u001b[39misinf()):\n\u001b[1;32m     64\u001b[0m     \u001b[39mraise\u001b[39;00m \u001b[39mRuntimeError\u001b[39;00m(\n\u001b[1;32m     65\u001b[0m         \u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mThe total norm of order \u001b[39m\u001b[39m{\u001b[39;00mnorm_type\u001b[39m}\u001b[39;00m\u001b[39m for gradients from \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     66\u001b[0m         \u001b[39m'\u001b[39m\u001b[39m`parameters` is non-finite, so it cannot be clipped. To disable \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     67\u001b[0m         \u001b[39m'\u001b[39m\u001b[39mthis error and scale the gradients by the non-finite norm anyway, \u001b[39m\u001b[39m'\u001b[39m\n\u001b[1;32m     68\u001b[0m         \u001b[39m'\u001b[39m\u001b[39mset `error_if_nonfinite=False`\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     69\u001b[0m clip_coef \u001b[39m=\u001b[39m max_norm \u001b[39m/\u001b[39m (total_norm \u001b[39m+\u001b[39m \u001b[39m1e-6\u001b[39m)\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "torch.set_float32_matmul_precision(\"medium\")\n",
+    "setup(\n",
+    "    devices=1,\n",
+    "    #train_data_dir=Path(\"data/lit-redpajama-sample\")\n",
+    "    train_data_dir=Path(\"/home/raghu/work/data/redpajama/data/lit-redpajama-sample\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

requirements.txt

@@ -0,0 +1,5 @@
+torch
+gradio
+lightning
+sentencepiece
+

tsai_gpt/__init__.py

@@ -0,0 +1,15 @@
+from tsai_gpt.model import GPT
+from tsai_gpt.config import Config
+from tsai_gpt.tokenizer import Tokenizer
+
+from lightning_utilities.core.imports import RequirementCache
+
+_LIGHTNING_AVAILABLE = RequirementCache("lightning>=2.1.0.dev0")
+if not bool(_LIGHTNING_AVAILABLE):
+    raise ImportError(
+        "Lit-GPT requires lightning==2.1. Please run:\n"
+        f" pip uninstall -y lightning; pip install -r requirements.txt\n{str(_LIGHTNING_AVAILABLE)}"
+    )
+
+
+__all__ = ["GPT", "Config", "Tokenizer"]

tsai_gpt/config.py

@@ -0,0 +1,1181 @@
+import json
+from copy import deepcopy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal, Optional, Type, Union
+
+import torch
+from typing_extensions import Self
+
+import tsai_gpt.model
+from tsai_gpt.utils import find_multiple
+
+
+@dataclass
+class Config:
+    name: str = ""
+    hf_config: dict = field(default_factory=dict)
+    block_size: int = 4096
+    vocab_size: int = 50254
+    padding_multiple: int = 512
+    padded_vocab_size: Optional[int] = None
+    n_layer: int = 16
+    n_head: int = 32
+    n_embd: int = 4096
+    rotary_percentage: float = 0.25
+    parallel_residual: bool = True
+    bias: bool = True
+    lm_head_bias: bool = False
+    # to use multi-head attention (MHA), set this to `n_head` (default)
+    # to use multi-query attention (MQA), set this to 1
+    # to use grouped-query attention (GQA), set this to a value in between
+    # Example with `n_head=4`
+    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+    # │ v ││ v ││ v ││ v │     │ v │    │ v │             │ v │
+    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+    #   │    │    │    │         │        │                 │
+    # ┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
+    # │ k ││ k ││ k ││ k │     │ k │    │ k │             │ k │
+    # └───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
+    #   │    │    │    │      ┌──┴──┐  ┌──┴──┐      ┌────┬──┴─┬────┐
+    # ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐
+    # │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │
+    # └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘
+    # ◀──────────────────▶  ◀──────────────────▶  ◀──────────────────▶
+    #         MHA                    GQA                   MQA
+    #   n_query_groups=4       n_query_groups=2      n_query_groups=1
+    #
+    # credit https://arxiv.org/pdf/2305.13245.pdf
+    n_query_groups: Optional[int] = None
+    shared_attention_norm: bool = False
+    _norm_class: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
+    norm_eps: float = 1e-5
+    _mlp_class: Literal["GptNeoxMLP", "LLaMAMLP"] = "GptNeoxMLP"
+    gelu_approximate: str = "none"
+    intermediate_size: Optional[int] = None
+    rope_condense_ratio: int = 1
+    rope_base: int = 10000
+
+    def __post_init__(self):
+        if not self.name:
+            self.name = self.hf_config.get("name", self.name)
+
+        assert self.n_embd % self.n_head == 0
+        self.head_size = self.n_embd // self.n_head
+
+        # vocab size should be a power of 2 to be optimal on hardware. compute the closest value
+        if self.padded_vocab_size is None:
+            self.padded_vocab_size = find_multiple(self.vocab_size, self.padding_multiple)
+        else:
+            # vocab size shouldn't be larger than padded vocab size
+            self.vocab_size = min(self.vocab_size, self.padded_vocab_size)
+
+        # compute the number of query groups
+        if self.n_query_groups is not None:
+            assert self.n_head % self.n_query_groups == 0
+        else:
+            self.n_query_groups = self.n_head
+
+        # compute the intermediate size for MLP if not set
+        if self.intermediate_size is None:
+            if self._mlp_class == "LLaMAMLP":
+                raise ValueError("The config needs to set the `intermediate_size`")
+            self.intermediate_size = 4 * self.n_embd
+
+        self.rope_n_elem = int(self.rotary_percentage * self.head_size)
+
+    @classmethod
+    def from_name(cls, name: str, **kwargs: Any) -> Self:
+        if name not in name_to_config:
+            # search through all `config['hf_config']['name']`
+            conf_dict = next(config for config in configs if name == config["hf_config"]["name"])
+        else:
+            conf_dict = name_to_config[name]
+
+        conf_dict = conf_dict.copy()
+        if "condense_ratio" in kwargs:  # legacy name
+            kwargs["rope_condense_ratio"] = kwargs.pop("condense_ratio")
+        conf_dict.update(kwargs)
+        return cls(**conf_dict)
+
+    @classmethod
+    def from_json(cls, path: Union[str, Path], **kwargs: Any) -> Self:
+        with open(path, encoding="utf-8") as fp:
+            json_kwargs = json.load(fp)
+        if "condense_ratio" in json_kwargs:  # legacy name
+            json_kwargs["rope_condense_ratio"] = json_kwargs.pop("condense_ratio")
+        if "condense_ratio" in kwargs:  # legacy name
+            kwargs["rope_condense_ratio"] = kwargs.pop("condense_ratio")
+        if "org" in json_kwargs:  # legacy name
+            json_kwargs["hf_config"] = {"name": json_kwargs["name"], "org": json_kwargs.pop("org")}
+        if "org" in kwargs:  # legacy name
+            kwargs["hf_config"] = {"name": kwargs.get("name", json_kwargs["name"]), "org": kwargs.pop("org")}
+        json_kwargs.update(kwargs)
+        return cls(**json_kwargs)
+
+    @property
+    def mlp_class(self) -> Type:
+        # `self._mlp_class` cannot be the type to keep the config json serializable
+        return getattr(tsai_gpt.model, self._mlp_class)
+
+    @property
+    def norm_class(self) -> Type:
+        # `self._norm_class` cannot be the type to keep the config json serializable
+        if self._norm_class == "RMSNorm":
+            from tsai_gpt.rmsnorm import RMSNorm
+
+            return RMSNorm
+        return getattr(torch.nn, self._norm_class)
+
+
+########################
+# Stability AI StableLM
+########################
+configs = [
+    # https://huggingface.co/stabilityai/stablelm-base-alpha-3b/blob/main/config.json
+    dict(name="stablelm-base-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-base-alpha-3b")),
+    # https://huggingface.co/stabilityai/stablelm-base-alpha-7b/blob/main/config.json
+    dict(
+        name="stablelm-base-alpha-7b",
+        hf_config=dict(org="stabilityai", name="stablelm-base-alpha-7b"),
+        n_head=48,
+        n_embd=6144,
+        padding_multiple=256,
+    ),
+    # https://huggingface.co/stabilityai/stablelm-tuned-alpha-3b/blob/main/config.json
+    dict(name="stablelm-tuned-alpha-3b", hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-3b"), n_head=32),
+    # https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b/blob/main/config.json
+    dict(
+        name="stablelm-tuned-alpha-7b",
+        hf_config=dict(org="stabilityai", name="stablelm-tuned-alpha-7b"),
+        n_head=48,
+        n_embd=6144,
+        padding_multiple=256,
+    ),
+]
+
+####################
+# EleutherAI Pythia
+####################
+pythia = [
+    # https://huggingface.co/EleutherAI/pythia-70m/blob/main/config.json
+    dict(
+        name="pythia-70m",
+        hf_config=dict(org="EleutherAI", name="pythia-70m"),
+        block_size=2048,
+        n_layer=6,
+        n_embd=512,
+        n_head=8,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-160m/blob/main/config.json
+    dict(
+        name="pythia-160m",
+        hf_config=dict(org="EleutherAI", name="pythia-160m"),
+        block_size=2048,
+        n_layer=12,
+        n_embd=768,
+        n_head=12,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-410m/blob/main/config.json
+    dict(
+        name="pythia-410m",
+        hf_config=dict(org="EleutherAI", name="pythia-410m"),
+        block_size=2048,
+        n_layer=24,
+        n_embd=1024,
+        n_head=16,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-1b/blob/main/config.json
+    dict(
+        name="pythia-1b",
+        hf_config=dict(org="EleutherAI", name="pythia-1b"),
+        block_size=2048,
+        n_embd=2048,
+        n_head=8,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-1.4b/blob/main/config.json
+    dict(
+        name="pythia-1.4b",
+        hf_config=dict(org="EleutherAI", name="pythia-1.4b"),
+        block_size=2048,
+        n_layer=24,
+        n_embd=2048,
+        n_head=16,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-2.8b/blob/main/config.json
+    dict(
+        name="pythia-2.8b",
+        hf_config=dict(org="EleutherAI", name="pythia-2.8b"),
+        block_size=2048,
+        n_layer=32,
+        n_embd=2560,
+        padding_multiple=128,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-6.9b/blob/main/config.json
+    dict(
+        name="pythia-6.9b",
+        hf_config=dict(org="EleutherAI", name="pythia-6.9b"),
+        block_size=2048,
+        n_layer=32,
+        padding_multiple=256,
+    ),
+    # https://huggingface.co/EleutherAI/pythia-12b/blob/main/config.json
+    dict(
+        name="pythia-12b",
+        hf_config=dict(org="EleutherAI", name="pythia-12b"),
+        block_size=2048,
+        n_layer=36,
+        n_embd=5120,
+        n_head=40,
+    ),
+]
+configs.extend(pythia)
+for c in pythia:
+    copy = c.copy()
+    copy["name"] = f"{c['name']}-deduped"
+    copy["hf_config"]["name"] = f"{c['hf_config']['name']}-deduped"
+    configs.append(copy)
+
+
+####################################
+# togethercomputer RedPajama INCITE
+####################################
+redpajama_incite = [
+    # https://huggingface.co/togethercomputer/RedPajama-INCITE-Base-3B-v1/blob/main/config.json
+    dict(
+        name="RedPajama-INCITE-{}-3B-v1",
+        hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-{}-3B-v1"),
+        block_size=2048,
+        n_layer=32,
+        n_embd=2560,
+        padding_multiple=256,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+    ),
+    # https://huggingface.co/togethercomputer/RedPajama-INCITE-7B-Base/blob/main/config.json
+    dict(
+        name="RedPajama-INCITE-7B-{}",
+        hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-7B-{}"),
+        block_size=2048,
+        n_layer=32,
+        padding_multiple=256,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+    ),
+    # this redirects to the checkpoint above. kept for those who had the old weights already downloaded
+    dict(
+        name="RedPajama-INCITE-{}-7B-v0.1",
+        hf_config=dict(org="togethercomputer", name="RedPajama-INCITE-{}-7B-v0.1"),
+        block_size=2048,
+        n_layer=32,
+        padding_multiple=256,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+    ),
+]
+for c in redpajama_incite:
+    for kind in ("Base", "Chat", "Instruct"):
+        copy = c.copy()
+        copy["name"] = c["name"].format(kind)
+        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
+        configs.append(copy)
+
+
+#################
+# TII UAE Falcon
+#################
+falcon = [
+    # https://huggingface.co/tiiuae/falcon-7b/blob/main/config.json
+    dict(
+        name="falcon-7b{}",
+        hf_config=dict(org="tiiuae", name="falcon-7b{}"),
+        block_size=2048,
+        vocab_size=65024,
+        padded_vocab_size=65024,
+        n_layer=32,
+        n_head=71,
+        n_embd=4544,
+        rotary_percentage=1.0,
+        n_query_groups=1,
+        bias=False,
+        # this is not in the config, but in the original model implementation, only for this config
+        shared_attention_norm=True,
+    ),
+    # https://huggingface.co/tiiuae/falcon-40b/blob/main/config.json
+    dict(
+        name="falcon-40b{}",
+        hf_config=dict(org="tiiuae", name="falcon-40b{}"),
+        block_size=2048,
+        vocab_size=65024,
+        padded_vocab_size=65024,
+        n_layer=60,
+        n_head=128,
+        n_embd=8192,
+        rotary_percentage=1.0,
+        n_query_groups=8,
+        bias=False,
+    ),
+]
+for c in falcon:
+    for kind in ("", "-instruct"):
+        copy = c.copy()
+        copy["name"] = c["name"].format(kind)
+        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
+        configs.append(copy)
+
+# https://huggingface.co/tiiuae/falcon-180b/blob/main/config.json
+falcon180b = dict(
+    name="falcon-180B{}",
+    hf_config=dict(org="tiiuae", name="falcon-180B{}"),
+    block_size=2048,
+    vocab_size=65024,
+    padded_vocab_size=65024,
+    n_layer=80,
+    n_head=232,
+    n_embd=14848,
+    rotary_percentage=1.0,
+    n_query_groups=8,
+    bias=False,
+)
+
+for kind in ("", "-chat"):
+    copy = falcon180b.copy()
+    copy["name"] = falcon180b["name"].format(kind)
+    copy["hf_config"]["name"] = falcon180b["hf_config"]["name"].format(kind)
+    configs.append(copy)
+
+
+#############################
+# OpenLM Research Open LLaMA
+#############################
+open_LLaMA = [
+    # https://huggingface.co/openlm-research/open_llama_3b/blob/main/config.json
+    dict(
+        name="open_llama_3b",
+        hf_config=dict(org="openlm-research", name="open_llama_3b"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=26,
+        n_embd=3200,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=8640,
+    ),
+    # https://huggingface.co/openlm-research/open_llama_7b/blob/main/config.json
+    dict(
+        name="open_llama_7b",
+        hf_config=dict(org="openlm-research", name="open_llama_7b"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/openlm-research/open_llama_13b/blob/main/config.json
+    dict(
+        name="open_llama_13b",
+        hf_config=dict(org="openlm-research", name="open_llama_13b"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+]
+configs.extend(open_LLaMA)
+
+
+###############
+# LMSYS Vicuna
+###############
+vicuna = [
+    # https://huggingface.co/lmsys/vicuna-7b-v1.3/blob/main/config.json
+    dict(
+        name="vicuna-7b-v1.3",
+        hf_config=dict(org="lmsys", name="vicuna-7b-v1.3"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/lmsys/vicuna-13b-v1.3/blob/main/config.json
+    dict(
+        name="vicuna-13b-v1.3",
+        hf_config=dict(org="lmsys", name="vicuna-13b-v1.3"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/lmsys/vicuna-33b-v1.3/blob/main/config.json
+    dict(
+        name="vicuna-33b-v1.3",
+        hf_config=dict(org="lmsys", name="vicuna-33b-v1.3"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=60,
+        n_head=52,
+        n_embd=6656,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=17920,
+    ),
+    # https://huggingface.co/lmsys/vicuna-7b-v1.5/blob/main/config.json
+    dict(
+        name="vicuna-7b-v1.5",
+        hf_config=dict(org="lmsys", name="vicuna-7b-v1.5"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/lmsys/vicuna-7b-v1.5-16k/blob/main/config.json
+    dict(
+        name="vicuna-7b-v1.5-16k",
+        hf_config=dict(org="lmsys", name="vicuna-7b-v1.5-16k"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_condense_ratio=4,
+    ),
+    # https://huggingface.co/lmsys/vicuna-13b-v1.5/blob/main/config.json
+    dict(
+        name="vicuna-13b-v1.5",
+        hf_config=dict(org="lmsys", name="vicuna-13b-v1.5"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/lmsys/vicuna-13b-v1.5-16k/blob/main/config.json
+    dict(
+        name="vicuna-13b-v1.5-16k",
+        hf_config=dict(org="lmsys", name="vicuna-13b-v1.5-16k"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+        rope_condense_ratio=4,
+    ),
+]
+configs.extend(vicuna)
+
+
+#################
+# LMSYS LongChat
+#################
+long_chat = [
+    # https://huggingface.co/lmsys/longchat-7b-16k/blob/main/config.json
+    dict(
+        name="longchat-7b-16k",
+        hf_config=dict(org="lmsys", name="longchat-7b-16k"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_condense_ratio=8,
+    ),
+    # https://huggingface.co/lmsys/longchat-13b-16k/blob/main/config.json
+    dict(
+        name="longchat-13b-16k",
+        hf_config=dict(org="lmsys", name="longchat-13b-16k"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+        rope_condense_ratio=8,
+    ),
+]
+configs.extend(long_chat)
+
+
+######################
+# NousResearch Hermes
+######################
+nous_research = [
+    # https://huggingface.co/NousResearch/Nous-Hermes-llama-2-7b/blob/main/config.json
+    dict(
+        name="Nous-Hermes-llama-2-7b",
+        hf_config=dict(org="NousResearch", name="Nous-Hermes-llama-2-7b"),
+        padded_vocab_size=32000,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/NousResearch/Nous-Hermes-13B/blob/main/config.json
+    dict(
+        name="Nous-Hermes-13b",
+        hf_config=dict(org="NousResearch", name="Nous-Hermes-13b"),
+        block_size=2048,
+        vocab_size=32000,
+        padded_vocab_size=32001,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-6,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/NousResearch/Nous-Hermes-Llama2-13b
+    dict(
+        name="Nous-Hermes-Llama2-13b",
+        hf_config=dict(org="NousResearch", name="Nous-Hermes-Llama2-13b"),
+        vocab_size=32000,
+        padded_vocab_size=32032,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+]
+configs.extend(nous_research)
+
+
+###############
+# Meta LLaMA 2
+###############
+llama_2 = [
+    # https://huggingface.co/meta-llama/Llama-2-7b-hf/blob/main/config.json
+    dict(
+        name="Llama-2-7b{}-hf",
+        hf_config=dict(org="meta-llama", name="Llama-2-7b{}-hf"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/meta-llama/Llama-2-13b-hf/blob/main/config.json
+    dict(
+        name="Llama-2-13b{}-hf",
+        hf_config=dict(org="meta-llama", name="Llama-2-13b{}-hf"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/meta-llama/Llama-2-70b-hf/blob/main/config.json
+    dict(
+        name="Llama-2-70b{}-hf",
+        hf_config=dict(org="meta-llama", name="Llama-2-70b{}-hf"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=80,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=28672,
+    ),
+]
+for c in llama_2:
+    for kind in ("", "-chat"):
+        copy = c.copy()
+        copy["name"] = c["name"].format(kind)
+        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
+        configs.append(copy)
+
+
+##########################
+# Stability AI FreeWilly2
+##########################
+freewilly_2 = [
+    # https://huggingface.co/stabilityai/FreeWilly2/blob/main/config.json
+    dict(
+        name="FreeWilly2",
+        hf_config=dict(org="stabilityai", name="FreeWilly2"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=80,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=28672,
+    )
+]
+configs.extend(freewilly_2)
+
+
+##################
+# Meta Code Llama
+##################
+code_llama = [
+    # https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-7b-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-7b-hf"),
+        block_size=16384,
+        vocab_size=32016,
+        padding_multiple=16,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-13b-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-13b-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-13b-hf"),
+        block_size=16384,
+        vocab_size=32016,
+        padding_multiple=16,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-34b-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-34b-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-34b-hf"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=48,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=22016,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-7b-Python-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-7b-Python-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-7b-Python-hf"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-13b-Python-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-13b-Python-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-13b-Python-hf"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-34b-Python-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-34b-Python-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-34b-Python-hf"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=48,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=22016,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-7b-Instruct-hf/tree/main/config.json
+    dict(
+        name="CodeLlama-7b-Instruct-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-7b-Instruct-hf"),
+        block_size=16384,
+        vocab_size=32016,
+        padding_multiple=16,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-13b-Instruct-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-13b-Instruct-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-13b-Instruct-hf"),
+        block_size=2048,
+        vocab_size=32016,
+        padding_multiple=16,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+        rope_base=1000000,
+    ),
+    # https://huggingface.co/codellama/CodeLlama-34b-Instruct-hf/blob/main/config.json
+    dict(
+        name="CodeLlama-34b-Instruct-hf",
+        hf_config=dict(org="codellama", name="CodeLlama-34b-Instruct-hf"),
+        block_size=16384,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=48,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=22016,
+        rope_base=1000000,
+    ),
+]
+configs.extend(code_llama)
+
+
+########################
+# garage-bAInd Platypus
+########################
+platypus = [
+    # https://huggingface.co/garage-bAInd/Platypus-30B/blob/main/config.json
+    dict(
+        name="Platypus-30B",
+        hf_config=dict(org="garage-bAInd", name="Platypus-30B"),
+        block_size=2048,
+        padded_vocab_size=32000,
+        n_layer=60,
+        n_head=52,
+        n_embd=6656,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-06,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=17920,
+    ),
+    # https://huggingface.co/garage-bAInd/Platypus2-7B/blob/main/config.json
+    dict(
+        name="Platypus2-7B",
+        hf_config=dict(org="garage-bAInd", name="Platypus2-7B"),
+        padded_vocab_size=32000,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+    ),
+    # https://huggingface.co/garage-bAInd/Platypus2-13B/blob/main/config.json
+    dict(
+        name="Platypus2-13B",
+        hf_config=dict(org="garage-bAInd", name="Platypus2-13B"),
+        padded_vocab_size=32000,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/garage-bAInd/Platypus2-70B/blob/main/config.json
+    dict(
+        name="Platypus2-70B",
+        hf_config=dict(org="garage-bAInd", name="Platypus2-70B"),
+        padded_vocab_size=32000,
+        n_layer=80,
+        n_head=64,
+        n_embd=8192,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=28672,
+    ),
+    # https://huggingface.co/garage-bAInd/Camel-Platypus2-13B/blob/main/config.json
+    dict(
+        name="Camel-Platypus2-13B",
+        hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-13B"),
+        padded_vocab_size=32000,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/garage-bAInd/Camel-Platypus2-70B/blob/main/config.json
+    dict(
+        name="Camel-Platypus2-70B",
+        hf_config=dict(org="garage-bAInd", name="Camel-Platypus2-70B"),
+        padded_vocab_size=32000,
+        n_layer=80,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=28672,
+    ),
+    # https://huggingface.co/garage-bAInd/Stable-Platypus2-13B/blob/main/config.json
+    dict(
+        name="Stable-Platypus2-13B",
+        hf_config=dict(org="garage-bAInd", name="Stable-Platypus2-13B"),
+        padded_vocab_size=32000,
+        n_layer=40,
+        n_head=40,
+        n_embd=5120,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=13824,
+    ),
+    # https://huggingface.co/garage-bAInd/Platypus2-70B-instruct/blob/main/config.json
+    dict(
+        name="Platypus2-70B-instruct",
+        hf_config=dict(org="garage-bAInd", name="Platypus2-70B-instruct"),
+        padded_vocab_size=32000,
+        n_layer=80,
+        n_head=64,
+        n_embd=8192,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=28672,
+    ),
+]
+configs.extend(platypus)
+
+
+##########################
+# Stability AI StableCode
+##########################
+stablecode = [
+    # https://huggingface.co/stabilityai/stablecode-completion-alpha-3b/blob/main/config.json
+    dict(
+        name="stablecode-completion-alpha-3b",
+        hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b"),
+        block_size=16384,
+        vocab_size=49152,
+        n_layer=32,
+        n_embd=2560,
+    ),
+    # https://huggingface.co/stabilityai/stablecode-completion-alpha-3b-4k/blob/main/config.json
+    dict(
+        name="stablecode-completion-alpha-3b-4k",
+        hf_config=dict(org="stabilityai", name="stablecode-completion-alpha-3b-4k"),
+        vocab_size=49152,
+        n_layer=32,
+        n_embd=2560,
+    ),
+    # https://huggingface.co/stabilityai/stablecode-instruct-alpha-3b/blob/main/config.json
+    dict(
+        name="stablecode-instruct-alpha-3b",
+        hf_config=dict(org="stabilityai", name="stablecode-instruct-alpha-3b"),
+        vocab_size=49152,
+        n_layer=32,
+        n_embd=2560,
+    ),
+]
+configs.extend(stablecode)
+
+
+##################################
+# togethercomputer LLaMA-2-7B-32K
+##################################
+together_llama2_32k = [
+    # https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/main/config.json
+    dict(
+        name="LLaMA-2-7B-32K",
+        hf_config=dict(org="togethercomputer", name="LLaMA-2-7B-32K"),
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=32,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        _mlp_class="LLaMAMLP",
+        intermediate_size=11008,
+        rope_condense_ratio=8,
+    )
+]
+configs.extend(together_llama2_32k)
+
+
+################
+# Microsoft Phi
+################
+phi = [
+    # https://huggingface.co/microsoft/phi-1_5/blob/main/config.json
+    dict(
+        name="phi-1_5",
+        hf_config=dict(org="microsoft", name="phi-1_5"),
+        vocab_size=50257,
+        padded_vocab_size=51200,
+        block_size=2048,
+        n_embd=2048,
+        n_layer=24,
+        rotary_percentage=0.5,  # 32 / (n_embd / n_head) = 32 / 64
+        shared_attention_norm=True,
+        lm_head_bias=True,
+        gelu_approximate="tanh",
+    )
+]
+configs.extend(phi)
+
+
+#############
+# Mistral AI
+#############
+mistral = [
+    # https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json
+    dict(
+        name="Mistral-7B-{}v0.1",
+        hf_config=dict(org="mistralai", name="Mistral-7B-{}v0.1"),
+        padded_vocab_size=32000,
+        block_size=4096,  # should be 32768 but sliding window attention is not implemented
+        n_layer=32,
+        n_query_groups=8,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",
+        norm_eps=1e-05,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=14336,
+    )
+]
+for c in mistral:
+    for kind in ("", "Instruct-"):
+        copy = c.copy()
+        copy["name"] = c["name"].format(kind)
+        copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind)
+        configs.append(copy)
+
+
+############
+# TinyLlama
+############
+tiny_llama = [
+    dict(
+        name="tiny-llama-1.1b",
+        hf_config=dict(org="PY007", name="TinyLlama-1.1B-intermediate-step-480k-1T"),
+        block_size=2048,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=22,
+        n_head=32,
+        n_embd=2048,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",  # original TinyLlama uses FusedRMSNorm
+        norm_eps=1e-5,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=5632,
+        n_query_groups=4,
+    ),
+    dict(
+        name="tiny-llama-new",
+        hf_config=dict(org="PY007", name="TinyLlama-1.1B-intermediate-step-480k-1T"),
+        block_size=768,
+        vocab_size=32000,
+        padding_multiple=64,
+        n_layer=18,
+        n_head=32,
+        n_embd=1024,
+        rotary_percentage=1.0,
+        parallel_residual=False,
+        bias=False,
+        _norm_class="RMSNorm",  # original TinyLlama uses FusedRMSNorm
+        norm_eps=1e-5,
+        _mlp_class="LLaMAMLP",
+        intermediate_size=5632,
+        n_query_groups=4,
+    ),
+]
+configs.extend(tiny_llama)
+
+
+name_to_config = {config["name"]: config for config in configs}

tsai_gpt/generate.ipynb

@@ -0,0 +1,173 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from generate import generate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release\n",
+      "  warnings.warn(\n",
+      "Loading model '/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth' with {'name': 'Llama-2-7b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-7b-chat-hf'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.21 seconds.\n",
+      "Time to load the model weights: 0.29 seconds.\n",
+      "Seed set to 1234\n",
+      "James Bond asked for a Vodka Martini, shaken and 17-year-old in a handheld family home, which he has since met with heavy duty clinton officials.\n",
+      "\"I don't'm an 20-year-old, and I do not want to hurt him because nobody wants him to exist. I've tried to make a real difference and it felt like superbly at the time. I was the only one who busts it around and I had no problem. I've gotten back to 91. \" (One wants to change my name, of course.) I think that I, being married, would be happy to have married after a couple of years, but I'd still not be older. I cannot have anything in common with my old family, as it was the life of the person who had his own brother,\" said Mark Stenyard, a psychiatrist and partner of the Treasury. \"When you told me that I'm divorced here, I'm OK. I am going to tell you that I, being divorced, would be happy to have married. I need one, my family, and I'm still late, but am just so disgusted.\"\n",
+      "Time for inference 1: 1.85 sec total, 135.39 tokens/sec\n",
+      "Memory used: 0.35 GB\n"
+     ]
+    }
+   ],
+   "source": [
+    " !python3 generate.py --prompt=\"James Bond asked for a Vodka Martini, shaken and \" --checkpoint_dir=/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/ --max_new_tokens=250 --temperature=0.9 --num_samples=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release\n",
+      "  warnings.warn(\n",
+      "Loading model '/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth' with {'name': 'Llama-2-7b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-7b-chat-hf'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.20 seconds.\n",
+      "Time to load the model weights: 0.13 seconds.\n",
+      "Seed set to 1234\n",
+      "there is a difference between a finitely generated group 1 and a 10-minute activity set by the 18-point approach. In the group 1, a group 1 is set to take an 20-minute activity set by the 20-point approach. The group 1 is set to take an 20-minute activity set by the 20-point approach.\n",
+      "A theorianized group is a group 1 of a 10-minute activity set by the 20-point approach and the 18-point approach is set to take an 20-minute activity set by the 20-point process. The 20-point approach is set to take an 90-point approach targeted by the 20-point approach. In this group 1, a couple exists to engage between two.\n",
+      "A band 1 is set to take an 20-point focus set by the 20-point approach. The group 1 is set to take an 30-point approach targeting each 10-point journey. Bersecar mode, 2, 20 and 3 are single\n",
+      "Time for inference 1: 1.58 sec total, 157.98 tokens/sec\n",
+      "Memory used: 0.35 GB\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python3 generate.py --prompt=\"there is a difference between a finitely generated group \" --checkpoint_dir=/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/ --max_new_tokens=250 --temperature=0.9 --num_samples=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release\n",
+      "  warnings.warn(\n",
+      "Loading model '/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth' with {'name': 'Llama-2-7b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-7b-chat-hf'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.20 seconds.\n",
+      "Time to load the model weights: 0.12 seconds.\n",
+      "Seed set to 1234\n",
+      "there are torsion-free hyperbolic groups that uniformly 100,000 times a day. The current 18-day study shows that the group’s social bias and racism are more likely to be more than just a high percentage of U.S. citizens. That’s because 10,000 people were exposed to torsion-free absences this year.\n",
+      "Many of the victims are the same age groups. The most recent study in U.S. history suggests that the group may be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be the most likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be\n",
+      "Time for inference 1: 1.51 sec total, 165.87 tokens/sec\n",
+      "Memory used: 0.35 GB\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python3 generate.py --prompt=\"there are torsion-free hyperbolic groups that uniformly \" --checkpoint_dir=/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/ --max_new_tokens=250 --temperature=0.8 --num_samples=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release\n",
+      "  warnings.warn(\n",
+      "Loading model '/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth' with {'name': 'Llama-2-7b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-7b-chat-hf'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.20 seconds.\n",
+      "Time to load the model weights: 0.13 seconds.\n",
+      "Seed set to 1234\n",
+      "Virginia Attorney General Backs Off Ballot Proposal on 17th Amendment\n",
+      "The Supreme Court of Virginia Motors called it unconstitutional because the Supreme Court was justified by the Supreme Court’s decision to leave office. The Court said the U.S. Supreme Court justified the application by the Supreme Court to proceed. The Supreme Court ruled that the Supreme Court’s decision to leave office in the United States was resolved by the Supreme Court’s decision to leave office.\n",
+      "The Supreme Court said the Supreme Court had agreed to leave office in the United States (U.S. District Judge Paul Cablet F.L.V.V.V.V.V.). The Supreme Court of Virginia Motors called the decision by the Supreme Court to leave office in the United States (U.S. District Judge Paul Cablet F.L.V.V.V.V.V.) and the Supreme Court under the U.S. Supreme Court. The Supreme Court also ruled that the Supreme Court’s decision to leave the office in the United States (U.S. District Judge Paul Cablet F.L.V.V.V.V.V.V.V.V.V.V.V.V.\n",
+      "Time for inference 1: 1.82 sec total, 137.00 tokens/sec\n",
+      "Memory used: 0.35 GB\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python3 generate.py --prompt=\"Virginia Attorney General Backs Off Ballot Proposal on \" --checkpoint_dir=/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/ --max_new_tokens=250 --temperature=0.8 --num_samples=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 1.1build1 is an invalid version and will not be supported in a future release\n",
+      "  warnings.warn(\n",
+      "Loading model '/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/lit_model.pth' with {'name': 'Llama-2-7b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-7b-chat-hf'}, 'block_size': 2048, 'vocab_size': 50254, 'padding_multiple': 128, 'padded_vocab_size': 50304, 'n_layer': 12, 'n_head': 12, 'n_embd': 768, 'rotary_percentage': 0.25, 'parallel_residual': True, 'bias': True, 'lm_head_bias': False, 'n_query_groups': 12, 'shared_attention_norm': False, '_norm_class': 'LayerNorm', 'norm_eps': 1e-05, '_mlp_class': 'GptNeoxMLP', 'gelu_approximate': 'none', 'intermediate_size': 3072, 'rope_condense_ratio': 1, 'rope_base': 10000, 'head_size': 64, 'rope_n_elem': 16}\n",
+      "Time to instantiate model: 0.19 seconds.\n",
+      "Time to load the model weights: 0.14 seconds.\n",
+      "Seed set to 1234\n",
+      "there are torsion-free hyperbolic groups that uniformly 100,000 times a day. The current 18-day study shows that the group’s social bias and racism are more likely to be more than just a high percentage of U.S. citizens. That’s because 10,000 people were exposed to torsion-free absences this year.\n",
+      "Many of the victims are the same age groups. The most recent study in U.S. history suggests that the group may be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be the most likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be more likely to be\n",
+      "Time for inference 1: 1.56 sec total, 160.43 tokens/sec\n",
+      "Memory used: 0.35 GB\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python3 generate.py --prompt=\"there are torsion-free hyperbolic groups that uniformly \" --checkpoint_dir=/home/raghu/work/ERA-V1-assignments/assignment-22/checkpoints/meta-llama/Llama-2-7b-chat-hf/ --max_new_tokens=250 --temperature=0.8 --num_samples=1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tsai_gpt/generate.py

@@ -0,0 +1,211 @@
+import sys
+import time
+from pathlib import Path
+from typing import Literal, Optional
+
+import lightning as L
+import torch
+from lightning.fabric.plugins import BitsandbytesPrecision
+from lightning.fabric.strategies import FSDPStrategy
+
+# support running without installing as a package
+wd = Path(__file__).parent.parent.resolve()
+sys.path.append(str(wd))
+
+#from lit_gpt import GPT, Config, Tokenizer
+from config import Config
+from tokenizer import Tokenizer
+from model import GPT, Block
+from utils import (
+    check_valid_checkpoint_dir,
+    get_default_supported_precision,
+    gptq_quantization,
+    load_checkpoint,
+)
+
+
+def sample(logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None) -> torch.Tensor:
+    logits = logits[0, -1]
+    # optionally crop the logits to only the top k options
+    if top_k is not None:
+        v, i = torch.topk(logits, min(top_k, logits.size(-1)))
+        # do not use `torch.where` as in nanogpt because it will repeat top-k collisions
+        logits = torch.full_like(logits, float("-inf")).scatter_(-1, i, v)
+    # optionally scale the logits and sample from a probability distribution
+    if temperature > 0.0:
+        probs = torch.nn.functional.softmax(logits / temperature, dim=-1)
+        return torch.multinomial(probs, num_samples=1)
+    return torch.argmax(logits, dim=-1, keepdim=True)
+
+
+@torch.inference_mode()
+def generate(
+    model: GPT,
+    idx: torch.Tensor,
+    max_returned_tokens: int,
+    *,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    eos_id: Optional[int] = None,
+) -> torch.Tensor:
+    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
+
+    The implementation of this function is modified from A. Karpathy's nanoGPT.
+
+    Args:
+        model: The model to use.
+        idx: Tensor of shape (T) with indices of the prompt sequence.
+        max_returned_tokens: The maximum number of tokens to return (given plus generated).
+        temperature: Scales the predicted logits by 1 / temperature.
+        top_k: If specified, only sample among the tokens with the k highest probabilities.
+        eos_id: If specified, stop generating any more token once the <eos> token is triggered.
+    """
+    T = idx.size(0)
+    assert max_returned_tokens > T
+    if model.max_seq_length < max_returned_tokens - 1:
+        # rolling the kv cache based on the `input_pos` value would be necessary. However, doing so would introduce a
+        # data dependency on the `input_pos` tensor and impact model compilation. Since this setting is uncommon, we do
+        # not support it to avoid negatively impacting the overall speed
+        raise NotImplementedError(f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}")
+
+    device, dtype = idx.device, idx.dtype
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    empty = torch.empty(max_returned_tokens, dtype=dtype, device=device)
+    empty[:T] = idx
+    idx = empty
+    input_pos = torch.arange(0, T, device=device)
+
+    # generate up to a fixed number of tokens
+    for _ in range(max_returned_tokens - T):
+        x = idx.index_select(0, input_pos).view(1, -1)
+
+        # forward
+        logits = model(x, input_pos)
+        idx_next = sample(logits, temperature, top_k).to(dtype=dtype)
+
+        # advance
+        input_pos = input_pos[-1:] + 1
+
+        # concatenate the new generation
+        idx = idx.index_copy(0, input_pos, idx_next)
+
+        # if <eos> token is triggered, return the output (stop generation)
+        if idx_next == eos_id:
+            return idx[:input_pos]  # include the EOS token
+
+    return idx
+
+
+def main(
+    prompt: str = "Hello, my name is",
+    *,
+    num_samples: int = 1,
+    max_new_tokens: int = 50,
+    top_k: Optional[int] = 200,
+    temperature: float = 0.8,
+    checkpoint_dir: Path = Path("checkpoints/stabilityai/stablelm-base-alpha-3b"),
+    quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8", "gptq.int4"]] = None,
+    strategy: str = "auto",
+    devices: int = 1,
+    precision: Optional[str] = None,
+) -> None:
+    """Generates text samples based on a pre-trained model and tokenizer.
+
+    Args:
+        prompt: The prompt string to use for generating the samples.
+        num_samples: The number of text samples to generate.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+        checkpoint_dir: The checkpoint directory to load.
+        quantize: Whether to quantize the model and using which method:
+            - bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes
+            - bnb.int8: 8-bit quantization from bitsandbytes
+            - gptq.int4: 4-bit quantization from GPTQ
+            for more details, see https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/quantize.md
+        strategy: Indicates the Fabric strategy setting to use.
+        devices: How many devices to use.
+        precision: Indicates the Fabric precision setting to use.
+    """
+    precision = precision or get_default_supported_precision(training=False)
+
+    plugins = None
+    if quantize is not None:
+        if devices > 1:
+            raise NotImplementedError(
+                "Quantization is currently not supported for multi-GPU training. Please set devices=1 when using the"
+                " --quantize flag."
+            )
+        if quantize.startswith("bnb."):
+            if "mixed" in precision:
+                raise ValueError("Quantization and mixed precision is not supported.")
+            dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision]
+            plugins = BitsandbytesPrecision(quantize[4:], dtype)
+            precision = None
+
+    if strategy == "fsdp":
+        strategy = FSDPStrategy(auto_wrap_policy={Block}, cpu_offload=False)
+
+    fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy, plugins=plugins)
+    fabric.launch()
+
+    check_valid_checkpoint_dir(checkpoint_dir)
+
+    config = Config.from_json(checkpoint_dir / "lit_config.json")
+
+    if quantize == "gptq.int4":
+        model_file = "lit_model_gptq.4bit.pth"
+        if not (checkpoint_dir / model_file).is_file():
+            raise ValueError("Please run `python quantize/gptq.py` first")
+    else:
+        model_file = "lit_model.pth" #"lit_model.pth"
+    checkpoint_path = checkpoint_dir / model_file
+
+    fabric.print(f"Loading model {str(checkpoint_path)!r} with {config.__dict__}", file=sys.stderr)
+    t0 = time.perf_counter()
+    with fabric.init_module(empty_init=True), gptq_quantization(quantize == "gptq.int4"):
+        model = GPT(config)
+    fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    t0 = time.perf_counter()
+    load_checkpoint(fabric, model, checkpoint_path)
+    fabric.print(f"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.", file=sys.stderr)
+
+    tokenizer = Tokenizer(checkpoint_dir)
+    encoded = tokenizer.encode(prompt, device=fabric.device)
+    prompt_length = encoded.size(0)
+    max_returned_tokens = prompt_length + max_new_tokens
+
+    with fabric.init_tensor():
+        # set the max_seq_length to limit the memory usage to what we need
+        model.max_seq_length = max_returned_tokens
+
+    L.seed_everything(1234)
+    for i in range(num_samples):
+        with fabric.init_tensor():
+            # enable the kv cache
+            model.set_kv_cache(batch_size=1)
+
+        t0 = time.perf_counter()
+        y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k)
+        t = time.perf_counter() - t0
+
+        fabric.print(tokenizer.decode(y))
+        tokens_generated = y.size(0) - prompt_length
+        fabric.print(
+            f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_generated / t:.02f} tokens/sec", file=sys.stderr
+        )
+    if fabric.device.type == "cuda":
+        fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB", file=sys.stderr)
+
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(main)

tsai_gpt/generate_for_app.py

@@ -0,0 +1,204 @@
+import sys
+import time
+from pathlib import Path
+from typing import Literal, Optional
+
+import lightning as L
+import torch
+from lightning.fabric.plugins import BitsandbytesPrecision
+from lightning.fabric.strategies import FSDPStrategy
+
+# support running without installing as a package
+wd = Path(__file__).parent.parent.resolve()
+sys.path.append(str(wd))
+
+#from lit_gpt import GPT, Config, Tokenizer
+from .config import Config
+from .tokenizer import Tokenizer
+from .model import GPT, Block
+from .utils import (
+    check_valid_checkpoint_dir,
+    get_default_supported_precision,
+    gptq_quantization,
+    load_checkpoint,
+)
+
+
+def sample(logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None) -> torch.Tensor:
+    logits = logits[0, -1]
+    # optionally crop the logits to only the top k options
+    if top_k is not None:
+        v, i = torch.topk(logits, min(top_k, logits.size(-1)))
+        # do not use `torch.where` as in nanogpt because it will repeat top-k collisions
+        logits = torch.full_like(logits, float("-inf")).scatter_(-1, i, v)
+    # optionally scale the logits and sample from a probability distribution
+    if temperature > 0.0:
+        probs = torch.nn.functional.softmax(logits / temperature, dim=-1)
+        return torch.multinomial(probs, num_samples=1)
+    return torch.argmax(logits, dim=-1, keepdim=True)
+
+
+@torch.inference_mode()
+def generate(
+    model: GPT,
+    idx: torch.Tensor,
+    max_returned_tokens: int,
+    *,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    eos_id: Optional[int] = None,
+) -> torch.Tensor:
+    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
+
+    The implementation of this function is modified from A. Karpathy's nanoGPT.
+
+    Args:
+        model: The model to use.
+        idx: Tensor of shape (T) with indices of the prompt sequence.
+        max_returned_tokens: The maximum number of tokens to return (given plus generated).
+        temperature: Scales the predicted logits by 1 / temperature.
+        top_k: If specified, only sample among the tokens with the k highest probabilities.
+        eos_id: If specified, stop generating any more token once the <eos> token is triggered.
+    """
+    T = idx.size(0)
+    assert max_returned_tokens > T
+    if model.max_seq_length < max_returned_tokens - 1:
+        # rolling the kv cache based on the `input_pos` value would be necessary. However, doing so would introduce a
+        # data dependency on the `input_pos` tensor and impact model compilation. Since this setting is uncommon, we do
+        # not support it to avoid negatively impacting the overall speed
+        raise NotImplementedError(f"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}")
+
+    device, dtype = idx.device, idx.dtype
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    empty = torch.empty(max_returned_tokens, dtype=dtype, device=device)
+    empty[:T] = idx
+    idx = empty
+    input_pos = torch.arange(0, T, device=device)
+
+    # generate up to a fixed number of tokens
+    for _ in range(max_returned_tokens - T):
+        x = idx.index_select(0, input_pos).view(1, -1)
+
+        # forward
+        logits = model(x, input_pos)
+        idx_next = sample(logits, temperature, top_k).to(dtype=dtype)
+
+        # advance
+        input_pos = input_pos[-1:] + 1
+
+        # concatenate the new generation
+        idx = idx.index_copy(0, input_pos, idx_next)
+
+        # if <eos> token is triggered, return the output (stop generation)
+        if idx_next == eos_id:
+            return idx[:input_pos]  # include the EOS token
+
+    return idx
+
+
+def generate_for_app(
+    prompt: str = "Hello, my name is",
+    *,
+    num_samples: int = 1,
+    max_new_tokens: int = 50,
+    top_k: Optional[int] = 200,
+    temperature: float = 0.8,
+    checkpoint_dir: Path = Path("checkpoints/stabilityai/stablelm-base-alpha-3b"),
+    quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8", "gptq.int4"]] = None,
+    strategy: str = "auto",
+    devices: int = 1,
+    precision: Optional[str] = None,
+) -> str:
+    """Generates text samples based on a pre-trained model and tokenizer.
+
+    Args:
+        prompt: The prompt string to use for generating the samples.
+        num_samples: The number of text samples to generate.
+        max_new_tokens: The number of generation steps to take.
+        top_k: The number of top most probable tokens to consider in the sampling process.
+        temperature: A value controlling the randomness of the sampling process. Higher values result in more random
+            samples.
+        checkpoint_dir: The checkpoint directory to load.
+        quantize: Whether to quantize the model and using which method:
+            - bnb.nf4, bnb.nf4-dq, bnb.fp4, bnb.fp4-dq: 4-bit quantization from bitsandbytes
+            - bnb.int8: 8-bit quantization from bitsandbytes
+            - gptq.int4: 4-bit quantization from GPTQ
+            for more details, see https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/quantize.md
+        strategy: Indicates the Fabric strategy setting to use.
+        devices: How many devices to use.
+        precision: Indicates the Fabric precision setting to use.
+    """
+    precision = precision or get_default_supported_precision(training=False)
+
+    plugins = None
+    if quantize is not None:
+        if devices > 1:
+            raise NotImplementedError(
+                "Quantization is currently not supported for multi-GPU training. Please set devices=1 when using the"
+                " --quantize flag."
+            )
+        if quantize.startswith("bnb."):
+            if "mixed" in precision:
+                raise ValueError("Quantization and mixed precision is not supported.")
+            dtype = {"16-true": torch.float16, "bf16-true": torch.bfloat16, "32-true": torch.float32}[precision]
+            plugins = BitsandbytesPrecision(quantize[4:], dtype)
+            precision = None
+
+
+    fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy, plugins=plugins)
+    fabric.launch()
+
+    check_valid_checkpoint_dir(checkpoint_dir)
+
+    config = Config.from_json(checkpoint_dir / "lit_config.json")
+
+    if quantize == "gptq.int4":
+        model_file = "lit_model_gptq.4bit.pth"
+        if not (checkpoint_dir / model_file).is_file():
+            raise ValueError("Please run `python quantize/gptq.py` first")
+    else:
+        model_file = "lit_model.pth" #"lit_model.pth"
+    checkpoint_path = checkpoint_dir / model_file
+
+    fabric.print(f"Loading model {str(checkpoint_path)!r} with {config.__dict__}", file=sys.stderr)
+    with fabric.init_module(empty_init=True), gptq_quantization(quantize == "gptq.int4"):
+        model = GPT(config)
+
+    model.eval()
+    model = fabric.setup_module(model)
+
+    tokenizer = Tokenizer(checkpoint_dir)
+    encoded = tokenizer.encode(prompt, device=fabric.device)
+    prompt_length = encoded.size(0)
+    max_returned_tokens = prompt_length + max_new_tokens
+
+    with fabric.init_tensor():
+        # set the max_seq_length to limit the memory usage to what we need
+        model.max_seq_length = max_returned_tokens
+
+    L.seed_everything(1234)
+    generated_text = []
+    for i in range(num_samples):
+        with fabric.init_tensor():
+            # enable the kv cache
+            model.set_kv_cache(batch_size=1)
+        y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k)
+
+        generated_text.append(tokenizer.decode(y))
+        tokens_generated = y.size(0) - prompt_length
+
+
+    if fabric.device.type == "cuda":
+        fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB", file=sys.stderr)
+
+    fabric.print(' '.join(generated_text))
+
+    return ' '.join(generated_text)
+
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    torch.set_float32_matmul_precision("high")
+    CLI(generate_for_app)

tsai_gpt/model.py

@@ -0,0 +1,342 @@
+"""Full definition of a GPT NeoX Language Model, all of it in this single file.
+
+Based on the nanoGPT implementation: https://github.com/karpathy/nanoGPT and
+https://github.com/EleutherAI/gpt-neox/tree/main/megatron/model.
+"""
+import math
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from typing_extensions import Self
+
+from tsai_gpt.config import Config
+
+
+
+class GPT(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        assert config.padded_vocab_size is not None
+        self.config = config
+
+        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=config.lm_head_bias)
+        self.transformer = nn.ModuleDict(
+            dict(
+                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
+                h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
+                ln_f=config.norm_class(config.n_embd, eps=config.norm_eps),
+            )
+        )
+        self.max_seq_length = self.config.block_size
+        self.mask_cache: Optional[torch.Tensor] = None
+
+    @property
+    def max_seq_length(self) -> int:
+        return self._max_seq_length
+
+    @max_seq_length.setter
+    def max_seq_length(self, value: int) -> None:
+        """
+        When doing inference, the sequences used might be shorter than the model's context length.
+        This allows setting a smaller number to avoid allocating unused memory
+        """
+        if value > self.config.block_size:
+            raise ValueError(f"Cannot attend to {value}, block size is only {self.config.block_size}")
+        self._max_seq_length = value
+        if not hasattr(self, "cos"):
+            # first call
+            cos, sin = self.rope_cache()
+            self.register_buffer("cos", cos, persistent=False)
+            self.register_buffer("sin", sin, persistent=False)
+        elif value != self.cos.size(0):
+            # override
+            self.cos, self.sin = self.rope_cache(device=self.cos.device)
+        # the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know
+        # if the kv cache is expected
+
+    def reset_parameters(self) -> None:
+        # Trigger resetting the rope-cache
+        self.max_seq_length = self.config.block_size
+
+    def _init_weights(self, module: nn.Module) -> None:
+        """Meant to be used with `gpt.apply(gpt._init_weights)`."""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx: torch.Tensor, input_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
+        T = idx.size(1)
+        if self.max_seq_length < T:
+            raise ValueError(f"Cannot forward sequence of length {T}, max seq length is only {self.max_seq_length}.")
+
+        if input_pos is not None:  # use the kv cache
+            cos = self.cos.index_select(0, input_pos)
+            sin = self.sin.index_select(0, input_pos)
+            if self.mask_cache is None:
+                raise TypeError("You need to call `gpt.set_kv_cache()`")
+            mask = self.mask_cache.index_select(2, input_pos)
+        else:
+            cos = self.cos[:T]
+            sin = self.sin[:T]
+            mask = None
+
+        x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
+        for block in self.transformer.h:
+            x = block(x, cos, sin, mask, input_pos)
+        x = self.transformer.ln_f(x)
+        return self.lm_head(x)  # (b, t, vocab_size)
+
+    @classmethod
+    def from_name(cls, name: str, **kwargs: Any) -> Self:
+        return cls(Config.from_name(name, **kwargs))
+
+    def rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        return build_rope_cache(
+            seq_len=self.max_seq_length,
+            n_elem=self.config.rope_n_elem,
+            device=device,
+            condense_ratio=self.config.rope_condense_ratio,
+            base=self.config.rope_base,
+        )
+
+    def set_kv_cache(
+        self,
+        batch_size: int,
+        rope_cache_length: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        if rope_cache_length is None:
+            rope_cache_length = self.cos.size(-1)
+        max_seq_length = self.max_seq_length
+
+        # initialize the kv cache for all blocks
+        for block in self.transformer.h:
+            block.attn.kv_cache = block.attn.build_kv_cache(
+                batch_size, max_seq_length, rope_cache_length, device, dtype
+            )
+
+        if self.mask_cache is None or self.mask_cache.size(3) != max_seq_length:
+            # passing `attn_mask` to SDPA downgrades it to use the inefficient implementation. since we only need the mask
+            # for the kv-cache support (only during inference), we only create it in that situation
+            # this will be resolved by https://github.com/pytorch/pytorch/issues/96099
+            ones = torch.ones((max_seq_length, max_seq_length), device=device, dtype=torch.bool)
+            self.mask_cache = torch.tril(ones).unsqueeze(0).unsqueeze(0)
+
+    def clear_kv_cache(self) -> None:
+        self.mask_cache = None
+        for block in self.transformer.h:
+            block.attn.kv_cache = None
+
+
+class Block(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.norm_1 = config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.attn = CausalSelfAttention(config)
+        self.norm_2 = None if config.shared_attention_norm else config.norm_class(config.n_embd, eps=config.norm_eps)
+        self.mlp = config.mlp_class(config)
+
+        self.config = config
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        n_1 = self.norm_1(x)
+        h = self.attn(n_1, cos, sin, mask, input_pos)
+        if self.config.parallel_residual:
+            n_2 = n_1 if self.config.shared_attention_norm else self.norm_2(x)
+            x = self.mlp(n_2) + h + x
+        else:
+            if self.config.shared_attention_norm:
+                raise NotImplementedError(
+                    "No checkpoint amongst the ones we support uses this configuration"
+                    " (non-parallel residual and shared attention norm)."
+                )
+            x = h + x
+            x = self.mlp(self.norm_2(x)) + x
+        return x
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        shape = (config.n_head + 2 * config.n_query_groups) * config.head_size
+        # key, query, value projections for all heads, but in a batch
+        self.attn = nn.Linear(config.n_embd, shape, bias=config.bias)
+        # output projection
+        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # disabled by default
+        self.kv_cache: Optional[KVCache] = None
+
+        self.config = config
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+
+        qkv = self.attn(x)
+
+        # assemble into a number of query groups to support MHA, MQA and GQA together (see `config.n_query_groups`)
+        q_per_kv = self.config.n_head // self.config.n_query_groups
+        total_qkv = q_per_kv + 2  # each group has 1+ queries, 1 key, and 1 value
+        qkv = qkv.view(B, T, self.config.n_query_groups, total_qkv, self.config.head_size)
+        qkv = qkv.permute(0, 2, 3, 1, 4)  # (B, n_query_groups, total_qkv, T, hs)
+
+        # split batched computation into three
+        q, k, v = qkv.split((q_per_kv, 1, 1), dim=2)
+
+        # maybe repeat k and v if for the non multi-head attention cases
+        # training: flash attention requires it
+        # inference: multi-query would require a full kv cache so avoid it to limit its memory usage
+        if self.config.n_query_groups != self.config.n_head and (input_pos is None or self.config.n_query_groups != 1):
+            k = k.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
+            v = v.expand(B, self.config.n_query_groups, q_per_kv, T, self.config.head_size)
+
+        q = q.reshape(B, -1, T, self.config.head_size)  # (B, nh_q, T, hs)
+        k = k.reshape(B, -1, T, self.config.head_size)  # (B, nh_k, T, hs)
+        v = v.reshape(B, -1, T, self.config.head_size)  # (B, nh_v, T, hs)
+
+        q_roped = apply_rope(q[..., : self.config.rope_n_elem], cos, sin)
+        k_roped = apply_rope(k[..., : self.config.rope_n_elem], cos, sin)
+        q = torch.cat((q_roped, q[..., self.config.rope_n_elem :]), dim=-1)
+        k = torch.cat((k_roped, k[..., self.config.rope_n_elem :]), dim=-1)
+
+        if input_pos is not None:
+            if not isinstance(self.kv_cache, KVCache):
+                raise TypeError("You need to call `gpt.set_kv_cache()`")
+            k, v = self.kv_cache(input_pos, k, v)
+
+        y = self.scaled_dot_product_attention(q, k, v, mask)
+
+        y = y.reshape(B, T, C)  # re-assemble all head outputs side by side
+
+        # output projection
+        return self.proj(y)
+
+    def scaled_dot_product_attention(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        scale = 1.0 / math.sqrt(self.config.head_size)
+        y = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=mask, dropout_p=0.0, scale=scale, is_causal=mask is None
+        )
+        return y.transpose(1, 2)
+
+    def build_kv_cache(
+        self,
+        batch_size: int,
+        max_seq_length: int,
+        rope_cache_length: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "KVCache":
+        heads = 1 if self.config.n_query_groups == 1 else self.config.n_head
+        v_shape = (batch_size, heads, max_seq_length, self.config.head_size)
+        if rope_cache_length is None:
+            if self.config.rotary_percentage != 1.0:
+                raise TypeError("Please pass the `rope_cache_length=gpt.cos.size(-1)` value")
+            k_shape = v_shape
+        else:
+            k_shape = (
+                batch_size,
+                heads,
+                max_seq_length,
+                rope_cache_length + self.config.head_size - self.config.rope_n_elem,
+            )
+        return KVCache(k_shape, v_shape, device=device, dtype=dtype)
+
+
+class GptNeoxMLP(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.fc = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
+        self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
+
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc(x)
+        x = torch.nn.functional.gelu(x, approximate=self.config.gelu_approximate)
+        return self.proj(x)
+
+
+class LLaMAMLP(nn.Module):
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+        self.fc_1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
+        self.fc_2 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
+        self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_fc_1 = self.fc_1(x)
+        x_fc_2 = self.fc_2(x)
+        x = torch.nn.functional.silu(x_fc_1) * x_fc_2
+        return self.proj(x)
+
+
+def build_rope_cache(
+    seq_len: int, n_elem: int, device: Optional[torch.device] = None, base: int = 10000, condense_ratio: int = 1
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Enhanced Transformer with Rotary Position Embedding.
+
+    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
+    transformers/rope/__init__.py. MIT License:
+    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
+    """
+    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
+    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))
+
+    # Create position indexes `[0, 1, ..., seq_len - 1]`
+    seq_idx = torch.arange(seq_len, device=device) / condense_ratio
+
+    # Calculate the product of position index and $\theta_i$
+    idx_theta = torch.outer(seq_idx, theta).repeat(1, 2)
+
+    return torch.cos(idx_theta), torch.sin(idx_theta)
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    head_size = x.size(-1)
+    x1 = x[..., : head_size // 2]  # (B, nh, T, hs/2)
+    x2 = x[..., head_size // 2 :]  # (B, nh, T, hs/2)
+    rotated = torch.cat((-x2, x1), dim=-1)  # (B, nh, T, hs)
+    roped = (x * cos) + (rotated * sin)
+    return roped.type_as(x)
+
+
+class KVCache(nn.Module):
+    def __init__(
+        self,
+        k_shape: Tuple[int, int, int, int],
+        v_shape: Tuple[int, int, int, int],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> None:
+        super().__init__()
+        self.register_buffer("k", torch.zeros(k_shape, device=device, dtype=dtype), persistent=False)
+        self.register_buffer("v", torch.zeros(v_shape, device=device, dtype=dtype), persistent=False)
+
+    def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        # move the buffer to the activation dtype for when AMP is used
+        self.k = self.k.to(k.dtype)
+        self.v = self.v.to(v.dtype)
+        # update the cache
+        k = self.k.index_copy_(2, input_pos, k)
+        v = self.v.index_copy_(2, input_pos, v)
+        return k, v

tsai_gpt/packed_dataset.py

@@ -0,0 +1,235 @@
+# Very loosely inspired by indexed_dataset in Fairseq, Megatron
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py
+
+
+import os
+import random
+import struct
+
+import numpy as np
+import torch
+from torch.utils.data import IterableDataset, get_worker_info
+
+dtypes = {1: np.uint8, 2: np.int8, 3: np.int16, 4: np.int32, 5: np.int64, 6: np.float32, 7: np.float64, 8: np.uint16}
+
+
+def code(dtype):
+    for k in dtypes:
+        if dtypes[k] == dtype:
+            return k
+    raise ValueError(dtype)
+
+
+HDR_MAGIC = b"LITPKDS"
+HDR_SIZE = 24  # bytes
+
+
+class PackedDataset(IterableDataset):
+    def __init__(
+        self, filenames, n_chunks, block_size, seed=12345, shuffle=True, wrap=False, num_processes=1, process_rank=0
+    ):
+        self._filenames = filenames
+        self._n_chunks = n_chunks
+        self._block_size = block_size
+        self._seed = seed
+        self._shuffle = shuffle
+        self._wrap = wrap
+        self._num_processes = num_processes
+        self._process_rank = process_rank
+
+    def __iter__(self):
+        worker_info = get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        num_shards = num_workers * self._num_processes
+        shard_id = self._process_rank * num_workers + worker_id
+
+        max_num_files = len(self._filenames) // num_shards * num_shards
+        filenames = self._filenames[shard_id:max_num_files:num_shards]
+
+        return PackedDatasetIterator(
+            filenames=filenames,
+            n_chunks=self._n_chunks,
+            block_size=self._block_size,
+            seed=self._seed,
+            shuffle=self._shuffle,
+            wrap=self._wrap,
+        )
+
+
+class PackedDatasetBuilder(object):
+    def __init__(self, outdir, prefix, chunk_size, sep_token, dtype="auto", vocab_size=None):
+        if dtype == "auto":
+            if vocab_size is None:
+                raise ValueError("vocab_size cannot be None when dtype='auto'")
+            if vocab_size is not None and vocab_size < 65500:
+                self._dtype = np.uint16
+            else:
+                self._dtype = np.int32
+        else:
+            self._dtype = dtype
+        self._counter = 0
+        self._chunk_size = chunk_size
+        self._outdir = outdir
+        self._prefix = prefix
+        self._sep_token = sep_token
+        self._arr = np.zeros(self._chunk_size, dtype=self._dtype)
+        self._arr.fill(self._sep_token)
+        self._idx = 0
+        self._version = 1
+        self._filenames = []
+
+    def _write_chunk(self):
+        filename = f"{self._prefix}_{self._counter:010d}.bin"
+        filename = os.path.join(self._outdir, filename)
+
+        with open(filename, "wb") as f:
+            f.write(HDR_MAGIC)
+            f.write(struct.pack("<Q", self._version))
+            f.write(struct.pack("<B", code(self._dtype)))
+            f.write(struct.pack("<Q", self._chunk_size))
+            f.write(self._arr.tobytes(order="C"))
+
+        self._filenames.append(filename)
+        self._counter += 1
+        self._arr.fill(self._sep_token)
+        self._idx = 0
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @property
+    def filenames(self):
+        return self._filenames.copy()
+
+    def add_array(self, arr):
+        while self._idx + arr.shape[0] > self._chunk_size:
+            part_len = self._chunk_size - self._idx
+            self._arr[self._idx : self._idx + part_len] = arr[:part_len]
+            self._write_chunk()
+            arr = arr[part_len:]
+
+        arr_len = arr.shape[0]
+        self._arr[self._idx : self._idx + arr_len] = arr
+        self._idx += arr_len
+
+    def write_reminder(self):
+        self._write_chunk()
+
+
+class PackedDatasetIterator:
+    def __init__(self, filenames, n_chunks, block_size, seed, shuffle, wrap):
+        self._seed = seed
+        self._shuffle = shuffle
+        self._rng = np.random.default_rng(seed) if shuffle else None
+        self._block_idxs = None
+
+        self._wrap = wrap
+
+        # TODO: instead of filenames, we could have a single text stream
+        #       (or text file) with the sequence of all files to be
+        #       fetched/loaded.
+        self._filenames = filenames
+        self._file_idx = 0
+
+        self._n_chunks = n_chunks
+
+        self._dtype = None
+        self._block_size = block_size
+        self._n_blocks = None
+
+        self._mmaps = []
+        self._buffers = []
+
+        self._block_idxs = []
+        self._curr_idx = 0
+
+        self._load_n_chunks()
+
+    def _read_header(self, path):
+        with open(path, "rb") as f:
+            magic = f.read(len(HDR_MAGIC))
+            assert magic == HDR_MAGIC, "File doesn't match expected format."
+            version = struct.unpack("<Q", f.read(8))
+            assert version == (1,)
+            (dtype_code,) = struct.unpack("<B", f.read(1))
+            dtype = dtypes[dtype_code]
+            (chunk_size,) = struct.unpack("<Q", f.read(8))
+        return dtype, chunk_size
+
+    def _close_mmaps(self):
+        for mmap in self._mmaps:
+            mmap._mmap.close()
+
+    def _load_n_chunks(self):
+        self._close_mmaps()
+        self._mmaps = []
+        self._buffers = []
+
+        if self._n_chunks > len(self._filenames[self._file_idx :]):
+            if not self._wrap:
+                raise StopIteration
+            self._file_idx = 0
+
+        for i in range(self._n_chunks):
+            filename = self._filenames[self._file_idx + i]
+            if self._dtype is None:
+                self._dtype, self._chunk_size = self._read_header(filename)
+                self._n_blocks = self._chunk_size // self._block_size
+            # TODO: check header matches with previous files
+            mmap = np.memmap(filename, mode="r", order="C", offset=HDR_SIZE)
+            self._mmaps.append(mmap)
+            self._buffers.append(memoryview(mmap))
+
+        self._file_idx += self._n_chunks
+        n_all_blocks = self._n_chunks * self._n_blocks
+
+        self._block_idxs = self._rng.permutation(n_all_blocks) if self._shuffle else range(n_all_blocks)
+
+        self._curr_idx = 0
+
+    def __del__(self):
+        self._close_mmaps()
+        del self._mmaps
+        del self._buffers
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._curr_idx >= len(self._block_idxs):
+            self._load_n_chunks()
+            # TODO: trigger fetching next next n_chunks if remote
+        block_idx = self._block_idxs[self._curr_idx]
+        chunk_id = block_idx // self._n_blocks
+        buffer = self._buffers[chunk_id]
+        elem_id = (block_idx % self._n_blocks) * self._block_size
+        offset = np.dtype(self._dtype).itemsize * elem_id
+        arr = np.frombuffer(buffer, dtype=self._dtype, count=self._block_size, offset=offset)
+        self._curr_idx += 1
+        return torch.from_numpy(arr.astype(np.int64))
+
+
+class CombinedDataset(IterableDataset):
+    def __init__(self, datasets, seed, weights=None):
+        self._seed = seed
+        self._datasets = datasets
+        self._weights = weights
+        n_datasets = len(datasets)
+        if weights is None:
+            self._weights = [1 / n_datasets] * n_datasets
+
+    def __iter__(self):
+        return CombinedDatasetIterator(self._datasets, self._seed, self._weights)
+
+
+class CombinedDatasetIterator:
+    def __init__(self, datasets, seed, weights):
+        self._datasets = [iter(el) for el in datasets]
+        self._weights = weights
+        self._rng = random.Random(seed)
+
+    def __next__(self):
+        (dataset,) = self._rng.choices(self._datasets, weights=self._weights, k=1)
+        return next(dataset)

tsai_gpt/rmsnorm.py

@@ -0,0 +1,26 @@
+import torch
+
+
+class RMSNorm(torch.nn.Module):
+    """Root Mean Square Layer Normalization.
+
+    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
+    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
+    """
+
+    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(size))
+        self.eps = eps
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x = x.float()
+        # NOTE: the original RMSNorm paper implementation is not equivalent
+        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
+        x_normed = x * torch.rsqrt(norm_x + self.eps)
+        return (self.weight * x_normed).to(dtype=dtype)
+
+    def reset_parameters(self) -> None:
+        torch.nn.init.ones_(self.weight)

tsai_gpt/speed_monitor.py

@@ -0,0 +1,425 @@
+import time
+from collections import deque
+from contextlib import nullcontext
+from typing import Any, Callable, Deque, Dict, Optional
+
+import torch
+from lightning import Callback, Fabric, LightningModule, Trainer
+from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1
+from lightning.fabric.plugins import (
+    BitsandbytesPrecision,
+    DoublePrecision,
+    FSDPPrecision,
+    HalfPrecision,
+    MixedPrecision,
+    Precision,
+    TransformerEnginePrecision,
+    XLAPrecision,
+)
+from lightning.fabric.utilities.rank_zero import rank_zero_only as fabric_rank_zero_only
+from lightning.pytorch.plugins import (
+    DoublePrecisionPlugin,
+    FSDPPrecisionPlugin,
+    HalfPrecisionPlugin,
+    MixedPrecisionPlugin,
+    XLAPrecisionPlugin,
+)
+from lightning.pytorch.utilities.rank_zero import rank_zero_only as trainer_rank_zero_only
+from torch.utils.flop_counter import FlopCounterMode
+
+from tsai_gpt import GPT
+from tsai_gpt.utils import num_parameters
+
+GPU_AVAILABLE_FLOPS = {
+    # source: https://resources.nvidia.com/en-us-tensor-core/nvidia-tensor-core-gpu-datasheet
+    # nvidia publishes spec sheet with a 2x sparsity factor
+    "h100-sxm": {
+        torch.float64: 67e12,
+        torch.float32: 67e12,
+        torch.bfloat16: 1.979e15 / 2,
+        torch.float16: 1.979e15 / 2,
+        torch.int8: 3.958e15 / 2,
+    },
+    "h100-pcie": {
+        torch.float64: 51e12,
+        torch.float32: 51e12,
+        torch.bfloat16: 1.513e15 / 2,
+        torch.float16: 1.513e15 / 2,
+        torch.int8: 3.026e15 / 2,
+    },
+    # source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf
+    # sxm and pcie have same flop counts
+    "a100": {torch.float64: 19.5e12, torch.float32: 19.5e12, torch.bfloat16: 312e12, torch.float16: 312e12},
+    # source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a10/pdf/a10-datasheet.pdf
+    "a10g": {torch.float32: 31.2e12, torch.bfloat16: 125e12, torch.float16: 125e12},
+    # source: https://images.nvidia.com/content/technologies/volta/pdf/volta-v100-datasheet-update-us-1165301-r5.pdf
+    "v100-sxm": {torch.float64: 7.8e12, torch.float32: 15.7e12, torch.float16: 125e12},
+    "v100-pcie": {torch.float64: 7e12, torch.float32: 14e12, torch.float16: 112e12},
+    "v100s-pcie": {torch.float64: 8.2e12, torch.float32: 16.4e12, torch.float16: 130e12},
+    # source: https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/tesla-t4/t4-tensor-core-datasheet-951643.pdf
+    # sxm and pcie have same flop counts
+    "t4": {torch.float32: 8.1e12, torch.float16: 65e12, torch.int8: 130e12},
+    # https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/quadro-product-literature/quadro-rtx-5000-data-sheet-us-nvidia-704120-r4-web.pdf
+    "quadro rtx 5000": {torch.float32: 11.2e12, torch.float16: 89.2e12},
+}
+
+TPU_AVAILABLE_FLOPS = {
+    # flop count for each TPU generation is the same for all precisions
+    # since bfloat16 precision is always used for performing matrix operations
+    # for more info: https://cloud.google.com/tpu/docs/bfloat16#choosing_bfloat16
+    # source: https://arxiv.org/pdf/1907.10701.pdf
+    "v2": 45e12,
+    # source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v3
+    "v3": 123e12,
+    # source: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v4
+    "v4": 275e12,
+    # source: https://cloud.google.com/tpu/docs/v5e-training
+    "v5litepod": 197e12,
+}
+
+
+def get_flops_available(device: torch.device, dtype: torch.dtype) -> Optional[float]:
+    if device.type == "cuda":
+        device_name = torch.cuda.get_device_name(device).lower()
+        if "h100" in device_name and "hbm3" in device_name:
+            device_name = "h100-sxm"
+        elif "h100" in device_name and ("pcie" in device_name or "hbm2e" in device_name):
+            device_name = "h100-pcie"
+        elif "a100" in device_name:
+            device_name = "a100"
+        elif "a10g" in device_name:
+            device_name = "a10g"
+        elif "v100-sxm" in device_name:
+            device_name = "v100-sxm"
+        elif "v100-pcie" in device_name:
+            device_name = "v100-pcie"
+        elif "t4" in device_name:
+            device_name = "t4"
+        elif "quadro rtx 5000" in device_name:
+            device_name = "quadro rtx 5000"
+        else:
+            device_name = None
+
+        if device_name is not None:
+            try:
+                return int(GPU_AVAILABLE_FLOPS[device_name][dtype])
+            except KeyError:
+                raise KeyError(
+                    f"flop count not found for {device_name} with dtype: {dtype}; "
+                    "MFU cannot be calculated and reported."
+                )
+    elif device.type == "xla":
+        if _XLA_GREATER_EQUAL_2_1:
+            from torch_xla._internal import tpu
+        else:
+            from torch_xla.experimental import tpu
+
+        device_name = tpu.get_tpu_env()["TYPE"].lower()
+        try:
+            return int(TPU_AVAILABLE_FLOPS[device_name])
+        except KeyError:
+            raise KeyError(
+                f"flop count not found for {device_name} with dtype: {dtype}; MFU cannot be calculated and reported."
+            )
+
+    return None
+
+
+# Adapted from https://github.com/mosaicml/composer/blob/f2a2dc820cb75023b9eb7c46fdfd25273712abd0/composer/callbacks/speed_monitor.py
+
+
+class SpeedMonitorBase:
+    """Logs the training throughput and utilization.
+
+    +-------------------------------------+-----------------------------------------------------------+
+    | Key                                 | Logged data                                               |
+    +=====================================+===========================================================+
+    |                                     | Rolling average (over `window_size` most recent           |
+    | `throughput/batches_per_sec`        | batches) of the number of batches processed per second    |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | Rolling average (over `window_size` most recent           |
+    | `throughput/samples_per_sec`        | batches) of the number of samples processed per second    |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | Rolling average (over `window_size` most recent           |
+    | `throughput/tokens_per_sec`         | batches) of the number of tokens processed per second.    |
+    |                                     | This may include padding depending on dataset             |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | Estimates flops by `flops_per_batch * batches_per_sec`    |
+    | `throughput/flops_per_sec`          |                                                           |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    | `throughput/device/batches_per_sec` | `throughput/batches_per_sec` divided by world size        |
+    +-------------------------------------+-----------------------------------------------------------+
+    | `throughput/device/samples_per_sec` | `throughput/samples_per_sec` divided by world size        |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | `throughput/tokens_per_sec` divided by world size. This   |
+    | `throughput/device/tokens_per_sec`  | may include pad tokens depending on dataset               |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | `throughput/flops_per_sec` divided by world size. Only    |
+    | `throughput/device/flops_per_sec`   | logged when model has attribute `flops_per_batch`         |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    |                                     | `throughput/device/flops_per_sec` divided by world size.  |
+    | `throughput/device/mfu`             |                                                           |
+    |                                     |                                                           |
+    +-------------------------------------+-----------------------------------------------------------+
+    | `time/train`                        | Total elapsed training time                               |
+    +-------------------------------------+-----------------------------------------------------------+
+    | `time/val`                          | Total elapsed validation time                             |
+    +-------------------------------------+-----------------------------------------------------------+
+    | `time/total`                        | Total elapsed time (time/train + time/val)                |
+    +-------------------------------------+-----------------------------------------------------------+
+
+    Notes:
+        - The implementation assumes that devices are homogeneous as it normalizes by the world size.
+        - Tokens/sec, flops/sec and MFU do not account for padding tokens if present. We suggest using samples/sec or
+          batches/sec to measure throughput under this circumstance.
+        - Be careful when comparing MFU numbers across projects, as this will highly depend on the ``flops_per_batch``.
+          There is no widespread, realistic, and reliable implementation to compute them.
+          We suggest using our ``measure_flops`` function, but many other works will use ``estimated_flops`` which
+          will almost always be an overestimate when compared to the true value.
+
+    Args:
+        window_size (int, optional): Number of batches to use for a rolling average of throughput.
+            Defaults to 100.
+        time_unit (str, optional): Time unit to use for `time` logging. Can be one of
+            'seconds', 'minutes', 'hours', or 'days'. Defaults to 'hours'.
+    """
+
+    def __init__(
+        self,
+        flops_available: float,
+        log_dict: Callable[[Dict, int], None],
+        window_size: int = 100,
+        time_unit: str = "hours",
+    ):
+        self.flops_available = flops_available
+        self.log_dict = log_dict
+
+        # Track the batch num samples and wct to compute throughput over a window of batches
+        self.history_samples: Deque[int] = deque(maxlen=window_size + 1)
+        self.history_wct: Deque[float] = deque(maxlen=window_size + 1)
+        self.history_lengths: Deque[int] = deque(maxlen=window_size + 1)
+        self.history_flops: Deque[int] = deque(maxlen=window_size + 1)
+
+        self.divider = 1
+        if time_unit == "seconds":
+            self.divider = 1
+        elif time_unit == "minutes":
+            self.divider = 60
+        elif time_unit == "hours":
+            self.divider = 60 * 60
+        elif time_unit == "days":
+            self.divider = 60 * 60 * 24
+        else:
+            raise ValueError(
+                f'Invalid time_unit: {time_unit}. Must be one of "seconds", "minutes", "hours", or "days".'
+            )
+
+        # Keep track of time spent evaluating
+        self.total_eval_wct = 0.0
+        self.step = -1
+
+    def on_train_batch_end(
+        self,
+        samples: int,  # total samples seen (per device)
+        train_elapsed: float,  # total training time (seconds)
+        world_size: int,
+        flops_per_batch: Optional[int] = None,  # (per device)
+        lengths: Optional[int] = None,  # total length of the samples seen (per device)
+    ) -> None:
+        self.step += 1
+        step = self.step
+        metrics = {}
+
+        self.history_samples.append(samples)
+        if lengths is not None:
+            self.history_lengths.append(lengths)
+            # if lengths are passed, there should be as many values as samples
+            assert len(self.history_samples) == len(self.history_lengths)
+        self.history_wct.append(train_elapsed)
+        if len(self.history_wct) == self.history_wct.maxlen:
+            elapsed_batches = len(self.history_samples) - 1
+            elapsed_samples = self.history_samples[-1] - self.history_samples[0]
+            elapsed_wct = self.history_wct[-1] - self.history_wct[0]
+            samples_per_sec = elapsed_samples * world_size / elapsed_wct
+            dev_samples_per_sec = elapsed_samples / elapsed_wct
+            metrics.update(
+                {
+                    "throughput/batches_per_sec": elapsed_batches * world_size / elapsed_wct,
+                    "throughput/samples_per_sec": samples_per_sec,
+                    "throughput/device/batches_per_sec": elapsed_batches / elapsed_wct,
+                    "throughput/device/samples_per_sec": dev_samples_per_sec,
+                }
+            )
+            if lengths is not None:
+                elapsed_lengths = int(self.history_lengths[-1]) - int(self.history_lengths[0])
+                avg_length = elapsed_lengths / elapsed_batches
+                metrics.update(
+                    {
+                        "throughput/tokens_per_sec": samples_per_sec * avg_length,
+                        "throughput/device/tokens_per_sec": dev_samples_per_sec * avg_length,
+                    }
+                )
+
+        if flops_per_batch is not None:
+            # sum of flops per batch across ranks
+            self.history_flops.append(flops_per_batch * world_size)
+        if len(self.history_flops) == self.history_flops.maxlen:
+            elapsed_flops = sum(self.history_flops) - self.history_flops[0]
+            elapsed_wct = self.history_wct[-1] - self.history_wct[0]
+            flops_per_sec = elapsed_flops / elapsed_wct
+            device_flops_per_sec = flops_per_sec / world_size
+            metrics.update(
+                {"throughput/flops_per_sec": flops_per_sec, "throughput/device/flops_per_sec": device_flops_per_sec}
+            )
+            if self.flops_available:
+                metrics["throughput/device/mfu"] = device_flops_per_sec / self.flops_available
+
+        metrics.update(
+            {
+                "time/train": train_elapsed / self.divider,
+                "time/val": self.total_eval_wct / self.divider,
+                "time/total": (train_elapsed + self.total_eval_wct) / self.divider,
+                "samples": samples,
+            }
+        )
+
+        self.log_dict(metrics, step)
+
+    def eval_end(self, eval_elapsed: float) -> None:
+        self.total_eval_wct += eval_elapsed  # seconds
+
+
+def plugin_to_compute_dtype(plugin: Precision) -> torch.dtype:
+    if isinstance(plugin, BitsandbytesPrecision):
+        return plugin.dtype
+    if isinstance(plugin, (HalfPrecision, MixedPrecision, HalfPrecisionPlugin)):
+        return plugin._desired_input_dtype
+    if isinstance(plugin, MixedPrecisionPlugin):
+        return torch.bfloat16 if plugin.precision == "bf16-mixed" else torch.half
+    if isinstance(plugin, (DoublePrecision, DoublePrecisionPlugin)):
+        return torch.double
+    if isinstance(plugin, (XLAPrecision, XLAPrecisionPlugin)):
+        return plugin._desired_dtype
+    if isinstance(plugin, TransformerEnginePrecision):
+        return torch.int8
+    if isinstance(plugin, (FSDPPrecision, FSDPPrecisionPlugin)):
+        return plugin.mixed_precision_config.reduce_dtype
+    if isinstance(plugin, Precision):
+        return torch.float32
+    raise NotImplementedError(plugin)
+
+
+class SpeedMonitorFabric(SpeedMonitorBase):
+    def __init__(self, fabric: Fabric, *args: Any, **kwargs: Any) -> None:
+        dtype = plugin_to_compute_dtype(fabric.strategy.precision)
+        flops_available = get_flops_available(fabric.device, dtype)
+        super().__init__(flops_available, fabric.log_dict, *args, **kwargs)
+
+    @fabric_rank_zero_only
+    def on_train_batch_end(self, *args: Any, **kwargs: Any) -> None:
+        super().on_train_batch_end(*args, **kwargs)
+
+
+class SpeedMonitorCallback(Callback):
+    def __init__(self, length_fn: Callable[[Any], int], batch_size: int, **kwargs: Any) -> None:
+        super().__init__()
+        self.speed_monitor: Optional[SpeedMonitorBase] = None
+        self.speed_monitor_kwargs = kwargs
+        self.length_fn = length_fn
+        self.batch_size = batch_size
+        self.eval_t0: int = 0
+        self.train_t0: int = 0
+        self.total_lengths: int = 0
+
+    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
+        if self.speed_monitor is not None:
+            return  # already setup
+        dtype = plugin_to_compute_dtype(trainer.precision_plugin)
+        flops_available = get_flops_available(trainer.strategy.root_device, dtype)
+        self.speed_monitor = SpeedMonitorBase(flops_available, trainer.logger.log_metrics, **self.speed_monitor_kwargs)
+
+    @trainer_rank_zero_only
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if trainer.fit_loop._should_accumulate():
+            return
+
+        self.train_t0 = time.perf_counter()
+
+    @trainer_rank_zero_only
+    def on_train_batch_end(
+        self, trainer: Trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int
+    ) -> None:
+        self.total_lengths += self.length_fn(batch)
+        if trainer.fit_loop._should_accumulate():
+            return
+        train_elapsed = time.perf_counter() - self.train_t0
+        assert self.speed_monitor is not None
+        iter_num = trainer.fit_loop.total_batch_idx
+        assert (measured_flops := pl_module.measured_flops) is not None
+        self.speed_monitor.on_train_batch_end(
+            (iter_num + 1) * self.batch_size,
+            train_elapsed,
+            # this assumes that device FLOPs are the same and that all devices have the same batch size
+            trainer.world_size,
+            flops_per_batch=measured_flops,
+            lengths=self.total_lengths,
+        )
+
+    @trainer_rank_zero_only
+    def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self.eval_t0 = time.perf_counter()
+
+    @trainer_rank_zero_only
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        eval_elapsed = time.perf_counter() - self.eval_t0
+        assert self.speed_monitor is not None
+        self.speed_monitor.eval_end(eval_elapsed)
+
+
+def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
+    flops_per_token = 2 * n_params  # each parameter is used for a MAC (2 FLOPS) per network operation
+    # this assumes that all samples have a fixed length equal to the block size
+    # which is most likely false during finetuning
+    flops_per_seq = flops_per_token * max_seq_length
+    attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2))
+    return flops_per_seq + attn_flops_per_seq
+
+
+def estimate_flops(model: GPT) -> int:
+    """Measures estimated FLOPs for MFU.
+
+    Refs:
+        * https://ar5iv.labs.arxiv.org/html/2205.05198#A1
+        * https://ar5iv.labs.arxiv.org/html/2204.02311#A2
+    """
+    # using all parameters for this is a naive over estimation because not all model parameters actually contribute to
+    # this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage
+    # (~10%) compared to the measured FLOPs, making those lower but more realistic.
+    # For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper.
+    n_trainable_params = num_parameters(model, requires_grad=True)
+    trainable_flops = flops_per_param(
+        model.max_seq_length, model.config.n_layer, model.config.n_embd, n_trainable_params
+    )
+    # forward + backward + gradients (assumes no gradient accumulation)
+    ops_per_step = 3 if model.training else 1
+    n_frozen_params = num_parameters(model, requires_grad=False)
+    frozen_flops = flops_per_param(model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params)
+    # forward + backward
+    frozen_ops_per_step = 2 if model.training else 1
+    return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops
+
+
+def measure_flops(model: GPT, x: torch.Tensor) -> int:
+    """Measures real FLOPs for HFU"""
+    flop_counter = FlopCounterMode(model, display=False)
+    ctx = nullcontext() if model.training else torch.no_grad()
+    with ctx, flop_counter:
+        y = model(x)
+        if model.training:
+            y.sum().backward()
+    return flop_counter.get_total_flops()

tsai_gpt/tokenizer.py

@@ -0,0 +1,103 @@
+import json
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+
+class Tokenizer:
+    def __init__(self, checkpoint_dir: Path) -> None:
+        self.use_bos = self.check_if_bos_token_used(checkpoint_dir)
+        self.bos_id = None
+        self.eos_id = None
+
+        # some checkpoints have both files, `.model` takes precedence
+        if (vocabulary_path := checkpoint_dir / "tokenizer.model").is_file():
+            from sentencepiece import SentencePieceProcessor
+
+            self.processor = SentencePieceProcessor(model_file=str(vocabulary_path))
+            self.backend = "sentencepiece"
+            self.bos_id = self.processor.bos_id()
+            self.eos_id = self.processor.eos_id()
+
+        elif (vocabulary_path := checkpoint_dir / "tokenizer.json").is_file():
+            from tokenizers import Tokenizer as HFTokenizer
+
+            self.processor = HFTokenizer.from_file(str(vocabulary_path))
+            self.backend = "huggingface"
+
+            if (special_tokens_path := checkpoint_dir / "tokenizer_config.json").is_file():
+                with open(special_tokens_path) as fp:
+                    config = json.load(fp)
+                bos_token = config.get("bos_token")
+                self.bos_id = self.token_to_id(bos_token) if bos_token is not None else None
+                eos_token = config.get("eos_token")
+                self.eos_id = self.token_to_id(eos_token) if eos_token is not None else None
+            if (special_tokens_path := checkpoint_dir / "generation_config.json").is_file():
+                with open(special_tokens_path) as fp:
+                    config = json.load(fp)
+                if self.bos_id is None:
+                    self.bos_id = config.get("bos_token_id")
+                if self.eos_id is None:
+                    self.eos_id = config.get("eos_token_id")
+        else:
+            raise NotImplementedError
+
+    @property
+    def vocab_size(self) -> int:
+        if self.backend == "huggingface":
+            return self.processor.get_vocab_size(with_added_tokens=False)
+        if self.backend == "sentencepiece":
+            return self.processor.vocab_size()
+        raise RuntimeError
+
+    def token_to_id(self, token: str) -> int:
+        if self.backend == "huggingface":
+            id_ = self.processor.token_to_id(token)
+        elif self.backend == "sentencepiece":
+            id_ = self.processor.piece_to_id(token)
+        else:
+            raise RuntimeError
+        if id_ is None:
+            raise ValueError(f"token {token!r} not found in the collection.")
+        return id_
+
+    def check_if_bos_token_used(self, checkpoint_dir: Path) -> bool:
+        if not (tokenizer_config_path := checkpoint_dir / "tokenizer_config.json").is_file():
+            return False
+        with open(tokenizer_config_path) as fp:
+            config = json.load(fp)
+        if any(config.get(check, False) for check in ("add_bos_token", "add_prefix_space")):
+            return True
+        # for examples that also use the Llama tokenizer, but do not have or set add_bos_token to True.
+        # ex: https://huggingface.co/stabilityai/StableBeluga2/blob/main/tokenizer_config.json#L2
+        return config.get("add_bos_token") is None and config.get("tokenizer_class") == "LlamaTokenizer"
+
+    def encode(
+        self,
+        string: str,
+        device: Optional[torch.device] = None,
+        bos: Optional[bool] = None,
+        eos: bool = False,
+        max_length: int = -1,
+    ) -> torch.Tensor:
+        if self.backend == "huggingface":
+            tokens = self.processor.encode(string).ids
+        elif self.backend == "sentencepiece":
+            tokens = self.processor.encode(string)
+        else:
+            raise RuntimeError
+        if bos or (bos is None and self.use_bos):
+            bos_id = self.bos_id
+            if bos_id is None:
+                raise NotImplementedError("This tokenizer does not have a defined a bos token")
+            tokens = [bos_id] + tokens
+        if eos:
+            tokens = tokens + [self.eos_id]
+        if max_length > 0:
+            tokens = tokens[:max_length]
+        return torch.tensor(tokens, dtype=torch.int, device=device)
+
+    def decode(self, tensor: torch.Tensor) -> str:
+        tokens = [tensor.item()] if tensor.ndim == 0 else tensor.tolist()
+        return self.processor.decode(tokens)

tsai_gpt/utils.py

@@ -0,0 +1,347 @@
+"""Utility functions for training and inference."""
+import math
+import pickle
+import sys
+from contextlib import nullcontext
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, ContextManager, Dict, List, Mapping, Optional, TypeVar, Union
+
+import lightning as L
+import torch
+import torch.nn as nn
+import torch.utils._device
+from lightning.fabric.strategies import FSDPStrategy
+from lightning.fabric.utilities.load import _lazy_load as lazy_load
+from torch.serialization import normalize_storage_type
+
+if TYPE_CHECKING:
+    from model import GPT
+
+
+def find_multiple(n: int, k: int) -> int:
+    assert k > 0
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def num_parameters(module: nn.Module, requires_grad: Optional[bool] = None) -> int:
+    total = 0
+    for p in module.parameters():
+        if requires_grad is None or p.requires_grad == requires_grad:
+            if hasattr(p, "quant_state"):
+                # bitsandbytes 4bit layer support
+                total += math.prod(p.quant_state[1])
+            else:
+                total += p.numel()
+    return total
+
+
+def gptq_quantization(enabled: bool = False) -> ContextManager:
+    if not enabled:
+        return nullcontext()
+
+    from lightning.fabric.plugins.precision.utils import _ClassReplacementContextManager
+
+    from quantize.gptq import ColBlockQuantizedLinear
+
+    class QuantizedLinear(ColBlockQuantizedLinear):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, bits=4, tile_cols=-1, **kwargs)
+
+    return _ClassReplacementContextManager({"torch.nn.Linear": QuantizedLinear})
+
+
+def check_valid_checkpoint_dir(checkpoint_dir: Path) -> None:
+    files = {
+        "lit_model.pth": (checkpoint_dir / "lit_model.pth").is_file(),
+        "lit_config.json": (checkpoint_dir / "lit_config.json").is_file(),
+        "tokenizer.json OR tokenizer.model": (checkpoint_dir / "tokenizer.json").is_file() or (
+            checkpoint_dir / "tokenizer.model"
+        ).is_file(),
+        "tokenizer_config.json": (checkpoint_dir / "tokenizer_config.json").is_file(),
+    }
+    if checkpoint_dir.is_dir():
+        if all(files.values()):
+            # we're good
+            return
+        problem = f" is missing the files: {[f for f, exists in files.items() if not exists]!r}"
+    else:
+        problem = " is not a checkpoint directory"
+
+    # list locally available checkpoints
+    available = list(Path("checkpoints").glob("*/*"))
+    if available:
+        options = "\n --checkpoint_dir ".join([""] + [repr(str(p.resolve())) for p in available])
+        extra = f"\nYou have downloaded locally:{options}\n"
+    else:
+        extra = ""
+
+    error_message = (
+        f"--checkpoint_dir {str(checkpoint_dir.absolute())!r}{problem}."
+        "\nFind download instructions at https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials\n"
+        f"{extra}\nSee all download options by running:\n python scripts/download.py"
+    )
+    print(error_message, file=sys.stderr)
+    raise SystemExit(1)
+
+
+class SavingProxyForStorage:
+    def __init__(self, obj, saver, protocol_version=5):
+        self.protocol_version = protocol_version
+        self.saver = saver
+        if not (isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj)):
+            raise TypeError(f"expected storage, not {type(obj)}")
+
+        # this logic is taken from PyTorch 2.0+ torch/serialization.py
+        if isinstance(obj, torch.storage.TypedStorage):
+            # PT upstream wants to deprecate this eventually...
+            storage = obj._untyped_storage
+            storage_type_str = obj._pickle_storage_type()
+            storage_type = getattr(torch, storage_type_str)
+            storage_numel = obj._size()
+        else:
+            storage = obj
+            storage_type = normalize_storage_type(type(obj))
+            storage_numel = storage.nbytes()
+
+        storage_key = saver._write_storage_and_return_key(storage)
+        location = torch.serialization.location_tag(storage)
+
+        self.storage_info = ("storage", storage_type, storage_key, location, storage_numel)
+
+    def __reduce_ex__(self, protocol_version):
+        assert False, "this should be handled with out of band"
+
+
+class SavingProxyForTensor:
+    def __init__(self, tensor, saver, protocol_version=5):
+        self.protocol_version = protocol_version
+        self.reduce_ret_fn, reduce_args = tensor.__reduce_ex__(protocol_version)
+        if reduce_args[0] == torch._utils._rebuild_tensor_v2:
+            # for Tensors with Python attributes
+            (a0, a1, (storage, *a2_other), *other_reduce_args) = reduce_args
+            assert isinstance(storage, torch.storage.TypedStorage), "Please check for updates"
+            storage_proxy = SavingProxyForStorage(storage, saver, protocol_version=protocol_version)
+            self.reduce_args = (a0, a1, (storage_proxy, *a2_other), *other_reduce_args)
+        else:
+            (storage, *other_reduce_args) = reduce_args
+            assert isinstance(storage, torch.storage.TypedStorage), "Please check for updates"
+            storage_proxy = SavingProxyForStorage(storage, saver, protocol_version=protocol_version)
+            self.reduce_args = (storage_proxy, *other_reduce_args)
+
+    def __reduce_ex__(self, protocol_version):
+        if protocol_version != self.protocol_version:
+            raise RuntimeError(f"Unexpected protocol version: expected {self.protocol_version}, got {protocol_version}")
+        return self.reduce_ret_fn, self.reduce_args
+
+
+class IncrementalPyTorchPickler(pickle.Pickler):
+    def __init__(self, saver, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.storage_dtypes = {}
+        self.saver = saver
+        self.id_map = {}
+
+    # this logic is taken from PyTorch 2.0+ torch/serialization.py
+    def persistent_id(self, obj):
+        # FIXME: the docs say that persistent_id should only return a string
+        # but torch store returns tuples. This works only in the binary protocol
+        # see
+        # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects
+        # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537
+        if isinstance(obj, SavingProxyForStorage):
+            return obj.storage_info
+
+        if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj):
+            if isinstance(obj, torch.storage.TypedStorage):
+                # TODO: Once we decide to break serialization FC, this case
+                # can be deleted
+                storage = obj._untyped_storage
+                storage_dtype = obj.dtype
+                storage_type_str = obj._pickle_storage_type()
+                storage_type = getattr(torch, storage_type_str)
+                storage_numel = obj._size()
+
+            else:
+                storage = obj
+                storage_dtype = torch.uint8
+                storage_type = normalize_storage_type(type(obj))
+                storage_numel = storage.nbytes()
+
+            # If storage is allocated, ensure that any other saved storages
+            # pointing to the same data all have the same dtype. If storage is
+            # not allocated, don't perform this check
+            if storage.data_ptr() != 0:
+                if storage.data_ptr() in self.storage_dtypes:
+                    if storage_dtype != self.storage_dtypes[storage.data_ptr()]:
+                        raise RuntimeError(
+                            "Cannot save multiple tensors or storages that view the same data as different types"
+                        )
+                else:
+                    self.storage_dtypes[storage.data_ptr()] = storage_dtype
+
+            storage_key = self.id_map.get(storage._cdata)
+            if storage_key is None:
+                storage_key = self.saver._write_storage_and_return_key(storage)
+                self.id_map[storage._cdata] = storage_key
+            location = torch.serialization.location_tag(storage)
+
+            return ("storage", storage_type, storage_key, location, storage_numel)
+
+        return None
+
+
+class incremental_save:
+    def __init__(self, name):
+        self.name = name
+        self.zipfile = torch._C.PyTorchFileWriter(str(name))
+        self.has_saved = False
+        self.next_key = 0
+
+    def __enter__(self):
+        return self
+
+    def store_early(self, tensor):
+        if isinstance(tensor, torch.Tensor):
+            return SavingProxyForTensor(tensor, self)
+        raise TypeError(f"can only store tensors early, not {type(tensor)}")
+
+    def save(self, obj):
+        if self.has_saved:
+            raise RuntimeError("have already saved")
+        # Write the pickle data for `obj`
+        data_buf = BytesIO()
+        pickler = IncrementalPyTorchPickler(self, data_buf, protocol=5)
+        pickler.dump(obj)
+        data_value = data_buf.getvalue()
+        self.zipfile.write_record("data.pkl", data_value, len(data_value))
+        self.has_saved = True
+
+    def _write_storage_and_return_key(self, storage):
+        if self.has_saved:
+            raise RuntimeError("have already saved")
+        key = self.next_key
+        self.next_key += 1
+        name = f"data/{key}"
+        if storage.device.type != "cpu":
+            storage = storage.cpu()
+        num_bytes = storage.nbytes()
+        self.zipfile.write_record(name, storage.data_ptr(), num_bytes)
+        return key
+
+    def __exit__(self, type, value, traceback):
+        self.zipfile.write_end_of_file()
+
+
+T = TypeVar("T")
+
+
+def chunked_cross_entropy(
+    logits: Union[torch.Tensor, List[torch.Tensor]], targets: torch.Tensor, chunk_size: int = 128
+) -> torch.Tensor:
+    # with large max_sequence_lengths, the beginning of `backward` allocates a large memory chunk which can dominate
+    # the memory usage in fine-tuning settings with low number of parameters.
+    # as a workaround hack, the cross entropy computation is chunked to force it to deallocate on the go, reducing
+    # the memory spike's magnitude
+
+    # lm_head was chunked (we are fine-tuning)
+    if isinstance(logits, list):
+        # don't want to chunk cross entropy
+        if chunk_size == 0:
+            logits = torch.cat(logits, dim=1)
+            logits = logits.reshape(-1, logits.size(-1))
+            targets = targets.reshape(-1)
+            return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1)
+
+        # chunk cross entropy
+        logit_chunks = [logit_chunk.reshape(-1, logit_chunk.size(-1)) for logit_chunk in logits]
+        target_chunks = [target_chunk.reshape(-1) for target_chunk in targets.split(logits[0].size(1), dim=1)]
+        loss_chunks = [
+            torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none")
+            for logit_chunk, target_chunk in zip(logit_chunks, target_chunks)
+        ]
+        return torch.cat(loss_chunks).mean()
+
+    # no chunking at all
+    logits = logits.reshape(-1, logits.size(-1))
+    targets = targets.reshape(-1)
+    if chunk_size == 0:
+        return torch.nn.functional.cross_entropy(logits, targets, ignore_index=-1)
+
+    # lm_head wasn't chunked, chunk cross entropy
+    logit_chunks = logits.split(chunk_size)
+    target_chunks = targets.split(chunk_size)
+    loss_chunks = [
+        torch.nn.functional.cross_entropy(logit_chunk, target_chunk, ignore_index=-1, reduction="none")
+        for logit_chunk, target_chunk in zip(logit_chunks, target_chunks)
+    ]
+    return torch.cat(loss_chunks).mean()
+
+
+def map_old_state_dict_weights(state_dict: Dict, mapping: Mapping, prefix: str) -> Dict:
+    for checkpoint_name, attribute_name in mapping.items():
+        full_checkpoint_name = prefix + checkpoint_name
+        if full_checkpoint_name in state_dict:
+            full_attribute_name = prefix + attribute_name
+            state_dict[full_attribute_name] = state_dict.pop(full_checkpoint_name)
+    return state_dict
+
+
+def get_default_supported_precision(training: bool) -> str:
+    """Return default precision that is supported by the hardware: either `bf16` or `16`.
+
+    Args:
+        training: `-mixed` or `-true` version of the precision to use
+
+    Returns:
+        default precision that is suitable for the task and is supported by the hardware
+    """
+    from lightning.fabric.accelerators import MPSAccelerator
+
+    if MPSAccelerator.is_available() or (torch.cuda.is_available() and not torch.cuda.is_bf16_supported()):
+        return "16-mixed" if training else "16-true"
+    return "bf16-mixed" if training else "bf16-true"
+
+
+def load_checkpoint(fabric: L.Fabric, model: nn.Module, checkpoint_path: Path, strict: bool = True) -> None:
+    if isinstance(fabric.strategy, FSDPStrategy):
+        fabric.load_raw(checkpoint_path, model, strict=strict)
+    else:
+        state_dict = lazy_load(checkpoint_path)
+        state_dict = state_dict.get("model", state_dict)
+        model.load_state_dict(state_dict, strict=strict)
+
+
+def flops_per_param(max_seq_length: int, n_layer: int, n_embd: int, n_params: int) -> int:
+    flops_per_token = 2 * n_params  # each parameter is used for a MAC (2 FLOPS) per network operation
+    # this assumes that all samples have a fixed length equal to the block size
+    # which is most likely false during finetuning
+    flops_per_seq = flops_per_token * max_seq_length
+    attn_flops_per_seq = n_layer * 2 * 2 * (n_embd * (max_seq_length**2))
+    return flops_per_seq + attn_flops_per_seq
+
+
+def estimate_flops(model: "GPT", training: bool) -> int:
+    """Measures estimated FLOPs for MFU.
+
+    Refs:
+        * https://ar5iv.labs.arxiv.org/html/2205.05198#A1
+        * https://ar5iv.labs.arxiv.org/html/2204.02311#A2
+    """
+    # using all parameters for this is a naive over estimation because not all model parameters actually contribute to
+    # this FLOP computation (e.g. embedding, norm). For this reason, the result will be higher by a fixed percentage
+    # (~10%) compared to the measured FLOPs, making those lower but more realistic.
+    # For a proper estimate, this needs a more fine-grained calculation as in Appendix A of the paper.
+    n_trainable_params = num_parameters(model, requires_grad=True)
+    trainable_flops = flops_per_param(
+        model.max_seq_length, model.config.n_layer, model.config.n_embd, n_trainable_params
+    )
+    # forward + backward + gradients (assumes no gradient accumulation)
+    ops_per_step = 3 if training else 1
+    n_frozen_params = num_parameters(model, requires_grad=False)
+    frozen_flops = flops_per_param(model.max_seq_length, model.config.n_layer, model.config.n_embd, n_frozen_params)
+    # forward + backward
+    frozen_ops_per_step = 2 if training else 1
+    return ops_per_step * trainable_flops + frozen_ops_per_step * frozen_flops