main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pip in /opt/conda/lib/python3.10/site-packages (23.3.1)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mLooking in indexes: https://download.pytorch.org/whl/cu118\n",
      "Requirement already satisfied: torch in /opt/conda/lib/python3.10/site-packages (2.1.1+cu118)\n",
      "Requirement already satisfied: torchaudio in /opt/conda/lib/python3.10/site-packages (2.1.1+cu118)\n",
      "Requirement already satisfied: torchvision in /opt/conda/lib/python3.10/site-packages (0.16.1+cu118)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch) (3.12.2)\n",
      "Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.10/site-packages (from torch) (4.7.1)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch) (3.1)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from torch) (2023.6.0)\n",
      "Requirement already satisfied: triton==2.1.0 in /opt/conda/lib/python3.10/site-packages (from torch) (2.1.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from torchvision) (1.24.4)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from torchvision) (2.31.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.10/site-packages (from torchvision) (10.0.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch) (2.1.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision) (3.1.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision) (2023.7.22)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch) (1.3.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395 (from -r requirements.txt (line 1))\n",
      "  Using cached lightning-2.2.0.dev0-py3-none-any.whl\n",
      "Requirement already satisfied: huggingface_hub in /opt/conda/lib/python3.10/site-packages (0.19.4)\n",
      "Requirement already satisfied: tokenizers in /opt/conda/lib/python3.10/site-packages (0.15.0)\n",
      "Requirement already satisfied: sentencepiece in /opt/conda/lib/python3.10/site-packages (0.1.99)\n",
      "Requirement already satisfied: jsonargparse[signatures] in /opt/conda/lib/python3.10/site-packages (from -r requirements.txt (line 2)) (4.27.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (3.12.2)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (2023.6.0)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (4.65.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (6.0.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (4.7.1)\n",
      "Requirement already satisfied: packaging>=20.9 in /opt/conda/lib/python3.10/site-packages (from huggingface_hub) (23.1)\n",
      "Requirement already satisfied: lightning-utilities<2.0,>=0.8.0 in /opt/conda/lib/python3.10/site-packages (from lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (0.10.0)\n",
      "Requirement already satisfied: numpy<3.0,>=1.17.2 in /opt/conda/lib/python3.10/site-packages (from lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.24.4)\n",
      "Requirement already satisfied: torch<4.0,>=1.12.0 in /opt/conda/lib/python3.10/site-packages (from lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (2.1.1+cu118)\n",
      "Requirement already satisfied: torchmetrics<3.0,>=0.7.0 in /opt/conda/lib/python3.10/site-packages (from lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.2.0)\n",
      "Requirement already satisfied: pytorch-lightning in /opt/conda/lib/python3.10/site-packages (from lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (2.1.2)\n",
      "Requirement already satisfied: docstring-parser>=0.15 in /opt/conda/lib/python3.10/site-packages (from jsonargparse[signatures]->-r requirements.txt (line 2)) (0.15)\n",
      "Requirement already satisfied: typeshed-client>=2.1.0 in /opt/conda/lib/python3.10/site-packages (from jsonargparse[signatures]->-r requirements.txt (line 2)) (2.4.0)\n",
      "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /opt/conda/lib/python3.10/site-packages (from fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (3.8.6)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.10/site-packages (from lightning-utilities<2.0,>=0.8.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (65.6.3)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (3.1)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (3.1.2)\n",
      "Requirement already satisfied: triton==2.1.0 in /opt/conda/lib/python3.10/site-packages (from torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (2.1.0)\n",
      "Requirement already satisfied: importlib-resources>=1.4.0 in /opt/conda/lib/python3.10/site-packages (from typeshed-client>=2.1.0->jsonargparse[signatures]->-r requirements.txt (line 2)) (6.1.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface_hub) (3.1.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface_hub) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface_hub) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->huggingface_hub) (2023.7.22)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (4.0.3)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.4.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.10/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2025.0,>2021.06.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.3.1)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch<4.0,>=1.12.0->lightning@ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395->-r requirements.txt (line 1)) (1.3.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install --upgrade pip\n",
    "!pip install torch torchaudio torchvision --upgrade --index-url https://download.pytorch.org/whl/cu118\n",
    "!pip install huggingface_hub tokenizers sentencepiece -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = \"pythia-160m\"\n",
    "name = \"redpajama\"\n",
    "out_dir = Path(\"out\") / name\n",
    "save_interval = 100\n",
    "eval_interval = 1000\n",
    "eval_iters = 100\n",
    "log_interval = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Hyperparameters\n",
    "learning_rate = 6e-3\n",
    "batch_size = 32\n",
    "micro_batch_size = 4\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 = 25000\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 = 6000\n",
    "lr_decay_iters = max_iters\n",
    "min_lr = 6e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "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": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_copy = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "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": 11,
   "metadata": {
    "tags": []
   },
   "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\"]+1) % 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)\n",
    "            \n",
    "        '''if loss.item() <= 4.0 and state['iter_num'] >= 2000:\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",
    "            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'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "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": 13,
   "metadata": {
    "tags": []
   },
   "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": 14,
   "metadata": {
    "tags": []
   },
   "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": 15,
   "metadata": {
    "tags": []
   },
   "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": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16-bit Automatic Mixed Precision (AMP)\n",
      "Seed set to 1337\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'model_name': 'pythia-160m', 'name': 'redpajama', 'save_interval': 100, 'eval_interval': 1000, 'eval_iters': 100, 'log_interval': 100, 'learning_rate': 0.006, 'batch_size': 32, 'micro_batch_size': 4, 'gradient_accumulation_steps': 8, 'max_iters': 25000, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'decay_lr': True, 'warmup_iters': 6000, 'lr_decay_iters': 25000, 'min_lr': 0.0006}\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: 1.85 seconds.\n",
      "Total parameters 162,322,944\n",
      "Resuming training from out/redpajama/iter-025000-ckpt.pth\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[16], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m torch\u001b[38;5;241m.\u001b[39mset_float32_matmul_precision(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmedium\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43msetup\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdevices\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrain_data_dir\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mPath\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mdata/lit-redpajama-sample\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresume\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\n\u001b[1;32m      6\u001b[0m \u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[8], line 27\u001b[0m, in \u001b[0;36msetup\u001b[0;34m(devices, train_data_dir, val_data_dir, precision, resume)\u001b[0m\n\u001b[1;32m     25\u001b[0m fabric \u001b[38;5;241m=\u001b[39m L\u001b[38;5;241m.\u001b[39mFabric(devices\u001b[38;5;241m=\u001b[39mdevices, strategy\u001b[38;5;241m=\u001b[39mstrategy, precision\u001b[38;5;241m=\u001b[39mprecision, loggers\u001b[38;5;241m=\u001b[39mlogger)\n\u001b[1;32m     26\u001b[0m fabric\u001b[38;5;241m.\u001b[39mprint(hparams)\n\u001b[0;32m---> 27\u001b[0m \u001b[43mfabric\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_data_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_data_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresume\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/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[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstrategy\u001b[38;5;241m.\u001b[39mlauncher, (_MultiProcessingLauncher, _XLALauncher)):\n\u001b[1;32m    830\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m    831\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTo use the `\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstrategy)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m` strategy, `.launch()` needs to be called with a function\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    832\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m that contains the code to launch in processes.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    833\u001b[0m     )\n\u001b[0;32m--> 834\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_wrap_and_launch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunction\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/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[38;5;28;01mif\u001b[39;00m (launcher \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_strategy\u001b[38;5;241m.\u001b[39mlauncher) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    919\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m launcher\u001b[38;5;241m.\u001b[39mlaunch(to_run, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 920\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mto_run\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/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[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_strategy\u001b[38;5;241m.\u001b[39msetup_environment()\n\u001b[1;32m    924\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m _replace_dunder_methods(DataLoader, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdataset\u001b[39m\u001b[38;5;124m\"\u001b[39m), _replace_dunder_methods(BatchSampler):\n\u001b[0;32m--> 925\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mto_run\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[10], line 68\u001b[0m, in \u001b[0;36mmain\u001b[0;34m(fabric, train_data_dir, val_data_dir, resume)\u001b[0m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m resume:\n\u001b[1;32m     67\u001b[0m     fabric\u001b[38;5;241m.\u001b[39mprint(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mResuming training from \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresume\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 68\u001b[0m     \u001b[43mfabric\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresume\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstate\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     70\u001b[0m train_time \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mperf_counter()\n\u001b[1;32m     71\u001b[0m train(fabric, state, train_dataloader, val_dataloader, speed_monitor)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/lightning/fabric/fabric.py:764\u001b[0m, in \u001b[0;36mFabric.load\u001b[0;34m(self, path, state, strict)\u001b[0m\n\u001b[1;32m    747\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Load a checkpoint from a file and restore the state of objects (modules, optimizers, etc.)\u001b[39;00m\n\u001b[1;32m    748\u001b[0m \n\u001b[1;32m    749\u001b[0m \u001b[38;5;124;03mHow and which processes load gets determined by the `strategy`.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    761\u001b[0m \n\u001b[1;32m    762\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    763\u001b[0m unwrapped_state \u001b[38;5;241m=\u001b[39m _unwrap_objects(state)\n\u001b[0;32m--> 764\u001b[0m remainder \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_strategy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstate\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43munwrapped_state\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstrict\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstrict\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    765\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbarrier()\n\u001b[1;32m    766\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m state \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    767\u001b[0m     \u001b[38;5;66;03m# We need to unwrap objects (see above) but this creates a new dictionary. In-place updates\u001b[39;00m\n\u001b[1;32m    768\u001b[0m     \u001b[38;5;66;03m# (for user metadata) wouldn't show up in the original dict, so we need to copy the data back.\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/lightning/fabric/strategies/strategy.py:333\u001b[0m, in \u001b[0;36mStrategy.load_checkpoint\u001b[0;34m(self, path, state, strict)\u001b[0m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Load the contents from a checkpoint and restore the state of the given objects.\u001b[39;00m\n\u001b[1;32m    315\u001b[0m \n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03mArgs:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    330\u001b[0m \n\u001b[1;32m    331\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    332\u001b[0m torch\u001b[38;5;241m.\u001b[39mcuda\u001b[38;5;241m.\u001b[39mempty_cache()\n\u001b[0;32m--> 333\u001b[0m checkpoint \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckpoint_io\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    334\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m state:\n\u001b[1;32m    335\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m checkpoint\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/lightning/fabric/plugins/io/torch_io.py:79\u001b[0m, in \u001b[0;36mTorchCheckpointIO.load_checkpoint\u001b[0;34m(self, path, map_location)\u001b[0m\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m fs\u001b[38;5;241m.\u001b[39mexists(path):\n\u001b[1;32m     77\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mFileNotFoundError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCheckpoint file not found: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpath\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 79\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpl_load\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmap_location\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmap_location\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/lightning/fabric/utilities/cloud_io.py:55\u001b[0m, in \u001b[0;36m_load\u001b[0;34m(path_or_url, map_location)\u001b[0m\n\u001b[1;32m     53\u001b[0m fs \u001b[38;5;241m=\u001b[39m get_filesystem(path_or_url)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m fs\u001b[38;5;241m.\u001b[39mopen(path_or_url, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrb\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[0;32m---> 55\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmap_location\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmap_location\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/serialization.py:1014\u001b[0m, in \u001b[0;36mload\u001b[0;34m(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)\u001b[0m\n\u001b[1;32m   1012\u001b[0m             \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1013\u001b[0m                 \u001b[38;5;28;01mraise\u001b[39;00m pickle\u001b[38;5;241m.\u001b[39mUnpicklingError(UNSAFE_MESSAGE \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(e)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 1014\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_load\u001b[49m\u001b[43m(\u001b[49m\u001b[43mopened_zipfile\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1015\u001b[0m \u001b[43m                     \u001b[49m\u001b[43mmap_location\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1016\u001b[0m \u001b[43m                     \u001b[49m\u001b[43mpickle_module\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1017\u001b[0m \u001b[43m                     \u001b[49m\u001b[43moverall_storage\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moverall_storage\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1018\u001b[0m \u001b[43m                     \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mpickle_load_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1019\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m mmap:\n\u001b[1;32m   1020\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmmap can only be used with files saved with \u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   1021\u001b[0m                        \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`torch.save(_use_new_zipfile_serialization=True), \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1022\u001b[0m                        \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mplease torch.save your checkpoint with this option in order to use mmap.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/serialization.py:1422\u001b[0m, in \u001b[0;36m_load\u001b[0;34m(zip_file, map_location, pickle_module, pickle_file, overall_storage, **pickle_load_args)\u001b[0m\n\u001b[1;32m   1420\u001b[0m unpickler \u001b[38;5;241m=\u001b[39m UnpicklerWrapper(data_file, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mpickle_load_args)\n\u001b[1;32m   1421\u001b[0m unpickler\u001b[38;5;241m.\u001b[39mpersistent_load \u001b[38;5;241m=\u001b[39m persistent_load\n\u001b[0;32m-> 1422\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43munpickler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1424\u001b[0m torch\u001b[38;5;241m.\u001b[39m_utils\u001b[38;5;241m.\u001b[39m_validate_loaded_sparse_tensors()\n\u001b[1;32m   1425\u001b[0m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39m_log_api_usage_metadata(\n\u001b[1;32m   1426\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch.load.metadata\u001b[39m\u001b[38;5;124m\"\u001b[39m, {\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mserialization_id\u001b[39m\u001b[38;5;124m\"\u001b[39m: zip_file\u001b[38;5;241m.\u001b[39mserialization_id()}\n\u001b[1;32m   1427\u001b[0m )\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/serialization.py:1392\u001b[0m, in \u001b[0;36m_load.<locals>.persistent_load\u001b[0;34m(saved_id)\u001b[0m\n\u001b[1;32m   1390\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1391\u001b[0m     nbytes \u001b[38;5;241m=\u001b[39m numel \u001b[38;5;241m*\u001b[39m torch\u001b[38;5;241m.\u001b[39m_utils\u001b[38;5;241m.\u001b[39m_element_size(dtype)\n\u001b[0;32m-> 1392\u001b[0m     typed_storage \u001b[38;5;241m=\u001b[39m \u001b[43mload_tensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnbytes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_maybe_decode_ascii\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlocation\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1394\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m typed_storage\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/serialization.py:1357\u001b[0m, in \u001b[0;36m_load.<locals>.load_tensor\u001b[0;34m(dtype, numel, key, location)\u001b[0m\n\u001b[1;32m   1355\u001b[0m     storage \u001b[38;5;241m=\u001b[39m overall_storage[storage_offset:storage_offset \u001b[38;5;241m+\u001b[39m numel]\n\u001b[1;32m   1356\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1357\u001b[0m     storage \u001b[38;5;241m=\u001b[39m \u001b[43mzip_file\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_storage_from_record\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnumel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mUntypedStorage\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39m_typed_storage()\u001b[38;5;241m.\u001b[39m_untyped_storage\n\u001b[1;32m   1358\u001b[0m \u001b[38;5;66;03m# swap here if byteswapping is needed\u001b[39;00m\n\u001b[1;32m   1359\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m byteorderdata \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ],
     "output_type": "error"
    }
   ],
   "source": [
    "torch.set_float32_matmul_precision(\"medium\")\n",
    "setup(\n",
    "    devices=1,\n",
    "    train_data_dir=Path(\"data/lit-redpajama-sample\"),\n",
    "    resume=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "generation_config.json\tmeta-llama\ttokenizer.model\n",
      "lit_config.json\t\ttokenizer.json\ttokenizer_config.json\n"
     ]
    }
   ],
   "source": [
    "!ls checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from out/redpajama/iter-025000-ckpt.pth\n",
      "The main reason for the financial 16 taxpayer tax increase, the number of individuals working in the first quarter of 2018 to take interest rates and it is very likely that an employee might receive that amount and the amount paid back income of his assets in the third quarter of 2018.\n",
      "It is not easy to give credit to the first quarter of 2018 to give the same amount that they can sell, as the amount paid back, and the amount paid back income of the shares of the ear pay will not be paid until the retirement is in 2019.\n",
      "The more than the more the amount paid back in the quarter of 2018 to be paying a year and a pay back income of the year the tax return of the dividend will be paid back. RSS 600 605 50th Anniversary\n",
      "The G200 605 603 \n",
      "--------------\n",
      "\n",
      "Covid19 pandemic gave the world new 1965 results to the National Public Health Protection Agency (SCO), the largest provider of scientific medical care program, the latest one of the scientific community news agency, and the National Public Health Protection Agency (SCO), is an investment cap.\n",
      "The National Public Health Protection Agency (SCO), one of the largest provider of scientific medical care available for the scientific community.\n",
      "The National Public Health Protection Agency (SCO) is a leading provider of clinical research from all medical staff and doctoral research to the National Public Health Protection Agency (SCO), a leading provider of medical care facilities and hospitals, a new patient care center for patients, and a medical care center for approximately 10,000 patients.\n",
      "The National Public Health Protection Agency (SCO) is a state-wide program providing evidence for Medicare patients to treat patients with chronic pain and death risk for patients.\n",
      "We hope to work\n",
      "--------------\n",
      "\n",
      "Biofuels can be used 100% of electric vehicles (KA13669) or 0.5% of electric vehicles. The power to maintain electric vehicles is always a great source of electric vehicles. It can also be used to install vehicle charging via vehicle charging. It also is used to produce a vehicle charging cable that can be used to operate it. It can be used on vehicle charging and charging.\n",
      "The vehicles can be used to protect the fuel and fuel consumption and are used to produce a fuel fuel efficiency.\n",
      "The device is used to protect the vehicle from the vehicle from the vehicle from the vehicle without the car from the vehicle from the vehicle from the vehicle at the vehicle. It also can also be used to store the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from being the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle from the vehicle\n",
      "--------------\n",
      "\n",
      "You believe it or not but the fact is that it is an important place for people that are not alone. This is a simple way to deal with people who are the only ones in the country that can be very different from the internet.\n",
      "Anywhere that can be created is that you have to be successful.\n",
      "A lot of people have to have to be able to put people on the Internet, people, as you are now having to use them. If we are working at the time of the year then, we will have to do that as soon as you have to be ready to work on a site. I am no longer looking to know how to improve the work of a site or an Internet site on the site.\n",
      "If you are planning a project or not, you will be ready to work with a site and would make it work as the site.\n",
      "How can I do that? It is a huge challenge for people who want to be in and that have been working the way they have to share in the\n",
      "--------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "from tsai_gpt.tokenizer import Tokenizer\n",
    "precision = get_default_supported_precision(False)\n",
    "logger = CSVLogger(\"out\", name, flush_logs_every_n_steps=log_interval)\n",
    "fabric = L.Fabric(devices=1, strategy=\"auto\", precision=precision, loggers=logger)\n",
    "\n",
    "config = Config.from_name(model_name)\n",
    "\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",
    "checkpoint_path = Path(\"out/redpajama/iter-025000-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",
    "weight_decay = 1e-1\n",
    "beta1 = 0.9\n",
    "beta2 = 0.95\n",
    "learning_rate = 6e-3\n",
    "hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith(\"_\")}\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",
    "resume = max(out_dir.glob(\"*.pth\"), key=lambda p: int(p.name.split(\"-\")[1]))\n",
    "if resume:\n",
    "    fabric.print(f\"Loading model from {resume}\")\n",
    "    fabric.load(resume, state)\n",
    "\n",
    "deviceType = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "m  = model.to(deviceType)\n",
    "tokenizer_gpt = Tokenizer(checkpoint_dir=Path(\"checkpoints/meta-llama/Llama-2-7b-chat-hf\")) \n",
    "    \n",
    "def generate_predictions(prompt, max_new_tokens=200, temperature=0.8, top_k=50):\n",
    "    m.eval()\n",
    "    encoded_text = tokenizer_gpt.encode(prompt)\n",
    "    #print('--------------------encoded text = ',encoded_text)\n",
    "    \n",
    "    reshaped_tensor = torch.unsqueeze(encoded_text, 0).to(deviceType)     \n",
    "    #print('--------------------reshaped_tensor = ',reshaped_tensor)\n",
    "    out_text = tokenizer_gpt.decode(m.generate(reshaped_tensor, max_new_tokens=max_new_tokens, temperature=0.8, top_k=50)[0])\n",
    "    m.train()\n",
    "    return out_text\n",
    "\n",
    "\n",
    "print(\n",
    "    generate_predictions(\n",
    "        \"The main reason for the financial \"\n",
    "    )\n",
    ")\n",
    "print(\"--------------\\n\")\n",
    "\n",
    "print(\n",
    "    generate_predictions(\n",
    "        \"Covid19 pandemic gave the world new \"\n",
    "    )\n",
    ")\n",
    "print(\"--------------\\n\")\n",
    "\n",
    "print(\n",
    "    generate_predictions(\n",
    "        \"Biofuels can be used \"\n",
    "    )\n",
    ")\n",
    "print(\"--------------\\n\")\n",
    "\n",
    "print(\n",
    "    generate_predictions(\n",
    "        \"You believe it or not but the fact is\"\n",
    "    )\n",
    ")\n",
    "print(\"--------------\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:2048\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0.0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4.0,
    "name": "ml.t3.medium",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 1.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8.0,
    "name": "ml.t3.large",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 2.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 3.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 4.0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8.0,
    "name": "ml.m5.large",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 5.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 6.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 7.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64.0,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 8.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128.0,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 9.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192.0,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 10.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256.0,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 11.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384.0,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 12.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8.0,
    "name": "ml.m5d.large",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 13.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 14.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 15.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64.0,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 16.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128.0,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 17.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192.0,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 18.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256.0,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 19.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384.0,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 20.0,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0.0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0.0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0.0
   },
   {
    "_defaultOrder": 21.0,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4.0,
    "name": "ml.c5.large",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 22.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8.0,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 23.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 24.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 25.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72.0,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36.0
   },
   {
    "_defaultOrder": 26.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96.0,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 27.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144.0,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72.0
   },
   {
    "_defaultOrder": 28.0,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192.0,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 29.0,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 30.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 31.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64.0,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 32.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128.0,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 33.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192.0,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 34.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256.0,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 35.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 61.0,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 36.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 244.0,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 37.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 488.0,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 38.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768.0,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 39.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.r5.large",
    "vcpuNum": 2.0
   },
   {
    "_defaultOrder": 40.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 41.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64.0,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 42.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128.0,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 43.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256.0,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 44.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384.0,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 45.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512.0,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 46.0,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768.0,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 47.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16.0,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4.0
   },
   {
    "_defaultOrder": 48.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32.0,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8.0
   },
   {
    "_defaultOrder": 49.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64.0,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16.0
   },
   {
    "_defaultOrder": 50.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128.0,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32.0
   },
   {
    "_defaultOrder": 51.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256.0,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64.0
   },
   {
    "_defaultOrder": 52.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192.0,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48.0
   },
   {
    "_defaultOrder": 53.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384.0,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 54.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768.0,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192.0
   },
   {
    "_defaultOrder": 55.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152.0,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96.0
   },
   {
    "_defaultOrder": 56.0,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8.0,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152.0,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96.0
   }
  ],
  "instance_type": "ml.g4dn.2xlarge",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 2.0.1 Python 3.10 GPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:ap-south-1:394103062818:image/pytorch-2.0.1-gpu-py310"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}