Untitled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3c1eb100",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /home/user1/miniconda3/envs/Prateek/lib/libtinfo.so.6: no version information available (required by /bin/bash)\n",
      "--2023-02-28 16:24:57--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.109.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 1115394 (1.1M) [text/plain]\n",
      "Saving to: ‘input.txt’\n",
      "\n",
      "input.txt           100%[===================>]   1.06M  --.-KB/s    in 0.1s    \n",
      "\n",
      "2023-02-28 16:24:58 (7.29 MB/s) - ‘input.txt’ saved [1115394/1115394]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# We always start with a dataset to train on. Let's download the tiny shakespeare dataset\n",
    "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4ca3b4c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read it in to inspect it\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b43cd226",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters:  1115394\n"
     ]
    }
   ],
   "source": [
    "print(\"length of dataset in characters: \", len(text))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ef9b3522",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You are all resolved rather to die than to famish?\n",
      "\n",
      "All:\n",
      "Resolved. resolved.\n",
      "\n",
      "First Citizen:\n",
      "First, you know Caius Marcius is chief enemy to the people.\n",
      "\n",
      "All:\n",
      "We know't, we know't.\n",
      "\n",
      "First Citizen:\n",
      "Let us kill him, and we'll have corn at our own price.\n",
      "Is't a verdict?\n",
      "\n",
      "All:\n",
      "No more talking on't; let it be done: away, away!\n",
      "\n",
      "Second Citizen:\n",
      "One word, good citizens.\n",
      "\n",
      "First Citizen:\n",
      "We are accounted poor citizens, the patricians good.\n",
      "What authority surfeits on would relieve us: if they\n",
      "would yield us but the superfluity, while it were\n",
      "wholesome, we might guess they relieved us humanely;\n",
      "but they think we are too dear: the leanness that\n",
      "afflicts us, the object of our misery, is as an\n",
      "inventory to particularise their abundance; our\n",
      "sufferance is a gain to them Let us revenge this with\n",
      "our pikes, ere we become rakes: for the gods know I\n",
      "speak this in hunger for bread, not in thirst for revenge.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# let's look at the first 1000 characters\n",
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bb03c9ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "65\n"
     ]
    }
   ],
   "source": [
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "86681abf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
      "hii there\n"
     ]
    }
   ],
   "source": [
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3b20eecb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115394]) torch.int64\n",
      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
      "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
      "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
      "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
      "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
      "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50, 50,\n",
      "         1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1, 58,\n",
      "        53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51, 47,\n",
      "        57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43, 42,\n",
      "         8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57, 58,\n",
      "         1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1, 63,\n",
      "        53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56, 41,\n",
      "        47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51, 63,\n",
      "         1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0, 13,\n",
      "        50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,  1,\n",
      "        49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57,  1, 49, 47, 50, 50,  1,\n",
      "        46, 47, 51,  6,  1, 39, 52, 42,  1, 61, 43,  5, 50, 50,  1, 46, 39, 60,\n",
      "        43,  1, 41, 53, 56, 52,  1, 39, 58,  1, 53, 59, 56,  1, 53, 61, 52,  1,\n",
      "        54, 56, 47, 41, 43,  8,  0, 21, 57,  5, 58,  1, 39,  1, 60, 43, 56, 42,\n",
      "        47, 41, 58, 12,  0,  0, 13, 50, 50, 10,  0, 26, 53,  1, 51, 53, 56, 43,\n",
      "         1, 58, 39, 50, 49, 47, 52, 45,  1, 53, 52,  5, 58, 11,  1, 50, 43, 58,\n",
      "         1, 47, 58,  1, 40, 43,  1, 42, 53, 52, 43, 10,  1, 39, 61, 39, 63,  6,\n",
      "         1, 39, 61, 39, 63,  2,  0,  0, 31, 43, 41, 53, 52, 42,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 27, 52, 43,  1, 61, 53, 56, 42,  6,  1, 45, 53,\n",
      "        53, 42,  1, 41, 47, 58, 47, 64, 43, 52, 57,  8,  0,  0, 18, 47, 56, 57,\n",
      "        58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 35, 43,  1, 39, 56, 43,  1,\n",
      "        39, 41, 41, 53, 59, 52, 58, 43, 42,  1, 54, 53, 53, 56,  1, 41, 47, 58,\n",
      "        47, 64, 43, 52, 57,  6,  1, 58, 46, 43,  1, 54, 39, 58, 56, 47, 41, 47,\n",
      "        39, 52, 57,  1, 45, 53, 53, 42,  8,  0, 35, 46, 39, 58,  1, 39, 59, 58,\n",
      "        46, 53, 56, 47, 58, 63,  1, 57, 59, 56, 44, 43, 47, 58, 57,  1, 53, 52,\n",
      "         1, 61, 53, 59, 50, 42,  1, 56, 43, 50, 47, 43, 60, 43,  1, 59, 57, 10,\n",
      "         1, 47, 44,  1, 58, 46, 43, 63,  0, 61, 53, 59, 50, 42,  1, 63, 47, 43,\n",
      "        50, 42,  1, 59, 57,  1, 40, 59, 58,  1, 58, 46, 43,  1, 57, 59, 54, 43,\n",
      "        56, 44, 50, 59, 47, 58, 63,  6,  1, 61, 46, 47, 50, 43,  1, 47, 58,  1,\n",
      "        61, 43, 56, 43,  0, 61, 46, 53, 50, 43, 57, 53, 51, 43,  6,  1, 61, 43,\n",
      "         1, 51, 47, 45, 46, 58,  1, 45, 59, 43, 57, 57,  1, 58, 46, 43, 63,  1,\n",
      "        56, 43, 50, 47, 43, 60, 43, 42,  1, 59, 57,  1, 46, 59, 51, 39, 52, 43,\n",
      "        50, 63, 11,  0, 40, 59, 58,  1, 58, 46, 43, 63,  1, 58, 46, 47, 52, 49,\n",
      "         1, 61, 43,  1, 39, 56, 43,  1, 58, 53, 53,  1, 42, 43, 39, 56, 10,  1,\n",
      "        58, 46, 43,  1, 50, 43, 39, 52, 52, 43, 57, 57,  1, 58, 46, 39, 58,  0,\n",
      "        39, 44, 44, 50, 47, 41, 58, 57,  1, 59, 57,  6,  1, 58, 46, 43,  1, 53,\n",
      "        40, 48, 43, 41, 58,  1, 53, 44,  1, 53, 59, 56,  1, 51, 47, 57, 43, 56,\n",
      "        63,  6,  1, 47, 57,  1, 39, 57,  1, 39, 52,  0, 47, 52, 60, 43, 52, 58,\n",
      "        53, 56, 63,  1, 58, 53,  1, 54, 39, 56, 58, 47, 41, 59, 50, 39, 56, 47,\n",
      "        57, 43,  1, 58, 46, 43, 47, 56,  1, 39, 40, 59, 52, 42, 39, 52, 41, 43,\n",
      "        11,  1, 53, 59, 56,  0, 57, 59, 44, 44, 43, 56, 39, 52, 41, 43,  1, 47,\n",
      "        57,  1, 39,  1, 45, 39, 47, 52,  1, 58, 53,  1, 58, 46, 43, 51,  1, 24,\n",
      "        43, 58,  1, 59, 57,  1, 56, 43, 60, 43, 52, 45, 43,  1, 58, 46, 47, 57,\n",
      "         1, 61, 47, 58, 46,  0, 53, 59, 56,  1, 54, 47, 49, 43, 57,  6,  1, 43,\n",
      "        56, 43,  1, 61, 43,  1, 40, 43, 41, 53, 51, 43,  1, 56, 39, 49, 43, 57,\n",
      "        10,  1, 44, 53, 56,  1, 58, 46, 43,  1, 45, 53, 42, 57,  1, 49, 52, 53,\n",
      "        61,  1, 21,  0, 57, 54, 43, 39, 49,  1, 58, 46, 47, 57,  1, 47, 52,  1,\n",
      "        46, 59, 52, 45, 43, 56,  1, 44, 53, 56,  1, 40, 56, 43, 39, 42,  6,  1,\n",
      "        52, 53, 58,  1, 47, 52,  1, 58, 46, 47, 56, 57, 58,  1, 44, 53, 56,  1,\n",
      "        56, 43, 60, 43, 52, 45, 43,  8,  0,  0])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user1/miniconda3/envs/ssl/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
    "import torch # we use PyTorch: https://pytorch.org\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000]) # the 1000 characters we looked at earier will to the GPT look like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d546fb00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's now split up the data into train and validation sets\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c5c3182e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "82547ea9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when input is tensor([18]) the target: 47\n",
      "when input is tensor([18, 47]) the target: 56\n",
      "when input is tensor([18, 47, 56]) the target: 57\n",
      "when input is tensor([18, 47, 56, 57]) the target: 58\n",
      "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
      "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"when input is {context} the target: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0a30ed98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      "torch.Size([4, 8])\n",
      "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
      "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
      "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
      "        [25, 17, 27, 10,  0, 21,  1, 54]])\n",
      "targets:\n",
      "torch.Size([4, 8])\n",
      "tensor([[43, 58,  5, 57,  1, 46, 43, 39],\n",
      "        [53, 56,  1, 58, 46, 39, 58,  1],\n",
      "        [58,  1, 58, 46, 39, 58,  1, 46],\n",
      "        [17, 27, 10,  0, 21,  1, 54, 39]])\n",
      "----\n",
      "when input is [24] the target: 43\n",
      "when input is [24, 43] the target: 58\n",
      "when input is [24, 43, 58] the target: 5\n",
      "when input is [24, 43, 58, 5] the target: 57\n",
      "when input is [24, 43, 58, 5, 57] the target: 1\n",
      "when input is [24, 43, 58, 5, 57, 1] the target: 46\n",
      "when input is [24, 43, 58, 5, 57, 1, 46] the target: 43\n",
      "when input is [24, 43, 58, 5, 57, 1, 46, 43] the target: 39\n",
      "when input is [44] the target: 53\n",
      "when input is [44, 53] the target: 56\n",
      "when input is [44, 53, 56] the target: 1\n",
      "when input is [44, 53, 56, 1] the target: 58\n",
      "when input is [44, 53, 56, 1, 58] the target: 46\n",
      "when input is [44, 53, 56, 1, 58, 46] the target: 39\n",
      "when input is [44, 53, 56, 1, 58, 46, 39] the target: 58\n",
      "when input is [44, 53, 56, 1, 58, 46, 39, 58] the target: 1\n",
      "when input is [52] the target: 58\n",
      "when input is [52, 58] the target: 1\n",
      "when input is [52, 58, 1] the target: 58\n",
      "when input is [52, 58, 1, 58] the target: 46\n",
      "when input is [52, 58, 1, 58, 46] the target: 39\n",
      "when input is [52, 58, 1, 58, 46, 39] the target: 58\n",
      "when input is [52, 58, 1, 58, 46, 39, 58] the target: 1\n",
      "when input is [52, 58, 1, 58, 46, 39, 58, 1] the target: 46\n",
      "when input is [25] the target: 17\n",
      "when input is [25, 17] the target: 27\n",
      "when input is [25, 17, 27] the target: 10\n",
      "when input is [25, 17, 27, 10] the target: 0\n",
      "when input is [25, 17, 27, 10, 0] the target: 21\n",
      "when input is [25, 17, 27, 10, 0, 21] the target: 1\n",
      "when input is [25, 17, 27, 10, 0, 21, 1] the target: 54\n",
      "when input is [25, 17, 27, 10, 0, 21, 1, 54] the target: 39\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "batch_size = 4 # how many independent sequences will we process in parallel?\n",
    "block_size = 8 # what is the maximum context length for predictions?\n",
    "\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "print('inputs:')\n",
    "print(xb.shape)\n",
    "print(xb)\n",
    "print('targets:')\n",
    "print(yb.shape)\n",
    "print(yb)\n",
    "\n",
    "print('----')\n",
    "\n",
    "for b in range(batch_size): # batch dimension\n",
    "    for t in range(block_size): # time dimension\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b,t]\n",
    "        print(f\"when input is {context.tolist()} the target: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8ad46fe7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[24, 43, 58,  5, 57,  1, 46, 43],\n",
      "        [44, 53, 56,  1, 58, 46, 39, 58],\n",
      "        [52, 58,  1, 58, 46, 39, 58,  1],\n",
      "        [25, 17, 27, 10,  0, 21,  1, 54]])\n"
     ]
    }
   ],
   "source": [
    "print(xb) # our input to the transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "73b40fb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 65])\n",
      "tensor(4.8786, grad_fn=<NllLossBackward0>)\n",
      "\n",
      "SKIcLT;AcELMoTbvZv C?nq-QE33:CJqkOKH-q;:la!oiywkHjgChzbQ?u!3bLIgwevmyFJGUGp\n",
      "wnYWmnxKWWev-tDqXErVKLgJ\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        logits = self.token_embedding_table(idx) # (B,T,C)\n",
    "        \n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "m = BigramLanguageModel(vocab_size)\n",
    "logits, loss = m(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=100)[0].tolist()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "837823eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4d5e2524",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.65630578994751\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "for steps in range(100): # increase number of steps for good results... \n",
    "    \n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = m(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "print(loss.item())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "567f7c67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "oTo.JUZ!!zqe!\n",
      "xBP qbs$Gy'AcOmrLwwt\n",
      "p$x;Seh-onQbfM?OjKbn'NwUAW -Np3fkz$FVwAUEa-wzWC -wQo-R!v -Mj?,SPiTyZ;o-opr$mOiPJEYD-CfigkzD3p3?zvS;ADz;.y?o,ivCuC'zqHxcVT cHA\n",
      "rT'Fd,SBMZyOslg!NXeF$sBe,juUzLq?w-wzP-h\n",
      "ERjjxlgJzPbHxf$ q,q,KCDCU fqBOQT\n",
      "SV&CW:xSVwZv'DG'NSPypDhKStKzC -$hslxIVzoivnp ,ethA:NCCGoi\n",
      "tN!ljjP3fwJMwNelgUzzPGJlgihJ!d?q.d\n",
      "pSPYgCuCJrIFtb\n",
      "jQXg\n",
      "pA.P LP,SPJi\n",
      "DBcuBM:CixjJ$Jzkq,OLf3KLQLMGph$O 3DfiPHnXKuHMlyjxEiyZib3FaHV-oJa!zoc'XSP :CKGUhd?lgCOF$;;DTHZMlvvcmZAm;:iv'MMgO&Ywbc;BLCUd&vZINLIzkuTGZa\n",
      "D.?\n"
     ]
    }
   ],
   "source": [
    "print(decode(m.generate(idx = torch.zeros((1, 1), dtype=torch.long), max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a782d7c3",
   "metadata": {},
   "source": [
    "### Maths trick "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8d88de7f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=\n",
      "tensor([[1.0000, 0.0000, 0.0000],\n",
      "        [0.5000, 0.5000, 0.0000],\n",
      "        [0.3333, 0.3333, 0.3333]])\n",
      "--\n",
      "b=\n",
      "tensor([[2., 7.],\n",
      "        [6., 4.],\n",
      "        [6., 5.]])\n",
      "--\n",
      "c=\n",
      "tensor([[2.0000, 7.0000],\n",
      "        [4.0000, 5.5000],\n",
      "        [4.6667, 5.3333]])\n"
     ]
    }
   ],
   "source": [
    "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
    "torch.manual_seed(42)\n",
    "a = torch.tril(torch.ones(3, 3))\n",
    "a = a / torch.sum(a, 1, keepdim=True)\n",
    "b = torch.randint(0,10,(3,2)).float()\n",
    "c = a @ b\n",
    "print('a=')\n",
    "print(a)\n",
    "print('--')\n",
    "print('b=')\n",
    "print(b)\n",
    "print('--')\n",
    "print('c=')\n",
    "print(c)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f0a63dc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 2])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# consider the following toy example:\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,2 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "19184b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
    "xbow = torch.zeros((B,T,C))\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1] # (t,C)\n",
    "        xbow[b,t] = torch.mean(xprev, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f2d92f3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 2: using matrix multiply for a weighted aggregation\n",
    "wei = torch.tril(torch.ones(T, T))\n",
    "wei = wei / wei.sum(1, keepdim=True)\n",
    "xbow2 = wei @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)\n",
    "torch.allclose(xbow, xbow2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "938f63fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 3: use Softmax\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "xbow3 = wei @ x\n",
    "torch.allclose(xbow, xbow3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5558552d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 4: self-attention!\n",
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# let's see a single Head perform self-attention\n",
    "head_size = 16\n",
    "key = nn.Linear(C, head_size, bias=False)\n",
    "query = nn.Linear(C, head_size, bias=False)\n",
    "value = nn.Linear(C, head_size, bias=False)\n",
    "k = key(x)   # (B, T, 16)\n",
    "q = query(x) # (B, T, 16)\n",
    "wei =  q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "#wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "\n",
    "v = value(x)\n",
    "out = wei @ v\n",
    "#out = wei @ x\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "533e8dba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.1574, 0.8426, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.2088, 0.1646, 0.6266, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5792, 0.1187, 0.1889, 0.1131, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0294, 0.1052, 0.0469, 0.0276, 0.7909, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0176, 0.2689, 0.0215, 0.0089, 0.6812, 0.0019, 0.0000, 0.0000],\n",
       "        [0.1691, 0.4066, 0.0438, 0.0416, 0.1048, 0.2012, 0.0329, 0.0000],\n",
       "        [0.0210, 0.0843, 0.0555, 0.2297, 0.0573, 0.0709, 0.2423, 0.2391]],\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wei[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "99a25058",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = torch.randn(B,T,head_size)\n",
    "q = torch.randn(B,T,head_size)\n",
    "wei = q @ k.transpose(-2, -1) * head_size**-0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c11dcda3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.0449)\n",
      "tensor(1.0700)\n",
      "tensor(1.0918)\n"
     ]
    }
   ],
   "source": [
    "print(k.var())\n",
    "print(q.var())\n",
    "print(wei.var())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "64661cb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.1925, 0.1426, 0.2351, 0.1426, 0.2872])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5]), dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a3803b1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0326, 0.0030, 0.1615, 0.0030, 0.8000])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(torch.tensor([0.1, -0.2, 0.3, -0.2, 0.5])*8, dim=-1) # gets too peaky, converges to one-hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "145a7620",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 100])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class LayerNorm1d: # (used to be BatchNorm1d)\n",
    "  \n",
    "  def __init__(self, dim, eps=1e-5, momentum=0.1):\n",
    "    self.eps = eps\n",
    "    self.gamma = torch.ones(dim)\n",
    "    self.beta = torch.zeros(dim)\n",
    "  \n",
    "  def __call__(self, x):\n",
    "    # calculate the forward pass\n",
    "    xmean = x.mean(1, keepdim=True) # batch mean\n",
    "    xvar = x.var(1, keepdim=True) # batch variance\n",
    "    xhat = (x - xmean) / torch.sqrt(xvar + self.eps) # normalize to unit variance\n",
    "    self.out = self.gamma * xhat + self.beta\n",
    "    return self.out\n",
    "  \n",
    "  def parameters(self):\n",
    "    return [self.gamma, self.beta]\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "module = LayerNorm1d(100)\n",
    "x = torch.randn(32, 100) # batch size 32 of 100-dimensional vectors\n",
    "x = module(x)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f2fb84d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(0.1469), tensor(0.8803))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[:,0].mean(), x[:,0].std() # mean,std of one feature across all batch inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "12a8e484",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-9.5367e-09), tensor(1.0000))"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0,:].mean(), x[0,:].std() # mean,std of a single input from the batch, of its features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "328ffb80",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Training and finshing                 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c7746785",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n",
      "0.209729 M parameters\n",
      "step 0: train loss 4.4116, val loss 4.4022\n",
      "step 100: train loss 2.6568, val loss 2.6670\n",
      "step 200: train loss 2.5089, val loss 2.5058\n",
      "step 300: train loss 2.4193, val loss 2.4333\n",
      "step 400: train loss 2.3497, val loss 2.3561\n",
      "step 500: train loss 2.2961, val loss 2.3126\n",
      "step 600: train loss 2.2405, val loss 2.2496\n",
      "step 700: train loss 2.2048, val loss 2.2188\n",
      "step 800: train loss 2.1637, val loss 2.1869\n",
      "step 900: train loss 2.1237, val loss 2.1494\n",
      "step 1000: train loss 2.1026, val loss 2.1291\n",
      "step 1100: train loss 2.0697, val loss 2.1180\n",
      "step 1200: train loss 2.0380, val loss 2.0787\n",
      "step 1300: train loss 2.0262, val loss 2.0645\n",
      "step 1400: train loss 1.9924, val loss 2.0358\n",
      "step 1500: train loss 1.9687, val loss 2.0287\n",
      "step 1600: train loss 1.9644, val loss 2.0484\n",
      "step 1700: train loss 1.9421, val loss 2.0132\n",
      "step 1800: train loss 1.9070, val loss 1.9944\n",
      "step 1900: train loss 1.9080, val loss 1.9887\n",
      "step 2000: train loss 1.8847, val loss 1.9942\n",
      "step 2100: train loss 1.8717, val loss 1.9741\n",
      "step 2200: train loss 1.8598, val loss 1.9623\n",
      "step 2300: train loss 1.8542, val loss 1.9514\n",
      "step 2400: train loss 1.8409, val loss 1.9420\n",
      "step 2500: train loss 1.8155, val loss 1.9429\n",
      "step 2600: train loss 1.8278, val loss 1.9387\n",
      "step 2700: train loss 1.8121, val loss 1.9325\n",
      "step 2800: train loss 1.8012, val loss 1.9188\n",
      "step 2900: train loss 1.8041, val loss 1.9318\n",
      "step 3000: train loss 1.7966, val loss 1.9217\n",
      "step 3100: train loss 1.7722, val loss 1.9222\n",
      "step 3200: train loss 1.7548, val loss 1.9106\n",
      "step 3300: train loss 1.7572, val loss 1.9085\n",
      "step 3400: train loss 1.7549, val loss 1.8944\n",
      "step 3500: train loss 1.7402, val loss 1.8982\n",
      "step 3600: train loss 1.7255, val loss 1.8907\n",
      "step 3700: train loss 1.7313, val loss 1.8885\n",
      "step 3800: train loss 1.7190, val loss 1.8956\n",
      "step 3900: train loss 1.7211, val loss 1.8797\n",
      "step 4000: train loss 1.7156, val loss 1.8661\n",
      "step 4100: train loss 1.7120, val loss 1.8773\n",
      "step 4200: train loss 1.7107, val loss 1.8725\n",
      "step 4300: train loss 1.7000, val loss 1.8479\n",
      "step 4400: train loss 1.7042, val loss 1.8616\n",
      "step 4500: train loss 1.6929, val loss 1.8568\n",
      "step 4600: train loss 1.6860, val loss 1.8353\n",
      "step 4700: train loss 1.6840, val loss 1.8457\n",
      "step 4800: train loss 1.6683, val loss 1.8441\n",
      "step 4900: train loss 1.6715, val loss 1.8396\n",
      "step 4999: train loss 1.6659, val loss 1.8283\n",
      "\n",
      "And they bride will to loves that set bube to take Our my called\n",
      "My art that us his vettar diless,\n",
      "Butwiced you, not to zormme\n",
      "Yout proof in heart mile diliblate is ensent,\n",
      "At at Herilieve to the vilt. Good you muself: now, and thus queen,\n",
      "Now up the speak you: non's nor\n",
      "To this destrity sworly that\n",
      "morly so whrow I by we arm doth will now;\n",
      "But poor of his butt kindnent, the son; if his shat thy flear the said.\n",
      "\n",
      "COMINIUS:\n",
      "He-buse!\n",
      "At you as ards become and this hence as is no ver tomI\n",
      "An you! who neet soke mary.\n",
      "\n",
      "HENRWORD OF GAUNTH:\n",
      "Those me? what the men.\n",
      "\n",
      "CORIOLANUS:\n",
      "What thy were on, wilth thy speakss\n",
      "Dects lome\n",
      "This shettup in strove I deed well.\n",
      "\n",
      "LORD CORIULE:\n",
      "Prack, in throye stingd his refest, as nothit's move\n",
      "Tunness, no dud tone. Mortar.\n",
      "\n",
      "KING RICHARD\n",
      "There's cowar so! who lump not, to so lack.\n",
      "\n",
      "PUTNGEY:\n",
      "Hath is brile.\n",
      "I know the not.\n",
      "\n",
      "POLINA:\n",
      "No, this of time the shope te summe wip.\n",
      "\n",
      "GRISABELLA:\n",
      "Why the sileet Blifenk thou so hate!\n",
      "\n",
      "SLY.\n",
      "\n",
      "BENVINTH:\n",
      "You there fort;\n",
      "To know our her and bruny heart, If\n",
      "twonce whe whill conve, do did tiland,\n",
      "And that honsury!\n",
      "\n",
      "AMINIUS:\n",
      "Yet not sever worting-deness of them Warwast did chial friend!\n",
      "\n",
      "BUCKINGIn:\n",
      "You arm in but may incoment.\n",
      "In to dot me, to librefeleing brothn the shall eath I die;\n",
      "Save thou for Heaven to Tumself:\n",
      "How the oldrefest jurse so vow\n",
      "Straught?\n",
      "\n",
      "Seliving:\n",
      "Give it is stad as lothry on the duty:\n",
      "I hurts begut,\n",
      "Good vonsure that thou that be sweet;\n",
      "And him yongety may's froes must,\n",
      "And I issome you, so are woll, we\n",
      "epen me is justibure themer, saye as;\n",
      "River thosk virtace my from my been;\n",
      "Our serval this broth ever clupber?\n",
      "\n",
      "MORCULET:\n",
      "There this is\n",
      "'tward them make nume while as I\n",
      "was that becomfort, themors Coling, this send; commant the stright flantess\n",
      "As I han cort, thou tatterrung 'I only talk we twretter tonguius, sir, I'll but the peepose: life?\n",
      "\n",
      "COMINIUS:\n",
      "Yet to the art; Sack you.\n",
      "\n",
      "CORIOLANUS:\n",
      "May, make, whose trong what Vence stan,\n",
      "Threse pating to swit shown'd for a dresss me; he heave degrave wi\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "# hyperparameters\n",
    "batch_size = 16 # how many independent sequences will we process in parallel?\n",
    "block_size = 32 # what is the maximum context length for predictions?\n",
    "max_iters = 5000\n",
    "eval_interval = 100\n",
    "learning_rate = 1e-3\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(device)\n",
    "eval_iters = 200\n",
    "n_embd = 64\n",
    "n_head = 4\n",
    "n_layer = 4\n",
    "dropout = 0.0\n",
    "# ------------\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "# wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "# Train and test splits\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "\n",
    "# data loading\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    return x, y\n",
    "\n",
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            logits, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out\n",
    "\n",
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)   # (B,T,C)\n",
    "        q = self.query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B,T,C)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
    "        return out\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embd, n_embd)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n",
    "\n",
    "class FeedFoward(nn.Module):\n",
    "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd, n_head):\n",
    "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
    "        super().__init__()\n",
    "        head_size = n_embd // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.ffwd = FeedFoward(n_embd)\n",
    "        self.ln1 = nn.LayerNorm(n_embd)\n",
    "        self.ln2 = nn.LayerNorm(n_embd)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.ffwd(self.ln2(x))\n",
    "        return x\n",
    "\n",
    "# super simple bigram model\n",
    "class BigramLanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
    "        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(n_embd) # final layer norm\n",
    "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)\n",
    "        x = tok_emb + pos_emb # (B,T,C)\n",
    "        x = self.blocks(x) # (B,T,C)\n",
    "        x = self.ln_f(x) # (B,T,C)\n",
    "        logits = self.lm_head(x) # (B,T,vocab_size)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cond)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n",
    "\n",
    "model = BigramLanguageModel()\n",
    "m = model.to(device)\n",
    "# print the number of parameters in the model\n",
    "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "\n",
    "    # every once in a while evaluate the loss on train and val sets\n",
    "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# generate from the model\n",
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d1099cdc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('token_embedding_table.weight',\n",
       "              tensor([[ 0.2455,  0.0870, -0.4051,  ...,  1.4293,  1.6123, -0.6675],\n",
       "                      [-0.7927,  0.5708, -0.0724,  ...,  0.4141, -0.5784, -1.5749],\n",
       "                      [-1.4926,  2.0634,  1.4399,  ...,  1.6246,  0.5758,  1.3313],\n",
       "                      ...,\n",
       "                      [ 1.2002, -0.3514,  1.2566,  ...,  0.2636,  0.3822, -1.6100],\n",
       "                      [-0.9023, -0.4427,  0.1572,  ...,  0.3966,  0.4549, -0.7686],\n",
       "                      [ 0.7100, -0.5336, -0.0900,  ...,  1.1289,  0.2578,  0.5086]],\n",
       "                     device='cuda:0')),\n",
       "             ('position_embedding_table.weight',\n",
       "              tensor([[-1.3082, -0.3272, -2.2798,  ...,  0.5896,  1.8284,  1.1728],\n",
       "                      [-0.1714,  1.3774,  0.2953,  ...,  0.0867, -1.0561,  1.5210],\n",
       "                      [ 0.4868, -0.8680, -0.6752,  ...,  1.9570,  0.0840,  1.2837],\n",
       "                      ...,\n",
       "                      [ 0.5011, -1.0728,  1.1017,  ..., -0.9873,  0.6694,  0.6457],\n",
       "                      [ 1.7188, -0.2847,  0.5471,  ..., -0.7362,  0.4446, -0.2157],\n",
       "                      [ 0.5420, -0.7003,  2.4963,  ..., -0.0614, -0.3255, -0.8274]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.0.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.0.key.weight',\n",
       "              tensor([[ 0.0142,  0.1017,  0.2047,  ...,  0.0117, -0.0376,  0.0397],\n",
       "                      [-0.0675,  0.0498,  0.1320,  ...,  0.0431, -0.0530, -0.0862],\n",
       "                      [ 0.0291, -0.1040,  0.0811,  ...,  0.0355,  0.0526, -0.1297],\n",
       "                      ...,\n",
       "                      [-0.0116, -0.0558,  0.0833,  ...,  0.1111, -0.1144,  0.1313],\n",
       "                      [ 0.0283, -0.0129, -0.0628,  ...,  0.0294, -0.0335, -0.1116],\n",
       "                      [ 0.1776,  0.0173, -0.1293,  ...,  0.0240,  0.0787, -0.0807]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.0.query.weight',\n",
       "              tensor([[-0.0747, -0.0447,  0.0677,  ..., -0.0718, -0.0020,  0.1236],\n",
       "                      [-0.0050,  0.0773,  0.0529,  ...,  0.0603,  0.1253, -0.1005],\n",
       "                      [ 0.0712,  0.1076,  0.0208,  ..., -0.1852,  0.0821,  0.0734],\n",
       "                      ...,\n",
       "                      [-0.1244, -0.0726, -0.1180,  ...,  0.0145, -0.0173, -0.0159],\n",
       "                      [-0.1628,  0.0925,  0.1004,  ...,  0.0825,  0.0305,  0.0688],\n",
       "                      [ 0.0233, -0.0218, -0.0678,  ..., -0.2253,  0.0445,  0.0690]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.0.value.weight',\n",
       "              tensor([[-0.0613,  0.1186,  0.1507,  ..., -0.1418, -0.1388, -0.0142],\n",
       "                      [ 0.0362, -0.1657, -0.0987,  ..., -0.0348,  0.1311, -0.0041],\n",
       "                      [-0.0037,  0.0664,  0.0488,  ...,  0.1324, -0.0094, -0.0549],\n",
       "                      ...,\n",
       "                      [-0.1043,  0.1192, -0.0395,  ...,  0.0870,  0.0598,  0.0782],\n",
       "                      [-0.1223, -0.0450,  0.0032,  ..., -0.0812, -0.0493,  0.0607],\n",
       "                      [-0.1257,  0.0274, -0.0531,  ..., -0.1205,  0.1156, -0.0698]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.1.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.1.key.weight',\n",
       "              tensor([[-0.0535,  0.0117, -0.1071,  ...,  0.0769,  0.1350,  0.0614],\n",
       "                      [-0.0509, -0.0009, -0.0703,  ...,  0.0489, -0.0748, -0.0002],\n",
       "                      [ 0.1309,  0.0834, -0.0438,  ...,  0.0470, -0.1343, -0.0534],\n",
       "                      ...,\n",
       "                      [ 0.0157,  0.1233, -0.1563,  ..., -0.0077,  0.1214, -0.1601],\n",
       "                      [ 0.0327,  0.1214,  0.1175,  ...,  0.1777,  0.0081,  0.0254],\n",
       "                      [-0.1680,  0.0435, -0.0182,  ..., -0.0751, -0.1260,  0.0696]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.1.query.weight',\n",
       "              tensor([[ 0.0233,  0.0394, -0.0004,  ..., -0.0617, -0.0156,  0.1277],\n",
       "                      [-0.0924,  0.0031,  0.0962,  ..., -0.0756,  0.0337,  0.0618],\n",
       "                      [ 0.0004,  0.0882,  0.1890,  ...,  0.0726, -0.1134,  0.0303],\n",
       "                      ...,\n",
       "                      [-0.0788,  0.0177, -0.0450,  ...,  0.0760, -0.0092,  0.2016],\n",
       "                      [-0.1048, -0.0638, -0.0979,  ...,  0.0392,  0.0087,  0.0009],\n",
       "                      [ 0.1675, -0.2349,  0.1428,  ..., -0.0033,  0.0102,  0.0715]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.1.value.weight',\n",
       "              tensor([[-0.1000, -0.0332,  0.0326,  ...,  0.0425,  0.1544, -0.0472],\n",
       "                      [ 0.0853,  0.0444,  0.0176,  ...,  0.0009, -0.1492,  0.0737],\n",
       "                      [ 0.0593, -0.1281,  0.0981,  ..., -0.0045, -0.1221, -0.1179],\n",
       "                      ...,\n",
       "                      [-0.1076, -0.0591,  0.1124,  ..., -0.0602, -0.0403,  0.0483],\n",
       "                      [-0.0180, -0.2200, -0.0034,  ..., -0.0889, -0.0287,  0.1028],\n",
       "                      [ 0.0695,  0.0374,  0.0006,  ...,  0.0146,  0.1038, -0.0456]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.2.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.2.key.weight',\n",
       "              tensor([[-0.0864,  0.0758, -0.2718,  ...,  0.0111,  0.1376, -0.0230],\n",
       "                      [-0.0015, -0.0926,  0.1398,  ...,  0.0484,  0.1394, -0.1225],\n",
       "                      [-0.0300,  0.1788, -0.0590,  ...,  0.0273,  0.0132, -0.0074],\n",
       "                      ...,\n",
       "                      [ 0.0189, -0.0550, -0.0109,  ...,  0.1302, -0.2090,  0.1781],\n",
       "                      [-0.0366, -0.1033, -0.1752,  ...,  0.0654,  0.1225, -0.0844],\n",
       "                      [ 0.1015,  0.0554,  0.0690,  ..., -0.1065,  0.0433,  0.1002]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.2.query.weight',\n",
       "              tensor([[-0.0921,  0.0747,  0.0683,  ..., -0.0583,  0.0589,  0.2188],\n",
       "                      [-0.0791, -0.0871,  0.0529,  ...,  0.1156, -0.0180, -0.0559],\n",
       "                      [-0.1976, -0.1670, -0.1147,  ...,  0.0791, -0.0733, -0.0541],\n",
       "                      ...,\n",
       "                      [-0.0798, -0.1661,  0.0292,  ..., -0.0177,  0.0812,  0.1590],\n",
       "                      [ 0.0093,  0.0030,  0.0763,  ...,  0.0910, -0.1443,  0.1348],\n",
       "                      [ 0.0278,  0.2109,  0.0769,  ..., -0.0761,  0.0308, -0.0073]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.2.value.weight',\n",
       "              tensor([[ 0.0823, -0.0307,  0.0433,  ...,  0.0642, -0.0151, -0.1240],\n",
       "                      [ 0.0632, -0.0556,  0.0314,  ..., -0.0750, -0.0225,  0.0143],\n",
       "                      [-0.0923, -0.0264, -0.0679,  ...,  0.0532,  0.0154, -0.0728],\n",
       "                      ...,\n",
       "                      [-0.0342, -0.0027,  0.0407,  ..., -0.0027, -0.0316, -0.1784],\n",
       "                      [-0.0725, -0.1238, -0.0795,  ...,  0.1122, -0.0545, -0.0950],\n",
       "                      [ 0.0343, -0.0312,  0.1069,  ..., -0.0388,  0.0370, -0.1411]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.3.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.3.key.weight',\n",
       "              tensor([[ 0.1626, -0.1699, -0.1493,  ..., -0.1004, -0.2252, -0.0531],\n",
       "                      [-0.3048,  0.0890, -0.1326,  ...,  0.1749,  0.0720, -0.1145],\n",
       "                      [-0.1682, -0.0204,  0.1876,  ..., -0.1685,  0.1313, -0.2741],\n",
       "                      ...,\n",
       "                      [-0.1879,  0.2321,  0.1512,  ...,  0.1086,  0.0363,  0.1680],\n",
       "                      [-0.0837, -0.0269, -0.2448,  ...,  0.0219, -0.2006,  0.1304],\n",
       "                      [ 0.1091,  0.0342, -0.0167,  ...,  0.0174,  0.2407, -0.1244]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.3.query.weight',\n",
       "              tensor([[-0.0994,  0.1137, -0.0169,  ...,  0.1429, -0.0847,  0.1948],\n",
       "                      [-0.1284,  0.0073,  0.0741,  ..., -0.0379,  0.0331, -0.0365],\n",
       "                      [-0.1182, -0.0753,  0.2778,  ...,  0.3122,  0.0764, -0.2038],\n",
       "                      ...,\n",
       "                      [-0.1151, -0.0604, -0.2630,  ...,  0.0674,  0.0942,  0.0040],\n",
       "                      [ 0.0751,  0.2262,  0.0404,  ...,  0.0259, -0.1320, -0.0967],\n",
       "                      [ 0.0567,  0.1524, -0.1154,  ...,  0.1067,  0.0203, -0.1023]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.heads.3.value.weight',\n",
       "              tensor([[-0.1107, -0.0428,  0.0787,  ..., -0.0532, -0.0504, -0.0123],\n",
       "                      [-0.0495,  0.0551, -0.1235,  ..., -0.0313, -0.0926, -0.0118],\n",
       "                      [-0.0936, -0.0802, -0.1006,  ...,  0.0011, -0.0020, -0.1397],\n",
       "                      ...,\n",
       "                      [ 0.0163, -0.1464, -0.1473,  ...,  0.0655, -0.0485, -0.0237],\n",
       "                      [ 0.0793,  0.0352,  0.0029,  ...,  0.0972,  0.1105, -0.0607],\n",
       "                      [-0.0867,  0.0774,  0.0464,  ...,  0.0363, -0.0985, -0.0947]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.proj.weight',\n",
       "              tensor([[-0.2200, -0.1119, -0.0827,  ...,  0.0765, -0.0626, -0.0590],\n",
       "                      [ 0.0383,  0.0786, -0.0179,  ...,  0.0432, -0.0050,  0.0038],\n",
       "                      [ 0.0966, -0.0166,  0.1416,  ..., -0.1444, -0.0754,  0.1822],\n",
       "                      ...,\n",
       "                      [-0.1344,  0.0546,  0.0706,  ...,  0.1610,  0.0169,  0.1325],\n",
       "                      [ 0.1193,  0.0966,  0.0633,  ..., -0.0180, -0.1164, -0.0213],\n",
       "                      [ 0.1628,  0.0152,  0.0947,  ...,  0.1199, -0.1246,  0.0173]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.sa.proj.bias',\n",
       "              tensor([ 0.1111, -0.0684,  0.0909, -0.0253, -0.0906,  0.0223,  0.0103,  0.0504,\n",
       "                       0.1183,  0.0308,  0.0056,  0.0923, -0.0954,  0.1058, -0.1002, -0.0293,\n",
       "                      -0.1105,  0.1510,  0.0283,  0.0056,  0.0017,  0.0891, -0.0680, -0.0565,\n",
       "                      -0.1607,  0.0340,  0.1066, -0.1141, -0.0094, -0.0983, -0.0414, -0.0449,\n",
       "                       0.0590,  0.0828,  0.0741, -0.0537, -0.0280, -0.0626,  0.0552, -0.0463,\n",
       "                       0.0183, -0.1182,  0.1241, -0.0168, -0.0278, -0.0138, -0.0721, -0.0719,\n",
       "                       0.1271,  0.0748,  0.1379,  0.1115, -0.0178,  0.0574,  0.0435,  0.0538,\n",
       "                       0.1008, -0.0320, -0.0469,  0.0966,  0.0584,  0.0863, -0.1180, -0.1034],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.ffwd.net.0.weight',\n",
       "              tensor([[-0.1692, -0.0754,  0.1036,  ...,  0.0862, -0.0434,  0.0491],\n",
       "                      [ 0.0518,  0.1067, -0.0770,  ..., -0.1064, -0.1290, -0.0513],\n",
       "                      [-0.2679,  0.0549,  0.0285,  ..., -0.0526, -0.1217, -0.0816],\n",
       "                      ...,\n",
       "                      [ 0.0810, -0.0491,  0.0125,  ..., -0.0431, -0.0641, -0.1029],\n",
       "                      [-0.0303, -0.0466, -0.1509,  ..., -0.0099,  0.1543,  0.1482],\n",
       "                      [-0.1250, -0.0288,  0.0076,  ..., -0.0594,  0.0635,  0.0722]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.ffwd.net.0.bias',\n",
       "              tensor([-9.3641e-02, -1.8075e-02, -6.5299e-02,  4.3750e-02, -7.8887e-02,\n",
       "                      -1.0220e-01,  1.1170e-02, -6.3058e-02, -1.1660e-02, -3.6382e-02,\n",
       "                      -6.3457e-02, -4.9641e-02, -5.0533e-02, -7.7724e-03, -1.2655e-01,\n",
       "                      -8.5901e-02,  5.9935e-02, -1.1923e-01, -2.7251e-02, -1.1677e-01,\n",
       "                      -6.9230e-02,  3.0953e-02, -3.9959e-02,  3.1773e-02, -4.7053e-02,\n",
       "                      -4.6065e-02, -1.6132e-01, -1.7270e-01, -2.0503e-01, -4.0579e-02,\n",
       "                      -1.7623e-01, -1.3295e-02, -1.0302e-01,  2.0238e-02, -8.1712e-02,\n",
       "                      -3.1726e-02, -1.8525e-02,  6.2829e-03, -3.2591e-02, -1.5445e-01,\n",
       "                      -1.8875e-01, -9.5143e-02, -6.2853e-02, -1.0420e-01, -1.8044e-01,\n",
       "                      -4.7940e-02, -1.4378e-01, -3.4493e-02, -7.3579e-02, -1.5875e-01,\n",
       "                       3.6695e-02,  1.8362e-02, -1.6483e-01, -8.5195e-02, -1.2113e-01,\n",
       "                      -9.1663e-02, -7.0825e-02,  2.5787e-02, -3.8433e-02, -1.5134e-02,\n",
       "                      -2.3255e-01, -7.9279e-02,  6.5524e-02, -6.3142e-02, -8.2643e-02,\n",
       "                      -7.1994e-02, -4.3829e-02, -1.1081e-01, -1.7041e-01, -3.6671e-02,\n",
       "                      -1.9265e-01, -1.1491e-01,  4.4615e-03, -1.1934e-01, -2.1172e-01,\n",
       "                      -1.1599e-01, -6.3704e-02,  8.2587e-03, -1.3614e-01, -6.4797e-02,\n",
       "                      -1.6492e-01, -1.9677e-01, -1.1685e-01, -3.2501e-02,  1.5106e-02,\n",
       "                      -6.5303e-02, -4.6378e-02, -4.5484e-02, -1.3778e-01, -1.7079e-01,\n",
       "                      -9.5130e-02, -1.9933e-01,  3.6463e-02, -2.9198e-03, -1.8819e-01,\n",
       "                      -7.8534e-02, -1.1650e-01, -1.1075e-01, -7.4441e-02, -7.8082e-02,\n",
       "                      -1.2387e-01, -4.7206e-02, -1.0934e-01,  1.4217e-02, -1.0062e-01,\n",
       "                       3.1468e-02,  5.0939e-02, -6.1927e-02, -1.0300e-01, -1.7690e-01,\n",
       "                      -1.4570e-01, -1.6044e-01, -1.4566e-01, -6.2750e-02, -6.2668e-02,\n",
       "                      -1.4117e-01, -1.3409e-01, -2.6587e-03, -7.7777e-02, -5.1584e-02,\n",
       "                      -9.0261e-02, -1.5376e-01, -1.3476e-01,  4.0472e-02, -6.4316e-02,\n",
       "                       6.1391e-03, -1.6525e-01, -1.0168e-01,  2.6735e-02, -7.0976e-02,\n",
       "                      -1.4564e-02, -6.5228e-02, -1.9042e-01, -7.4502e-02, -1.2293e-01,\n",
       "                      -1.3013e-01,  5.1964e-02, -3.8060e-02, -9.6581e-03, -3.2112e-02,\n",
       "                      -1.9187e-01,  2.5718e-02, -1.5150e-01,  9.8106e-02, -1.2846e-01,\n",
       "                      -1.6555e-01, -1.5138e-01, -1.1421e-01, -1.8540e-01, -4.7490e-02,\n",
       "                      -4.1826e-02, -1.1403e-01,  1.0015e-03, -6.4667e-02, -1.6575e-01,\n",
       "                      -1.2470e-01, -1.1531e-01, -4.6254e-02, -2.9798e-02,  1.3226e-01,\n",
       "                      -1.2418e-01, -3.2242e-02, -1.2059e-01, -1.9732e-01, -1.3800e-01,\n",
       "                       1.7427e-02,  2.2090e-02, -1.5735e-01, -4.8364e-02, -1.5064e-01,\n",
       "                      -8.7687e-02, -1.5581e-01, -1.6753e-01, -4.8957e-02,  7.0032e-02,\n",
       "                      -5.0370e-02, -8.0516e-02, -1.0552e-01,  2.8329e-02, -1.3517e-01,\n",
       "                      -9.0754e-02, -1.0216e-01, -1.7625e-01,  5.1985e-05,  3.2638e-02,\n",
       "                      -2.1281e-01, -2.2649e-01, -4.0125e-03, -1.0758e-01, -2.1236e-02,\n",
       "                      -1.2340e-01, -2.3641e-01, -1.7914e-01, -1.2445e-01, -3.6621e-02,\n",
       "                      -1.1069e-01, -5.1167e-02, -8.3441e-02, -1.3781e-01, -1.1032e-01,\n",
       "                      -4.8267e-02, -3.6854e-02, -2.5165e-02, -3.2702e-02, -1.3283e-01,\n",
       "                       2.1036e-02, -3.3808e-02, -1.8182e-01, -8.9518e-02, -1.3109e-01,\n",
       "                      -4.2062e-02, -1.2427e-01, -1.2321e-01,  2.7086e-02, -1.3414e-01,\n",
       "                      -1.1706e-02, -1.1253e-01, -1.3604e-02, -1.0775e-01, -7.1758e-02,\n",
       "                       6.5471e-02, -5.4388e-02, -1.7092e-01, -1.7075e-01, -4.9307e-02,\n",
       "                      -7.7769e-02, -4.7703e-02, -1.5658e-01, -3.6917e-03, -2.5377e-02,\n",
       "                       3.4374e-02, -1.2747e-01, -6.0491e-02, -1.3708e-02, -2.0835e-01,\n",
       "                      -2.1566e-01, -1.1263e-01, -2.5316e-02, -5.8362e-03, -1.5112e-01,\n",
       "                       5.1195e-03, -1.6659e-01, -6.4608e-02, -2.8780e-02, -6.0426e-02,\n",
       "                      -9.3351e-02, -8.6870e-02, -1.2986e-01,  2.9376e-02, -1.6608e-01,\n",
       "                      -1.9676e-01, -1.8089e-01, -7.8994e-02, -1.1279e-01, -1.6545e-01,\n",
       "                      -9.0822e-02], device='cuda:0')),\n",
       "             ('blocks.0.ffwd.net.2.weight',\n",
       "              tensor([[-0.0894,  0.0460, -0.0389,  ..., -0.0051, -0.0294,  0.0294],\n",
       "                      [ 0.1298,  0.0210, -0.0977,  ..., -0.0269,  0.0959, -0.0151],\n",
       "                      [-0.0173,  0.0035, -0.0887,  ...,  0.0470, -0.0254, -0.0364],\n",
       "                      ...,\n",
       "                      [-0.0241, -0.0365,  0.0114,  ...,  0.0478, -0.0923, -0.0328],\n",
       "                      [ 0.0114,  0.0738, -0.0598,  ...,  0.0419,  0.0058, -0.0886],\n",
       "                      [-0.0533,  0.0351, -0.0786,  ...,  0.0159, -0.0066, -0.0426]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.ffwd.net.2.bias',\n",
       "              tensor([-0.0824, -0.0594,  0.0519,  0.0557,  0.0029,  0.0181,  0.0537,  0.0198,\n",
       "                       0.0340, -0.0867, -0.0558, -0.0537,  0.0388,  0.0504,  0.0713,  0.0020,\n",
       "                       0.0323, -0.0492,  0.1105, -0.0408, -0.0476, -0.0217,  0.0697, -0.0292,\n",
       "                       0.0268,  0.0378,  0.0301, -0.0231, -0.0420, -0.0564,  0.0387,  0.0265,\n",
       "                       0.0325,  0.0312,  0.0341, -0.0033, -0.0481, -0.0230, -0.0081,  0.0139,\n",
       "                       0.0309,  0.0354, -0.0363,  0.0041,  0.0767, -0.0068,  0.0153,  0.0143,\n",
       "                       0.0935,  0.0001, -0.0061, -0.0347, -0.0117, -0.0218,  0.0420,  0.0374,\n",
       "                      -0.0348, -0.0320,  0.0275,  0.0252, -0.0651,  0.0615, -0.0595,  0.0056],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.ln1.weight',\n",
       "              tensor([0.8790, 0.9510, 0.9947, 0.9311, 0.9951, 0.8903, 1.1264, 0.8058, 0.8551,\n",
       "                      0.9159, 0.8606, 0.8188, 0.9224, 0.9935, 0.9019, 0.8729, 0.9244, 0.8383,\n",
       "                      0.8196, 0.9905, 0.9886, 1.0018, 0.9553, 0.9641, 0.8609, 1.0249, 0.9351,\n",
       "                      0.9614, 0.9793, 1.0046, 1.0061, 0.8824, 0.8629, 1.0310, 0.9234, 0.9247,\n",
       "                      0.8604, 0.9750, 0.9540, 0.8807, 1.0032, 0.9881, 0.8954, 0.8694, 1.0604,\n",
       "                      0.9157, 0.9454, 0.9869, 0.8277, 1.0447, 0.8927, 0.8942, 0.8672, 0.9384,\n",
       "                      0.9871, 0.8828, 0.9034, 0.9945, 0.9017, 0.9718, 0.9183, 0.8676, 0.8991,\n",
       "                      0.9434], device='cuda:0')),\n",
       "             ('blocks.0.ln1.bias',\n",
       "              tensor([ 0.0613,  0.0050,  0.0483,  0.1395,  0.0158, -0.0619, -0.0328,  0.0292,\n",
       "                      -0.0648,  0.0090, -0.0340,  0.0033,  0.0763,  0.0961,  0.0945, -0.1017,\n",
       "                      -0.1350,  0.0008, -0.0138,  0.0114,  0.0748,  0.0035,  0.0332,  0.0041,\n",
       "                      -0.0451, -0.0511,  0.0866, -0.0185, -0.1472, -0.0237,  0.0089, -0.0274,\n",
       "                       0.0548,  0.0621,  0.0861,  0.0150, -0.0636, -0.0642,  0.0072,  0.0476,\n",
       "                      -0.0130, -0.0133, -0.0658, -0.0817, -0.0630,  0.0270, -0.0055,  0.1628,\n",
       "                       0.0680,  0.0583,  0.0017,  0.0202, -0.0720,  0.0563, -0.0873, -0.0176,\n",
       "                      -0.0383, -0.0643, -0.1753,  0.0196,  0.0178, -0.0188,  0.0084, -0.0593],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.0.ln2.weight',\n",
       "              tensor([0.8695, 0.8624, 0.8796, 0.7782, 0.8445, 0.8801, 0.9025, 0.9278, 0.8306,\n",
       "                      0.8945, 0.8804, 0.7542, 0.9090, 0.8404, 0.9239, 0.9213, 0.8723, 0.8737,\n",
       "                      0.9057, 0.8868, 0.7666, 0.8448, 0.8365, 0.9552, 0.8717, 0.7959, 0.9076,\n",
       "                      0.8655, 0.8501, 0.8809, 0.8313, 0.8106, 0.7949, 0.8737, 0.9197, 0.8866,\n",
       "                      0.8844, 0.8777, 0.8628, 0.8211, 0.8916, 0.7686, 0.8372, 0.8506, 0.8471,\n",
       "                      0.8871, 0.9132, 0.8317, 0.8138, 0.7409, 0.8395, 0.9099, 0.8855, 0.8786,\n",
       "                      0.8137, 0.8498, 0.8163, 0.8814, 0.8681, 0.8795, 0.8754, 0.8315, 0.6749,\n",
       "                      0.9095], device='cuda:0')),\n",
       "             ('blocks.0.ln2.bias',\n",
       "              tensor([ 0.1000, -0.0357,  0.0358, -0.0037,  0.0006,  0.0060, -0.0267, -0.0161,\n",
       "                      -0.0405, -0.0299, -0.1126, -0.0618,  0.0461,  0.0462, -0.0614, -0.0269,\n",
       "                      -0.0179,  0.0395,  0.0156,  0.0194, -0.0341,  0.0492,  0.0417, -0.0031,\n",
       "                      -0.1411, -0.0482,  0.0271, -0.0723, -0.0556, -0.0441,  0.0520,  0.0438,\n",
       "                       0.0313, -0.0196,  0.0542, -0.0206,  0.0313, -0.0274,  0.0225, -0.0069,\n",
       "                      -0.0599,  0.0024,  0.0208, -0.0311, -0.0297, -0.0230, -0.0089, -0.0181,\n",
       "                       0.0327,  0.0174, -0.0009,  0.0346,  0.0604,  0.0358, -0.0206,  0.0640,\n",
       "                       0.0519,  0.0131, -0.1517, -0.0065,  0.0120,  0.0304, -0.0203, -0.0301],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.0.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.0.key.weight',\n",
       "              tensor([[ 0.0747, -0.0772, -0.2825,  ..., -0.1016, -0.1220,  0.2249],\n",
       "                      [ 0.0854,  0.0724, -0.1866,  ..., -0.2101,  0.0926,  0.1290],\n",
       "                      [ 0.0755,  0.1257,  0.0547,  ..., -0.0379,  0.0074, -0.2915],\n",
       "                      ...,\n",
       "                      [ 0.0471, -0.2120,  0.0606,  ...,  0.3132, -0.2037, -0.1109],\n",
       "                      [-0.1250, -0.0606,  0.0705,  ..., -0.0203, -0.0500, -0.1239],\n",
       "                      [ 0.1249,  0.2675,  0.1002,  ..., -0.1612,  0.1200, -0.0585]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.0.query.weight',\n",
       "              tensor([[ 0.1625,  0.1599, -0.2043,  ..., -0.0934,  0.1656, -0.1817],\n",
       "                      [-0.0650,  0.2611, -0.0437,  ..., -0.1270,  0.0168, -0.0444],\n",
       "                      [-0.2373, -0.2994,  0.1361,  ..., -0.0119, -0.0732,  0.2035],\n",
       "                      ...,\n",
       "                      [ 0.1784,  0.0945,  0.0159,  ..., -0.1381, -0.1667,  0.0087],\n",
       "                      [ 0.0549,  0.0451,  0.0373,  ..., -0.0354,  0.0971,  0.1204],\n",
       "                      [-0.2415,  0.1142,  0.3086,  ...,  0.0856,  0.1400, -0.2769]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.0.value.weight',\n",
       "              tensor([[ 0.1308,  0.2366,  0.1472,  ..., -0.1138, -0.0403, -0.0425],\n",
       "                      [-0.0051, -0.0475, -0.0589,  ...,  0.0490,  0.0852,  0.0566],\n",
       "                      [ 0.1173,  0.1239,  0.2045,  ..., -0.1364,  0.1735, -0.0056],\n",
       "                      ...,\n",
       "                      [-0.1798,  0.0881, -0.1000,  ...,  0.0821,  0.0579, -0.1903],\n",
       "                      [-0.0253, -0.0359,  0.1312,  ..., -0.1049, -0.1207, -0.0171],\n",
       "                      [ 0.0800,  0.0058, -0.0405,  ..., -0.0901,  0.0450, -0.0543]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.1.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.1.key.weight',\n",
       "              tensor([[-0.0185, -0.1696, -0.1233,  ..., -0.0343,  0.0097, -0.0954],\n",
       "                      [-0.0995, -0.3759, -0.0181,  ..., -0.0409,  0.1806, -0.0987],\n",
       "                      [-0.0639, -0.0473, -0.3480,  ..., -0.0798,  0.0903,  0.2450],\n",
       "                      ...,\n",
       "                      [-0.1611,  0.3480, -0.1057,  ..., -0.1452,  0.1801, -0.0555],\n",
       "                      [-0.3582, -0.3579, -0.5056,  ..., -0.0552, -0.0255,  0.1462],\n",
       "                      [ 0.2622, -0.0295,  0.2346,  ..., -0.3058,  0.0218, -0.0094]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.1.query.weight',\n",
       "              tensor([[-0.0993,  0.0518,  0.1372,  ...,  0.0774, -0.0675,  0.1525],\n",
       "                      [ 0.2518,  0.2964,  0.0282,  ..., -0.0662, -0.0127, -0.0058],\n",
       "                      [ 0.1743,  0.1503,  0.0755,  ..., -0.1535,  0.4824,  0.0417],\n",
       "                      ...,\n",
       "                      [ 0.0205, -0.0110,  0.1765,  ..., -0.1592,  0.1267, -0.2226],\n",
       "                      [ 0.1527,  0.0165,  0.2648,  ..., -0.0459,  0.2393,  0.1111],\n",
       "                      [ 0.2160,  0.0007,  0.2522,  ..., -0.0619,  0.1192,  0.0840]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.1.value.weight',\n",
       "              tensor([[-0.0277,  0.0624,  0.1513,  ...,  0.0518, -0.1195,  0.0479],\n",
       "                      [ 0.0197,  0.2015,  0.0188,  ..., -0.0419,  0.1684, -0.2290],\n",
       "                      [ 0.0006,  0.2013,  0.1372,  ..., -0.0386, -0.0749,  0.1597],\n",
       "                      ...,\n",
       "                      [-0.0870, -0.1923, -0.1649,  ...,  0.0381, -0.0784, -0.1235],\n",
       "                      [-0.0564,  0.0653,  0.1109,  ..., -0.1142, -0.1456, -0.2111],\n",
       "                      [-0.1600,  0.0254,  0.1187,  ..., -0.1499, -0.0284, -0.0079]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.2.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.2.key.weight',\n",
       "              tensor([[ 0.0267,  0.1337,  0.0090,  ..., -0.0472, -0.0720, -0.0137],\n",
       "                      [ 0.0287, -0.1432, -0.1863,  ...,  0.0501, -0.1564,  0.2060],\n",
       "                      [ 0.2294,  0.1978, -0.0114,  ..., -0.1615,  0.2221,  0.0718],\n",
       "                      ...,\n",
       "                      [ 0.0760,  0.0546,  0.1694,  ...,  0.0567, -0.0286,  0.0915],\n",
       "                      [ 0.0727, -0.2809, -0.1342,  ...,  0.0048, -0.1268,  0.3077],\n",
       "                      [ 0.0886, -0.0331, -0.0403,  ..., -0.0501, -0.1920,  0.1285]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.2.query.weight',\n",
       "              tensor([[-0.2423,  0.0343,  0.1980,  ..., -0.0136, -0.2656, -0.0048],\n",
       "                      [ 0.0473,  0.1254, -0.1179,  ..., -0.1441,  0.0682,  0.3755],\n",
       "                      [-0.2281,  0.1629, -0.0572,  ...,  0.3243, -0.1570, -0.2063],\n",
       "                      ...,\n",
       "                      [ 0.1765,  0.0344,  0.0039,  ...,  0.0053, -0.1623,  0.1159],\n",
       "                      [ 0.2689,  0.1711, -0.2473,  ..., -0.2058, -0.0761, -0.0618],\n",
       "                      [ 0.0676, -0.0790, -0.0011,  ...,  0.0876,  0.0322,  0.1930]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.2.value.weight',\n",
       "              tensor([[-0.0833, -0.0607, -0.0775,  ..., -0.0439,  0.1338,  0.0115],\n",
       "                      [ 0.0495,  0.1907, -0.1786,  ...,  0.1847,  0.0005,  0.0982],\n",
       "                      [ 0.0732,  0.0811, -0.1146,  ..., -0.0524,  0.0715, -0.0569],\n",
       "                      ...,\n",
       "                      [ 0.0177, -0.0816,  0.0320,  ...,  0.1524, -0.0521, -0.0903],\n",
       "                      [ 0.1459,  0.0298,  0.0017,  ..., -0.0098, -0.0721,  0.1791],\n",
       "                      [-0.1266, -0.1423,  0.0759,  ..., -0.0525, -0.1243, -0.0915]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.3.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.3.key.weight',\n",
       "              tensor([[-0.1490, -0.2082, -0.0772,  ...,  0.0901, -0.0822,  0.0960],\n",
       "                      [-0.1838,  0.0331,  0.0968,  ..., -0.0764,  0.1367, -0.1941],\n",
       "                      [ 0.2349,  0.1280, -0.0316,  ..., -0.2568,  0.0404,  0.1274],\n",
       "                      ...,\n",
       "                      [-0.0179, -0.1971,  0.0628,  ..., -0.1487, -0.0735,  0.0741],\n",
       "                      [ 0.0819,  0.1466,  0.1145,  ..., -0.0488,  0.2059,  0.0091],\n",
       "                      [-0.1575, -0.2478, -0.3046,  ...,  0.0602,  0.0917,  0.0672]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.3.query.weight',\n",
       "              tensor([[-0.0142,  0.0228,  0.0956,  ...,  0.1271,  0.1071,  0.1931],\n",
       "                      [ 0.0129, -0.2757,  0.1928,  ...,  0.0543,  0.1237,  0.1013],\n",
       "                      [ 0.2403,  0.1352,  0.0027,  ...,  0.2373, -0.0340,  0.0122],\n",
       "                      ...,\n",
       "                      [ 0.0359,  0.1292,  0.0745,  ..., -0.1361,  0.0694,  0.2179],\n",
       "                      [-0.0950,  0.0926,  0.0962,  ..., -0.1022,  0.1438, -0.1873],\n",
       "                      [ 0.1970,  0.1196, -0.0343,  ..., -0.0598, -0.3031,  0.1007]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.heads.3.value.weight',\n",
       "              tensor([[ 0.0165, -0.1578,  0.2088,  ...,  0.0864, -0.0243,  0.0582],\n",
       "                      [ 0.1433,  0.0197,  0.1145,  ..., -0.0600,  0.0405,  0.1053],\n",
       "                      [ 0.1140, -0.0061, -0.0054,  ...,  0.1435, -0.1144,  0.0554],\n",
       "                      ...,\n",
       "                      [ 0.1550,  0.0203, -0.0671,  ...,  0.2012, -0.0572, -0.0700],\n",
       "                      [ 0.0214, -0.1301,  0.0177,  ..., -0.0669,  0.1724,  0.1306],\n",
       "                      [ 0.1324, -0.0178,  0.1331,  ..., -0.0017, -0.0352, -0.0797]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.proj.weight',\n",
       "              tensor([[ 0.0714, -0.1320,  0.1000,  ...,  0.0914,  0.0760,  0.1181],\n",
       "                      [-0.1008, -0.0272, -0.0988,  ..., -0.0343,  0.0759, -0.0586],\n",
       "                      [-0.0547, -0.1214,  0.0150,  ..., -0.0033,  0.0125, -0.0189],\n",
       "                      ...,\n",
       "                      [-0.0435, -0.0533, -0.0893,  ..., -0.1987, -0.0130, -0.0138],\n",
       "                      [ 0.1825, -0.0006, -0.0504,  ...,  0.0295,  0.0285,  0.0633],\n",
       "                      [-0.0606, -0.0361, -0.0401,  ..., -0.0010,  0.0824,  0.0820]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.sa.proj.bias',\n",
       "              tensor([-0.0367, -0.0033,  0.0057, -0.0125, -0.0865, -0.0753, -0.1138,  0.1466,\n",
       "                      -0.0135, -0.0743,  0.0377, -0.0399,  0.0550,  0.1317, -0.0396,  0.0908,\n",
       "                       0.0406,  0.0606,  0.0762,  0.0581,  0.0145,  0.0584,  0.1028,  0.1090,\n",
       "                      -0.1011,  0.0925, -0.0664,  0.0976,  0.0490, -0.0971, -0.0845,  0.0310,\n",
       "                      -0.0340,  0.1137,  0.0509,  0.0479,  0.0526,  0.0709, -0.0420,  0.0636,\n",
       "                      -0.0090,  0.0404, -0.1097, -0.0917,  0.0745,  0.0214, -0.0476, -0.0056,\n",
       "                       0.0530, -0.0151,  0.0670,  0.1280, -0.0412, -0.0761, -0.0245,  0.0524,\n",
       "                       0.1080, -0.0764,  0.0415,  0.0168, -0.0177, -0.1004, -0.1477, -0.0447],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ffwd.net.0.weight',\n",
       "              tensor([[-0.0766, -0.0465, -0.1148,  ..., -0.0687,  0.0771,  0.2568],\n",
       "                      [-0.1803, -0.1442,  0.1089,  ..., -0.1083, -0.0466,  0.0709],\n",
       "                      [ 0.0276,  0.0804, -0.0057,  ..., -0.0274, -0.0059, -0.0830],\n",
       "                      ...,\n",
       "                      [ 0.1320, -0.0495,  0.0064,  ..., -0.0254, -0.1579, -0.1528],\n",
       "                      [ 0.0936,  0.0445, -0.1156,  ..., -0.2031, -0.0477,  0.1299],\n",
       "                      [-0.0341,  0.1116, -0.1065,  ..., -0.0468, -0.0057, -0.0060]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ffwd.net.0.bias',\n",
       "              tensor([-0.0587, -0.0489, -0.1368, -0.0105, -0.1426, -0.2355,  0.0231, -0.1780,\n",
       "                      -0.1630, -0.1642, -0.1758, -0.1433, -0.1530, -0.0870, -0.1839, -0.0776,\n",
       "                      -0.1691, -0.0174, -0.0479, -0.0665,  0.0064, -0.0812, -0.1910, -0.1554,\n",
       "                      -0.2214,  0.0247, -0.1890, -0.1386, -0.0232, -0.0243, -0.0480, -0.2360,\n",
       "                      -0.0844, -0.1140, -0.0096,  0.0256, -0.1429, -0.2059, -0.2268, -0.1016,\n",
       "                      -0.1578, -0.0543,  0.0399, -0.1123, -0.1824, -0.1347, -0.1217, -0.2048,\n",
       "                       0.0302, -0.0583, -0.1931, -0.0168,  0.0630, -0.0630, -0.1321, -0.0120,\n",
       "                      -0.0974, -0.1657, -0.1265, -0.1069, -0.1196, -0.1598, -0.1703, -0.1073,\n",
       "                      -0.0578, -0.1376,  0.0459, -0.1615, -0.1305, -0.0994, -0.1102, -0.1146,\n",
       "                      -0.0676, -0.1026, -0.0772, -0.1512, -0.0039,  0.0593, -0.1096, -0.2446,\n",
       "                      -0.1772, -0.2111, -0.1630, -0.0906,  0.0549,  0.0011,  0.0115, -0.1117,\n",
       "                      -0.1413, -0.0998, -0.1089, -0.0206, -0.0765, -0.1790,  0.0364, -0.0148,\n",
       "                      -0.0270, -0.2161, -0.0405, -0.0205,  0.0011, -0.1075,  0.0027, -0.1309,\n",
       "                      -0.0110, -0.2218, -0.0837, -0.1256, -0.1783, -0.0786, -0.0952, -0.0050,\n",
       "                      -0.1162, -0.0273, -0.1441, -0.1195, -0.1904, -0.1848, -0.1145,  0.0756,\n",
       "                      -0.1437, -0.0111, -0.0857,  0.0068,  0.0137, -0.2165, -0.1281, -0.0726,\n",
       "                      -0.1672, -0.0842, -0.2043, -0.1922,  0.0047, -0.0259, -0.1525, -0.1435,\n",
       "                      -0.0573, -0.0254, -0.0048, -0.0521, -0.1858, -0.1345,  0.0315, -0.0663,\n",
       "                      -0.0115, -0.0896, -0.1041, -0.1768, -0.0325, -0.0152, -0.0041, -0.1382,\n",
       "                       0.0081, -0.1166, -0.0576, -0.1510, -0.1206, -0.0400, -0.0735, -0.0879,\n",
       "                      -0.0433, -0.1252, -0.0848,  0.0324,  0.0259, -0.1640,  0.0272, -0.0858,\n",
       "                      -0.1195, -0.1963, -0.0185, -0.0106,  0.0364, -0.0739, -0.1981, -0.1863,\n",
       "                      -0.1447, -0.0790, -0.0376,  0.0308, -0.1406, -0.0587, -0.0494, -0.0345,\n",
       "                      -0.1204,  0.0384, -0.0731, -0.1375, -0.2033, -0.0594,  0.0251, -0.1486,\n",
       "                      -0.0833, -0.1566, -0.1295,  0.0443, -0.0058, -0.0890, -0.0122, -0.1625,\n",
       "                      -0.0753, -0.1296, -0.2200, -0.0936, -0.0426, -0.1325,  0.0053, -0.1935,\n",
       "                      -0.0522, -0.1859, -0.0312, -0.1118, -0.1494, -0.0358, -0.0275, -0.1017,\n",
       "                      -0.0861, -0.1206,  0.0145, -0.0554,  0.0024, -0.1440, -0.1096,  0.0357,\n",
       "                      -0.1390, -0.0484, -0.1075, -0.0351, -0.1874,  0.0436,  0.0098, -0.2128,\n",
       "                      -0.1278, -0.0461,  0.0730, -0.0444, -0.1890, -0.1834, -0.1026, -0.1221,\n",
       "                      -0.0151, -0.0655, -0.0364,  0.0063,  0.0067, -0.2338, -0.1674, -0.0443,\n",
       "                      -0.0416, -0.0041, -0.1847, -0.1919, -0.0795, -0.1065, -0.1719, -0.0900],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ffwd.net.2.weight',\n",
       "              tensor([[-0.0552,  0.0531, -0.0770,  ..., -0.1029,  0.0502, -0.0513],\n",
       "                      [-0.0419,  0.0562,  0.0367,  ...,  0.1373,  0.0564,  0.1458],\n",
       "                      [-0.0643, -0.0093, -0.1594,  ...,  0.0498, -0.0474, -0.0415],\n",
       "                      ...,\n",
       "                      [ 0.1934,  0.0377, -0.0058,  ...,  0.0870, -0.0191, -0.0092],\n",
       "                      [-0.0020,  0.0830,  0.1192,  ...,  0.0715, -0.0309, -0.0725],\n",
       "                      [-0.1137, -0.0512, -0.0861,  ...,  0.0084,  0.1226, -0.0111]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ffwd.net.2.bias',\n",
       "              tensor([-0.0425, -0.0470,  0.0328,  0.0329,  0.0210, -0.0159,  0.0829,  0.0439,\n",
       "                       0.0071, -0.0210, -0.0218, -0.0469, -0.0600,  0.0052,  0.0158,  0.0466,\n",
       "                      -0.0326, -0.0301, -0.0688, -0.0385, -0.0115,  0.0122,  0.0445, -0.0329,\n",
       "                       0.0851, -0.0168,  0.0060, -0.0279,  0.0440, -0.0731,  0.0296, -0.0544,\n",
       "                      -0.0498, -0.0438, -0.0576, -0.0287,  0.0127,  0.0577, -0.0149,  0.0507,\n",
       "                       0.0228, -0.0052,  0.0476,  0.0476,  0.0179, -0.0101,  0.0518, -0.0258,\n",
       "                       0.0597, -0.0071, -0.0504, -0.0368,  0.0321, -0.0038,  0.0712,  0.0427,\n",
       "                      -0.0122,  0.0190,  0.0336, -0.0139, -0.0638, -0.0264,  0.0437, -0.0319],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ln1.weight',\n",
       "              tensor([1.2611, 1.3188, 1.2932, 1.2838, 1.2651, 1.2243, 1.2187, 1.1806, 1.0494,\n",
       "                      1.2537, 1.1120, 1.1266, 1.2540, 1.1943, 1.3302, 1.2559, 1.1950, 1.1512,\n",
       "                      1.1415, 1.0731, 1.2206, 1.3139, 1.0765, 1.2082, 1.2015, 1.1085, 1.0972,\n",
       "                      1.1152, 1.2602, 1.2106, 1.2353, 1.0890, 1.1502, 1.3054, 1.2005, 1.1665,\n",
       "                      1.3033, 1.2505, 1.2609, 1.1378, 1.3073, 1.0909, 1.1583, 1.0717, 1.3007,\n",
       "                      1.1300, 1.1773, 1.0257, 1.1934, 1.1505, 1.2211, 1.1381, 1.3257, 1.2357,\n",
       "                      1.1296, 1.2887, 1.1351, 1.1911, 1.2992, 1.2574, 1.0837, 1.1153, 1.1898,\n",
       "                      1.1974], device='cuda:0')),\n",
       "             ('blocks.1.ln1.bias',\n",
       "              tensor([ 0.0424, -0.0075,  0.0194,  0.1393, -0.0092,  0.0189,  0.0113,  0.0414,\n",
       "                      -0.0701, -0.0846, -0.0416, -0.0125,  0.0357, -0.0048,  0.0113, -0.0132,\n",
       "                      -0.1383,  0.0080,  0.0998, -0.0884, -0.0495, -0.0754, -0.0023, -0.0827,\n",
       "                      -0.0578,  0.0957, -0.0073, -0.0374, -0.0698, -0.0204,  0.0101, -0.1048,\n",
       "                      -0.0024,  0.1012, -0.0230,  0.0223, -0.0163,  0.0172,  0.0291,  0.0634,\n",
       "                      -0.1788, -0.0267,  0.0676,  0.0431, -0.0107,  0.0421,  0.1146,  0.0370,\n",
       "                       0.0575,  0.0018,  0.1047,  0.0545, -0.0739, -0.0358,  0.0636,  0.0223,\n",
       "                      -0.0458, -0.0617, -0.0923,  0.0568, -0.0307,  0.0165,  0.0830,  0.0318],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.1.ln2.weight',\n",
       "              tensor([0.9762, 0.9801, 1.0320, 0.9519, 0.9621, 1.0271, 0.9674, 1.1371, 1.1194,\n",
       "                      0.9745, 1.0640, 1.0467, 0.9851, 1.0247, 1.0623, 1.1241, 0.9393, 1.0205,\n",
       "                      0.9604, 0.9755, 0.9248, 1.0055, 0.9792, 1.0767, 1.1936, 1.0413, 1.1333,\n",
       "                      1.0222, 0.9332, 1.0776, 1.0293, 1.1236, 1.0054, 1.0335, 1.0899, 1.0175,\n",
       "                      1.0835, 0.9317, 1.0801, 1.0502, 0.9009, 1.0146, 1.0538, 1.1073, 0.9542,\n",
       "                      1.0024, 1.0724, 0.9306, 1.0870, 0.9126, 1.0133, 1.0687, 0.9532, 1.0681,\n",
       "                      0.9160, 1.0727, 1.0485, 1.0312, 1.0101, 0.9947, 1.0713, 0.9790, 1.0149,\n",
       "                      1.0622], device='cuda:0')),\n",
       "             ('blocks.1.ln2.bias',\n",
       "              tensor([-0.0408,  0.0044,  0.0180,  0.0236, -0.1089, -0.0136, -0.1028,  0.0558,\n",
       "                      -0.1189,  0.0241, -0.0973, -0.0067, -0.0006,  0.0452, -0.0651,  0.1265,\n",
       "                      -0.0199, -0.0655,  0.0549, -0.0405, -0.0665, -0.0743, -0.0187,  0.1054,\n",
       "                      -0.1196,  0.0484,  0.0476,  0.0206, -0.0144,  0.0122,  0.0721,  0.1444,\n",
       "                       0.0652,  0.0059,  0.0418,  0.0693, -0.0472,  0.0012,  0.0148,  0.0125,\n",
       "                       0.0585,  0.0238, -0.0249, -0.1345,  0.0289,  0.0453,  0.1273,  0.0692,\n",
       "                       0.1386, -0.0406, -0.0275, -0.0033, -0.0702,  0.0882,  0.0867, -0.0246,\n",
       "                      -0.0284, -0.1098, -0.0775, -0.0699, -0.1585, -0.0334, -0.0653,  0.0004],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.0.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.0.key.weight',\n",
       "              tensor([[-0.0810, -0.2115, -0.1302,  ...,  0.2077, -0.1438,  0.0227],\n",
       "                      [-0.0847, -0.0799,  0.0015,  ...,  0.3001, -0.1604, -0.0009],\n",
       "                      [-0.0944,  0.0915,  0.1272,  ..., -0.0767,  0.0918, -0.0504],\n",
       "                      ...,\n",
       "                      [-0.0997, -0.1079, -0.0471,  ...,  0.0857,  0.0194,  0.0393],\n",
       "                      [ 0.0389, -0.1677, -0.1818,  ..., -0.2372, -0.1313, -0.0254],\n",
       "                      [ 0.1166, -0.0424, -0.3600,  ...,  0.0959, -0.1574, -0.1675]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.0.query.weight',\n",
       "              tensor([[ 0.1253,  0.0582, -0.1786,  ..., -0.0063, -0.0846,  0.1658],\n",
       "                      [-0.0215,  0.0648, -0.1500,  ..., -0.2182, -0.1152,  0.1588],\n",
       "                      [-0.1561,  0.0829,  0.2339,  ...,  0.1650,  0.3778, -0.1111],\n",
       "                      ...,\n",
       "                      [ 0.1418,  0.0674,  0.0316,  ..., -0.0421, -0.1196, -0.0088],\n",
       "                      [ 0.0841,  0.0381,  0.0925,  ...,  0.1281,  0.0502,  0.1411],\n",
       "                      [-0.0908, -0.0762, -0.0454,  ...,  0.1547,  0.2558,  0.1433]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.0.value.weight',\n",
       "              tensor([[ 0.0908,  0.0770, -0.0014,  ...,  0.0027, -0.0870,  0.0668],\n",
       "                      [ 0.0340, -0.1661, -0.1300,  ..., -0.0272, -0.1287, -0.0692],\n",
       "                      [ 0.0293,  0.0791,  0.0802,  ..., -0.1887,  0.0107, -0.1039],\n",
       "                      ...,\n",
       "                      [-0.0340, -0.0285,  0.1063,  ...,  0.0386, -0.1625,  0.1190],\n",
       "                      [-0.0555, -0.0753,  0.1532,  ...,  0.1471, -0.0662,  0.0994],\n",
       "                      [-0.0472,  0.0849, -0.0620,  ..., -0.0066, -0.0429,  0.1160]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.1.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.1.key.weight',\n",
       "              tensor([[ 0.0012,  0.1115, -0.0783,  ...,  0.0857, -0.0637, -0.1251],\n",
       "                      [-0.0079, -0.2981, -0.1203,  ...,  0.3986, -0.0252,  0.1974],\n",
       "                      [ 0.0451, -0.2559, -0.1964,  ...,  0.3849, -0.2874,  0.0787],\n",
       "                      ...,\n",
       "                      [-0.1391, -0.2333,  0.0903,  ...,  0.0606, -0.3399,  0.1219],\n",
       "                      [-0.0850, -0.1398, -0.1370,  ...,  0.2255, -0.1462,  0.0859],\n",
       "                      [ 0.0105,  0.0790,  0.1232,  ..., -0.2681,  0.1007,  0.0059]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.1.query.weight',\n",
       "              tensor([[ 0.3008,  0.1189,  0.0177,  ..., -0.1158, -0.0945, -0.0595],\n",
       "                      [-0.0301,  0.2334, -0.3027,  ..., -0.0423,  0.1060, -0.0268],\n",
       "                      [ 0.2198, -0.1260, -0.0697,  ...,  0.1049,  0.2093, -0.0072],\n",
       "                      ...,\n",
       "                      [ 0.0052, -0.2455, -0.0627,  ...,  0.1117, -0.1198,  0.0503],\n",
       "                      [ 0.4514,  0.2569, -0.1841,  ..., -0.2646, -0.0582,  0.1076],\n",
       "                      [-0.1824,  0.0483, -0.0033,  ..., -0.0143,  0.0953,  0.0220]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.1.value.weight',\n",
       "              tensor([[-0.1245, -0.1304,  0.0135,  ...,  0.0355, -0.1565, -0.0166],\n",
       "                      [ 0.1050, -0.0231,  0.1112,  ..., -0.0667, -0.0206,  0.0323],\n",
       "                      [-0.0662,  0.0534, -0.1600,  ..., -0.0329,  0.0706,  0.1108],\n",
       "                      ...,\n",
       "                      [-0.0149, -0.0648, -0.0591,  ..., -0.0386,  0.1394, -0.0027],\n",
       "                      [-0.0040,  0.0763, -0.0214,  ..., -0.0874,  0.0773, -0.0389],\n",
       "                      [ 0.0809,  0.0351, -0.0681,  ..., -0.0938,  0.1703, -0.0412]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.2.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.2.key.weight',\n",
       "              tensor([[-7.7507e-02,  2.9503e-01,  1.3299e-01,  ..., -6.1062e-02,\n",
       "                        2.1278e-01,  6.7423e-02],\n",
       "                      [-1.2501e-01,  1.1058e-01,  1.9408e-01,  ..., -3.2195e-01,\n",
       "                        9.6450e-02, -1.8270e-01],\n",
       "                      [ 9.4133e-05, -1.7241e-01,  8.3144e-02,  ...,  1.5734e-02,\n",
       "                       -1.4223e-01,  8.6740e-02],\n",
       "                      ...,\n",
       "                      [ 3.3517e-02,  2.0830e-01, -2.0627e-01,  ..., -2.5876e-01,\n",
       "                        8.3850e-02,  2.2444e-03],\n",
       "                      [-2.1541e-01, -2.2443e-01, -7.9787e-02,  ...,  1.4610e-01,\n",
       "                       -2.6951e-01, -9.4777e-03],\n",
       "                      [-3.3618e-02, -2.0146e-01, -9.5574e-02,  ..., -9.5880e-02,\n",
       "                       -1.1649e-03,  1.6629e-01]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.2.query.weight',\n",
       "              tensor([[ 5.5331e-02,  8.3287e-02,  1.9282e-01,  ..., -5.0788e-02,\n",
       "                        2.5936e-02,  1.4586e-01],\n",
       "                      [-2.8450e-01, -3.1621e-01,  1.9098e-02,  ...,  3.5217e-01,\n",
       "                       -2.6093e-01, -1.1002e-01],\n",
       "                      [ 2.3282e-01,  1.0336e-01,  4.0855e-02,  ..., -6.7764e-02,\n",
       "                        2.8833e-01, -1.0864e-02],\n",
       "                      ...,\n",
       "                      [-2.2553e-01, -8.8827e-02,  2.0651e-01,  ...,  1.9527e-01,\n",
       "                       -3.1823e-04, -5.1654e-02],\n",
       "                      [ 1.3400e-01, -1.5126e-01,  2.0332e-02,  ..., -2.5365e-02,\n",
       "                       -1.7831e-01, -8.2060e-02],\n",
       "                      [ 1.9495e-01,  1.0886e-01, -3.2694e-01,  ...,  2.2754e-01,\n",
       "                        1.0118e-01, -9.1683e-02]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.2.value.weight',\n",
       "              tensor([[ 0.0192, -0.1118,  0.0683,  ...,  0.0681, -0.1069,  0.0454],\n",
       "                      [ 0.1591, -0.0327,  0.0010,  ...,  0.0434, -0.1913, -0.1765],\n",
       "                      [-0.1370,  0.1128, -0.0831,  ...,  0.0007,  0.0099, -0.1855],\n",
       "                      ...,\n",
       "                      [ 0.1530,  0.1768, -0.0961,  ..., -0.0464,  0.0580,  0.0323],\n",
       "                      [-0.0017,  0.0059, -0.3043,  ...,  0.1534, -0.0378, -0.2256],\n",
       "                      [-0.2205, -0.0857,  0.0140,  ..., -0.0557,  0.1145,  0.0431]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.3.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.3.key.weight',\n",
       "              tensor([[-0.0749,  0.0511, -0.1088,  ..., -0.2153,  0.0451,  0.0872],\n",
       "                      [-0.1100,  0.2902,  0.1384,  ..., -0.1867,  0.3507,  0.0139],\n",
       "                      [ 0.0101,  0.2876,  0.0636,  ..., -0.1957,  0.0546,  0.1506],\n",
       "                      ...,\n",
       "                      [ 0.1878, -0.1629, -0.0210,  ...,  0.3985, -0.1497,  0.1768],\n",
       "                      [ 0.1036,  0.2370,  0.1606,  ..., -0.1180,  0.2949, -0.1143],\n",
       "                      [-0.0670,  0.1044,  0.2142,  ..., -0.2529,  0.1157, -0.0640]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.3.query.weight',\n",
       "              tensor([[-0.0877,  0.0706,  0.0664,  ..., -0.0671,  0.1301,  0.2359],\n",
       "                      [-0.0022, -0.0978,  0.1917,  ..., -0.0559,  0.0595, -0.0407],\n",
       "                      [-0.0173,  0.0587,  0.0298,  ..., -0.0404,  0.1942, -0.1536],\n",
       "                      ...,\n",
       "                      [ 0.1461,  0.1863, -0.2036,  ..., -0.0571, -0.1345,  0.0769],\n",
       "                      [-0.1067, -0.0652,  0.1577,  ...,  0.2222,  0.1060, -0.2366],\n",
       "                      [-0.2186, -0.3047,  0.1496,  ...,  0.3303,  0.0858,  0.1838]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.heads.3.value.weight',\n",
       "              tensor([[-0.0513, -0.1124,  0.0264,  ...,  0.1971, -0.0690, -0.0698],\n",
       "                      [ 0.0480,  0.0611,  0.0828,  ...,  0.0701,  0.0316,  0.0334],\n",
       "                      [ 0.2545, -0.1576,  0.0764,  ..., -0.0254,  0.0668,  0.1575],\n",
       "                      ...,\n",
       "                      [ 0.0215,  0.1592, -0.0232,  ...,  0.1470, -0.0393,  0.0213],\n",
       "                      [-0.0707, -0.2269,  0.0601,  ...,  0.1091,  0.0838,  0.0824],\n",
       "                      [-0.0816, -0.1039,  0.0172,  ...,  0.0403,  0.1813, -0.0526]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.proj.weight',\n",
       "              tensor([[-0.0832, -0.0440,  0.2091,  ..., -0.0097,  0.0193,  0.0058],\n",
       "                      [ 0.0790, -0.1265,  0.1553,  ...,  0.0400,  0.0134, -0.0917],\n",
       "                      [-0.1647,  0.1960, -0.1186,  ...,  0.1216,  0.1895, -0.2842],\n",
       "                      ...,\n",
       "                      [-0.0312,  0.0463,  0.0806,  ...,  0.0199,  0.0148,  0.1494],\n",
       "                      [ 0.1528,  0.0618,  0.0059,  ..., -0.2203,  0.1021, -0.0840],\n",
       "                      [ 0.0048, -0.1277, -0.0211,  ...,  0.0139, -0.0131, -0.0105]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.sa.proj.bias',\n",
       "              tensor([ 0.0904, -0.0589, -0.0143,  0.0480, -0.0413, -0.1116, -0.0604, -0.0599,\n",
       "                       0.1006,  0.0662, -0.0201,  0.0319,  0.1011,  0.0159,  0.0237,  0.1033,\n",
       "                      -0.0381,  0.0164, -0.0607,  0.0772, -0.0007, -0.0205, -0.1218,  0.0967,\n",
       "                      -0.0034,  0.0928, -0.0668, -0.0461, -0.0021,  0.0379, -0.0256,  0.0013,\n",
       "                      -0.0890, -0.0779,  0.0187,  0.0818, -0.1090, -0.0988, -0.1324, -0.0502,\n",
       "                      -0.0795, -0.0320,  0.0174,  0.0288,  0.0959, -0.0167,  0.0469,  0.0293,\n",
       "                       0.0278, -0.0922,  0.0978, -0.0294, -0.0942,  0.0025,  0.0764, -0.1271,\n",
       "                       0.1260, -0.0887,  0.1120,  0.0113,  0.0710, -0.0552,  0.0414, -0.0613],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ffwd.net.0.weight',\n",
       "              tensor([[ 0.1260, -0.0849, -0.0262,  ...,  0.0059,  0.1161, -0.0428],\n",
       "                      [ 0.1343,  0.0254, -0.2200,  ...,  0.1067,  0.0153, -0.0756],\n",
       "                      [-0.1142, -0.0589, -0.1784,  ...,  0.1828,  0.1872, -0.0563],\n",
       "                      ...,\n",
       "                      [-0.0234,  0.0861,  0.0099,  ..., -0.1834, -0.0042,  0.0843],\n",
       "                      [ 0.2538, -0.1464, -0.0126,  ..., -0.0010,  0.1764,  0.1605],\n",
       "                      [ 0.0130,  0.1022,  0.0445,  ...,  0.0939, -0.0525, -0.0061]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ffwd.net.0.bias',\n",
       "              tensor([ 0.0702, -0.1803, -0.1245,  0.0514, -0.0894, -0.2476, -0.0324,  0.0092,\n",
       "                      -0.0893, -0.0222,  0.0221, -0.0928, -0.2630, -0.0059, -0.1581, -0.0296,\n",
       "                      -0.0713, -0.1318, -0.1040, -0.0800, -0.0995, -0.0009, -0.1454,  0.0133,\n",
       "                      -0.1009, -0.1743, -0.1423,  0.0147, -0.0653, -0.1220, -0.1002, -0.0789,\n",
       "                      -0.0185, -0.0279, -0.1300,  0.0013, -0.0106, -0.1569, -0.1405, -0.0762,\n",
       "                      -0.1194, -0.2536, -0.1740, -0.0257, -0.1138, -0.1447, -0.0671, -0.0525,\n",
       "                      -0.1317, -0.1624, -0.1207, -0.1855, -0.1650,  0.0425, -0.0774, -0.1613,\n",
       "                      -0.0966,  0.0026, -0.2385, -0.1591, -0.2071, -0.1837, -0.0299, -0.1708,\n",
       "                      -0.0852, -0.1348, -0.0787, -0.2257, -0.0800, -0.0667, -0.1030,  0.0350,\n",
       "                       0.0270, -0.0162, -0.0956, -0.1382, -0.1551, -0.1155, -0.1525, -0.0505,\n",
       "                      -0.1345,  0.0244, -0.1661, -0.1836, -0.0087, -0.1631, -0.1432, -0.1573,\n",
       "                      -0.1260, -0.0170, -0.1761, -0.0352, -0.0257, -0.0081, -0.1854,  0.0044,\n",
       "                      -0.1190, -0.1250,  0.0077, -0.0673, -0.1307, -0.1182, -0.1405, -0.1725,\n",
       "                      -0.0313, -0.0477, -0.1895, -0.0826, -0.0401,  0.0149, -0.0362, -0.1938,\n",
       "                      -0.0738, -0.0551,  0.0892, -0.1620, -0.1556, -0.1498, -0.0569, -0.1139,\n",
       "                       0.0527, -0.1190, -0.0149, -0.0395, -0.0880, -0.0174, -0.1099, -0.0491,\n",
       "                      -0.1330, -0.2151, -0.1350, -0.0587,  0.0092, -0.0372, -0.1202, -0.1553,\n",
       "                      -0.0206, -0.0103, -0.0714, -0.1122, -0.1144,  0.0456, -0.1224, -0.0174,\n",
       "                      -0.1549, -0.2440, -0.1349,  0.0098, -0.0711, -0.0755, -0.0896, -0.0453,\n",
       "                       0.0103,  0.0273, -0.0205, -0.0382, -0.1521, -0.1815, -0.0724, -0.1294,\n",
       "                      -0.0713, -0.1365, -0.0087, -0.1735, -0.1171, -0.1612, -0.1489, -0.0494,\n",
       "                      -0.0706, -0.1141, -0.1888, -0.2226, -0.0578, -0.1434, -0.0835, -0.1530,\n",
       "                      -0.1527, -0.1308, -0.0584, -0.2323, -0.1532, -0.0974, -0.0528, -0.1763,\n",
       "                      -0.0723, -0.0835, -0.0909, -0.2513, -0.0895, -0.1544, -0.1427, -0.1425,\n",
       "                      -0.0452, -0.1066, -0.3200, -0.1431, -0.0654, -0.1154,  0.0220, -0.0602,\n",
       "                      -0.2125, -0.1473, -0.1052, -0.0598,  0.0466, -0.1394,  0.0857, -0.1182,\n",
       "                      -0.1853,  0.0028, -0.2292, -0.1610, -0.0151, -0.1185, -0.0866, -0.1158,\n",
       "                      -0.1704, -0.0153, -0.0633, -0.0346,  0.0138, -0.0224, -0.1506, -0.1492,\n",
       "                      -0.1312, -0.0467, -0.1071, -0.0682, -0.1059, -0.2091, -0.1225, -0.0263,\n",
       "                      -0.0655, -0.2126, -0.1291, -0.1300, -0.1732, -0.0836, -0.2435, -0.1665,\n",
       "                       0.0359, -0.0585, -0.0817, -0.0287, -0.0051,  0.0332, -0.1949, -0.0796,\n",
       "                      -0.0788, -0.1401, -0.0412, -0.1465, -0.0456, -0.1867, -0.0485, -0.0580],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ffwd.net.2.weight',\n",
       "              tensor([[-0.1772, -0.1787, -0.1082,  ..., -0.0801, -0.0338,  0.1858],\n",
       "                      [-0.0108, -0.1092,  0.1010,  ...,  0.0568,  0.0496, -0.0970],\n",
       "                      [ 0.0042,  0.1940,  0.0334,  ..., -0.0005,  0.0545,  0.0576],\n",
       "                      ...,\n",
       "                      [-0.0576, -0.0047, -0.0205,  ...,  0.0144,  0.1744, -0.0428],\n",
       "                      [-0.0697,  0.0075, -0.1041,  ...,  0.0684,  0.0816,  0.0601],\n",
       "                      [-0.0946, -0.0498,  0.1034,  ..., -0.0897, -0.0292, -0.0326]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ffwd.net.2.bias',\n",
       "              tensor([ 0.0228, -0.0410,  0.0089, -0.0435,  0.0410, -0.0428,  0.0380,  0.0686,\n",
       "                      -0.0065,  0.0339,  0.0387,  0.0376,  0.0186,  0.0218, -0.0144,  0.0288,\n",
       "                      -0.0336, -0.0788, -0.0058, -0.0025,  0.0214, -0.0097, -0.0357, -0.0663,\n",
       "                      -0.0069, -0.0250, -0.0604, -0.0128,  0.0523,  0.0165,  0.0298, -0.0231,\n",
       "                       0.0546,  0.0178,  0.0390, -0.0077,  0.0037,  0.0255,  0.0009, -0.0326,\n",
       "                      -0.0252,  0.0712,  0.0420,  0.0643, -0.0659,  0.0048, -0.0395, -0.0692,\n",
       "                      -0.0379, -0.0676, -0.0696, -0.0074, -0.0066,  0.0149,  0.0116, -0.0381,\n",
       "                      -0.0124,  0.0507, -0.0160, -0.0396,  0.0666,  0.0196,  0.0323,  0.0095],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ln1.weight',\n",
       "              tensor([1.0372, 1.2233, 1.0687, 1.2062, 1.0245, 1.0292, 1.0956, 1.1086, 1.0567,\n",
       "                      1.0429, 1.1318, 1.1538, 1.1212, 1.1505, 1.1198, 0.9973, 1.1862, 1.0944,\n",
       "                      1.1072, 1.1012, 1.0953, 1.1239, 0.9626, 1.1511, 1.1795, 1.0127, 1.0999,\n",
       "                      1.0261, 1.1295, 1.1125, 1.0717, 1.1124, 1.1909, 1.2037, 1.0697, 1.1878,\n",
       "                      1.1741, 1.0125, 1.0634, 1.1793, 1.1570, 1.0305, 1.0746, 1.0892, 1.2193,\n",
       "                      1.0500, 1.0985, 1.1768, 1.1565, 1.1511, 0.9701, 1.2370, 1.1587, 1.1345,\n",
       "                      1.0765, 1.1093, 1.1602, 1.1500, 1.1677, 1.0385, 1.0675, 1.1911, 1.2323,\n",
       "                      1.0384], device='cuda:0')),\n",
       "             ('blocks.2.ln1.bias',\n",
       "              tensor([-0.0253, -0.0831,  0.0558,  0.0759, -0.0237,  0.0620,  0.0392,  0.0083,\n",
       "                      -0.0770,  0.0557, -0.0248, -0.0020,  0.0088, -0.0055,  0.0438, -0.0314,\n",
       "                      -0.0577, -0.0130,  0.0442, -0.0662,  0.0259,  0.0395, -0.0516, -0.0281,\n",
       "                       0.0038,  0.0397,  0.0140, -0.0221, -0.0628, -0.0158, -0.0022, -0.1719,\n",
       "                      -0.0217,  0.0633, -0.0261,  0.0240, -0.0904,  0.0445, -0.0473, -0.0054,\n",
       "                      -0.0841,  0.0374,  0.0321,  0.0035,  0.0008, -0.0251,  0.0492,  0.1280,\n",
       "                       0.0956,  0.0716, -0.0873,  0.0177, -0.0345, -0.0448, -0.0181, -0.0186,\n",
       "                      -0.0557, -0.0445, -0.0574, -0.0578, -0.0268,  0.0332,  0.0131,  0.0273],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.2.ln2.weight',\n",
       "              tensor([1.1953, 1.0485, 1.0732, 1.0727, 1.2163, 1.1177, 1.1496, 1.2574, 1.0809,\n",
       "                      1.0756, 1.2818, 1.2283, 1.0860, 1.0478, 1.0901, 1.2612, 1.0939, 1.2410,\n",
       "                      1.0956, 1.2587, 0.8941, 1.1600, 0.9961, 1.1318, 1.2313, 1.0842, 1.2337,\n",
       "                      1.0697, 1.0231, 1.2260, 1.1296, 1.1205, 1.0694, 1.2474, 1.1891, 1.1976,\n",
       "                      1.2110, 1.0491, 1.1137, 1.1110, 0.9997, 1.1001, 1.0989, 1.1235, 1.2588,\n",
       "                      1.0690, 1.3009, 1.1399, 1.0936, 1.0117, 1.0552, 1.3180, 1.0641, 1.1581,\n",
       "                      1.2328, 1.2720, 1.1322, 1.1963, 1.1403, 1.1597, 1.1887, 1.0231, 1.1098,\n",
       "                      1.0714], device='cuda:0')),\n",
       "             ('blocks.2.ln2.bias',\n",
       "              tensor([-0.0786,  0.0134, -0.0240, -0.0381,  0.0112,  0.0046,  0.0066,  0.0259,\n",
       "                      -0.1035, -0.1401,  0.0718, -0.1026, -0.0076,  0.0766, -0.1077,  0.0775,\n",
       "                      -0.0621,  0.0858,  0.1035, -0.1367, -0.0937, -0.0112, -0.1137,  0.0582,\n",
       "                      -0.0282,  0.0031,  0.0327,  0.0413, -0.0028, -0.1033, -0.0110, -0.0367,\n",
       "                       0.0501,  0.0799, -0.1318,  0.0758, -0.0626,  0.0438,  0.0464,  0.0317,\n",
       "                      -0.0024, -0.0457,  0.0598,  0.0029,  0.1193,  0.1692,  0.0588,  0.2380,\n",
       "                       0.0599,  0.0616,  0.0561,  0.0233, -0.0549, -0.0738,  0.1683,  0.0300,\n",
       "                      -0.0260, -0.0782, -0.0455,  0.0367, -0.0916, -0.0177, -0.0322, -0.0307],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.0.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.0.key.weight',\n",
       "              tensor([[ 0.0379, -0.0629, -0.1288,  ..., -0.1445,  0.2355,  0.1290],\n",
       "                      [-0.2855, -0.2809, -0.1212,  ...,  0.4174, -0.2719,  0.2253],\n",
       "                      [ 0.0889, -0.1185, -0.0645,  ..., -0.0856, -0.1621,  0.0445],\n",
       "                      ...,\n",
       "                      [-0.1045,  0.1689,  0.0174,  ..., -0.1019,  0.0896, -0.0230],\n",
       "                      [-0.0168,  0.1689,  0.2001,  ...,  0.1091, -0.0745,  0.0876],\n",
       "                      [ 0.0591, -0.1723,  0.0315,  ..., -0.1059, -0.0804,  0.0444]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.0.query.weight',\n",
       "              tensor([[ 0.2222,  0.0849,  0.0161,  ...,  0.1600,  0.3662, -0.3451],\n",
       "                      [ 0.0863,  0.1945, -0.2245,  ..., -0.2182, -0.2355, -0.0706],\n",
       "                      [ 0.0828,  0.0507, -0.1652,  ..., -0.0998,  0.1920, -0.0519],\n",
       "                      ...,\n",
       "                      [ 0.0806,  0.0278,  0.1192,  ..., -0.0652,  0.1229,  0.1877],\n",
       "                      [-0.0863, -0.0891, -0.0699,  ..., -0.0503, -0.3451,  0.3445],\n",
       "                      [ 0.0559, -0.0523, -0.0820,  ...,  0.2666,  0.2269, -0.2787]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.0.value.weight',\n",
       "              tensor([[-0.1867, -0.1686, -0.0696,  ...,  0.0361, -0.0888,  0.0338],\n",
       "                      [-0.0266,  0.1691,  0.0633,  ...,  0.0418,  0.0749,  0.0865],\n",
       "                      [-0.0205, -0.1388,  0.0246,  ..., -0.1358, -0.0150, -0.0113],\n",
       "                      ...,\n",
       "                      [-0.0264,  0.1175, -0.0584,  ..., -0.0816,  0.1014,  0.0356],\n",
       "                      [ 0.1404, -0.0840, -0.0094,  ...,  0.0832, -0.2350, -0.1155],\n",
       "                      [ 0.1905,  0.0849,  0.0307,  ..., -0.0705, -0.1141, -0.0275]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.1.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.1.key.weight',\n",
       "              tensor([[-0.2265, -0.1472, -0.1220,  ...,  0.1513,  0.0716, -0.0524],\n",
       "                      [ 0.1131, -0.1861,  0.0360,  ..., -0.3803, -0.0354,  0.1105],\n",
       "                      [-0.1663,  0.0599,  0.1475,  ...,  0.2757,  0.0975,  0.0775],\n",
       "                      ...,\n",
       "                      [ 0.2028,  0.3603,  0.1859,  ..., -0.0966,  0.3782,  0.0487],\n",
       "                      [ 0.0671, -0.0537, -0.0252,  ...,  0.0927, -0.1739, -0.1286],\n",
       "                      [ 0.1891, -0.2222,  0.0132,  ..., -0.1850,  0.1175, -0.0628]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.1.query.weight',\n",
       "              tensor([[-0.0231,  0.3218,  0.1956,  ..., -0.1833, -0.0303, -0.3024],\n",
       "                      [ 0.0430, -0.1304,  0.3380,  ...,  0.2146,  0.1250,  0.2683],\n",
       "                      [-0.0009,  0.2114,  0.1286,  ...,  0.1729,  0.1281,  0.1899],\n",
       "                      ...,\n",
       "                      [-0.1125,  0.0385, -0.0054,  ..., -0.0104, -0.0485,  0.1198],\n",
       "                      [-0.0005, -0.1644, -0.1195,  ..., -0.0379,  0.0048,  0.1118],\n",
       "                      [-0.0791, -0.2204, -0.2981,  ...,  0.0407, -0.3026, -0.0859]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.1.value.weight',\n",
       "              tensor([[ 0.0156, -0.0876, -0.0952,  ...,  0.0492, -0.0414,  0.1774],\n",
       "                      [ 0.0055,  0.0219,  0.0793,  ..., -0.1615,  0.0485,  0.0650],\n",
       "                      [ 0.0013, -0.2420,  0.1402,  ..., -0.0831,  0.0754,  0.0617],\n",
       "                      ...,\n",
       "                      [ 0.0757, -0.1797,  0.0173,  ...,  0.0819,  0.0836,  0.0962],\n",
       "                      [ 0.0324,  0.0764,  0.0717,  ...,  0.0646, -0.1407,  0.2147],\n",
       "                      [-0.0126,  0.0527, -0.0963,  ..., -0.0025,  0.1259,  0.0835]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.2.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.2.key.weight',\n",
       "              tensor([[-0.0006, -0.0945,  0.0157,  ...,  0.1437, -0.0547,  0.1311],\n",
       "                      [-0.3057,  0.0243, -0.0712,  ...,  0.1136, -0.1066,  0.1019],\n",
       "                      [-0.0439,  0.0109, -0.1845,  ..., -0.2556, -0.1824, -0.0664],\n",
       "                      ...,\n",
       "                      [ 0.2224,  0.0519, -0.0050,  ..., -0.1060,  0.1277, -0.2265],\n",
       "                      [-0.1604, -0.0545,  0.2297,  ..., -0.0813,  0.2152,  0.0706],\n",
       "                      [-0.1335,  0.0256,  0.1671,  ...,  0.0355, -0.0158, -0.1272]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.2.query.weight',\n",
       "              tensor([[ 0.0644, -0.0465,  0.1524,  ..., -0.0591,  0.0819,  0.2289],\n",
       "                      [ 0.2119, -0.0304,  0.1108,  ..., -0.3626,  0.0354,  0.3282],\n",
       "                      [-0.0488, -0.0739,  0.2001,  ..., -0.1236,  0.0944,  0.2968],\n",
       "                      ...,\n",
       "                      [-0.1784,  0.0878, -0.0756,  ...,  0.1168,  0.0952, -0.2133],\n",
       "                      [ 0.0542,  0.0900,  0.1012,  ...,  0.0504, -0.0265,  0.0037],\n",
       "                      [ 0.2331, -0.0296,  0.0913,  ..., -0.2338,  0.0143,  0.0721]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.2.value.weight',\n",
       "              tensor([[-0.0912, -0.0228,  0.0547,  ..., -0.1297, -0.1185, -0.1369],\n",
       "                      [-0.0374,  0.0596, -0.0231,  ..., -0.1974,  0.0564, -0.1136],\n",
       "                      [-0.1757, -0.0609,  0.2125,  ..., -0.1013,  0.0472,  0.0676],\n",
       "                      ...,\n",
       "                      [-0.1106, -0.1468,  0.2184,  ...,  0.1026,  0.1424, -0.0136],\n",
       "                      [ 0.0636,  0.0501,  0.1402,  ...,  0.0440, -0.2297,  0.2058],\n",
       "                      [ 0.0685,  0.0417,  0.0628,  ..., -0.0383,  0.2676, -0.0215]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.3.tril',\n",
       "              tensor([[1., 0., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 0.,  ..., 0., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 0., 0., 0.],\n",
       "                      ...,\n",
       "                      [1., 1., 1.,  ..., 1., 0., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 0.],\n",
       "                      [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.3.key.weight',\n",
       "              tensor([[ 0.0505, -0.0117, -0.0679,  ...,  0.0538, -0.1498,  0.1256],\n",
       "                      [ 0.0569, -0.1057,  0.0310,  ...,  0.0849, -0.1263,  0.2060],\n",
       "                      [ 0.0373,  0.2919,  0.1226,  ..., -0.3367,  0.1457, -0.1795],\n",
       "                      ...,\n",
       "                      [ 0.0717,  0.1751,  0.2186,  ..., -0.1895,  0.0511, -0.0388],\n",
       "                      [-0.0422, -0.1639,  0.0685,  ...,  0.1179, -0.0934, -0.2133],\n",
       "                      [-0.1009, -0.4096,  0.0843,  ...,  0.0891, -0.1612, -0.0168]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.3.query.weight',\n",
       "              tensor([[ 1.7774e-01,  2.1403e-01,  1.0305e-02,  ..., -1.0296e-01,\n",
       "                        1.5821e-01, -9.0286e-02],\n",
       "                      [-1.8982e-01,  1.4882e-01, -4.5420e-02,  ..., -4.8680e-02,\n",
       "                        7.1875e-02,  1.4310e-01],\n",
       "                      [-1.5347e-01, -1.0192e-02,  7.9166e-02,  ...,  1.7360e-01,\n",
       "                        4.7306e-02,  6.5517e-03],\n",
       "                      ...,\n",
       "                      [-1.1189e-01, -3.3478e-02, -3.8914e-03,  ...,  1.2338e-01,\n",
       "                        1.3293e-01, -1.2833e-01],\n",
       "                      [-2.9355e-02, -1.6263e-01,  2.1297e-02,  ...,  3.1744e-02,\n",
       "                       -1.0745e-01,  8.8325e-02],\n",
       "                      [ 2.3233e-01, -1.6296e-02, -1.4091e-01,  ...,  7.8796e-05,\n",
       "                        1.4413e-01,  1.2077e-01]], device='cuda:0')),\n",
       "             ('blocks.3.sa.heads.3.value.weight',\n",
       "              tensor([[ 0.1673,  0.1056,  0.0069,  ..., -0.0205,  0.0867, -0.0607],\n",
       "                      [ 0.0687,  0.0762, -0.1712,  ..., -0.1236, -0.0310, -0.0432],\n",
       "                      [-0.0380,  0.0056,  0.0625,  ...,  0.0577, -0.0030, -0.0624],\n",
       "                      ...,\n",
       "                      [ 0.1431, -0.0775, -0.2202,  ..., -0.0659, -0.0495, -0.0679],\n",
       "                      [-0.0028,  0.1310, -0.0568,  ..., -0.1091, -0.0390, -0.0834],\n",
       "                      [ 0.0016,  0.0826,  0.1074,  ...,  0.0605,  0.0504, -0.0510]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.proj.weight',\n",
       "              tensor([[ 0.0732, -0.0939, -0.2625,  ...,  0.0611, -0.0513, -0.0819],\n",
       "                      [ 0.0079,  0.1194, -0.0200,  ...,  0.1833, -0.1527, -0.0651],\n",
       "                      [ 0.1225, -0.2888,  0.0077,  ...,  0.0854,  0.0985, -0.0161],\n",
       "                      ...,\n",
       "                      [-0.1446,  0.0060,  0.0296,  ...,  0.0757,  0.0545, -0.0190],\n",
       "                      [ 0.0240, -0.1016,  0.0311,  ...,  0.1272, -0.0499, -0.0150],\n",
       "                      [ 0.1532, -0.0631, -0.0721,  ...,  0.1067,  0.0983, -0.0989]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.sa.proj.bias',\n",
       "              tensor([ 0.0978, -0.0096,  0.0394, -0.0474, -0.0427,  0.0562,  0.0313,  0.0316,\n",
       "                       0.1221,  0.0773, -0.0835, -0.0084, -0.1097, -0.0839, -0.0949, -0.0081,\n",
       "                       0.1267, -0.1067,  0.0334,  0.0919,  0.0796, -0.0768,  0.1193, -0.1060,\n",
       "                      -0.0457,  0.0194, -0.1402,  0.0262,  0.0868, -0.1111,  0.0954, -0.0637,\n",
       "                       0.0367, -0.0453,  0.1210,  0.0982, -0.0049, -0.0717,  0.1021, -0.0796,\n",
       "                       0.0216,  0.0012,  0.0177, -0.0491, -0.0418, -0.0999,  0.1238, -0.0274,\n",
       "                      -0.0996,  0.0668,  0.0093,  0.0426, -0.0981, -0.0709, -0.0370,  0.0523,\n",
       "                       0.1242,  0.0406,  0.0046, -0.0306,  0.1026,  0.0335, -0.0331, -0.0773],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ffwd.net.0.weight',\n",
       "              tensor([[ 0.0718,  0.0482, -0.0997,  ..., -0.0120,  0.0480,  0.1347],\n",
       "                      [ 0.1990, -0.0748, -0.0430,  ..., -0.2071, -0.1286,  0.0160],\n",
       "                      [-0.0014, -0.1822, -0.1099,  ..., -0.3726, -0.0161, -0.0164],\n",
       "                      ...,\n",
       "                      [ 0.2143,  0.0643, -0.1580,  ..., -0.0686,  0.1399, -0.0622],\n",
       "                      [-0.0504, -0.0567,  0.0323,  ...,  0.1488, -0.1229,  0.1394],\n",
       "                      [-0.0294,  0.0356,  0.1089,  ..., -0.1927,  0.1447,  0.0662]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ffwd.net.0.bias',\n",
       "              tensor([-0.0992,  0.0095, -0.1090, -0.1745, -0.0734,  0.0370, -0.0819, -0.0930,\n",
       "                      -0.0278, -0.1629,  0.0035, -0.0613, -0.0724, -0.2016, -0.0309, -0.0495,\n",
       "                      -0.0695, -0.0684, -0.0288, -0.0959, -0.1445, -0.1311, -0.0934, -0.1823,\n",
       "                      -0.1476, -0.0147, -0.2079,  0.0443, -0.1245, -0.0143, -0.1325, -0.1121,\n",
       "                      -0.0966, -0.1484, -0.0632, -0.1997, -0.0222, -0.0373, -0.0471, -0.2149,\n",
       "                      -0.2396, -0.1711, -0.1213, -0.1121, -0.0295, -0.1708, -0.1347, -0.0363,\n",
       "                      -0.0642, -0.1370, -0.1730, -0.0884, -0.0351, -0.1231, -0.2107, -0.0783,\n",
       "                      -0.1770, -0.1627, -0.1996,  0.0784, -0.1354, -0.0022, -0.0698, -0.2401,\n",
       "                       0.0085, -0.0144, -0.1683, -0.0217, -0.1030, -0.0805, -0.0436, -0.0906,\n",
       "                      -0.0568, -0.1409, -0.0434, -0.1671, -0.0728, -0.0253, -0.1086, -0.0123,\n",
       "                      -0.0085,  0.0126, -0.0707, -0.0862, -0.1300, -0.0717, -0.1468, -0.0410,\n",
       "                      -0.2335, -0.1300, -0.0332, -0.1230, -0.0162, -0.0552,  0.0028,  0.0037,\n",
       "                      -0.1441, -0.1923, -0.0494, -0.1406, -0.1105, -0.0935, -0.0535, -0.1341,\n",
       "                      -0.0023,  0.0159, -0.1879, -0.1139, -0.1145, -0.1329, -0.0078, -0.0060,\n",
       "                      -0.1742, -0.0319, -0.1132, -0.1142, -0.1900, -0.1166, -0.1661, -0.2296,\n",
       "                      -0.1221, -0.1590, -0.0959, -0.0882, -0.0935, -0.0191, -0.1325, -0.0561,\n",
       "                      -0.2605, -0.0420,  0.0448, -0.0471, -0.1227, -0.1286, -0.1398, -0.1267,\n",
       "                       0.0190, -0.0974,  0.0199,  0.0607, -0.1326, -0.1055, -0.0183, -0.0110,\n",
       "                      -0.0688, -0.2180,  0.0191, -0.0784, -0.0874, -0.0375, -0.0898, -0.0019,\n",
       "                       0.0165, -0.1809,  0.0544, -0.0666, -0.1360, -0.0999, -0.2604, -0.1033,\n",
       "                      -0.1147, -0.1034, -0.1970, -0.1495, -0.1066, -0.0954, -0.1507, -0.1239,\n",
       "                      -0.1711,  0.0888, -0.1193, -0.1771,  0.0014, -0.1203, -0.1196, -0.0114,\n",
       "                       0.0196,  0.0291, -0.1673, -0.0233, -0.2494, -0.0530,  0.0030, -0.1604,\n",
       "                      -0.2111, -0.0148, -0.1750, -0.1000, -0.0584, -0.1413, -0.1373,  0.0104,\n",
       "                      -0.1746, -0.1655, -0.1656, -0.0332, -0.0844,  0.0872, -0.1583, -0.1678,\n",
       "                       0.0180,  0.0161, -0.0603, -0.0756, -0.0224, -0.0892, -0.0753, -0.1614,\n",
       "                      -0.0503, -0.1296, -0.1099, -0.1789,  0.0312, -0.0576, -0.0111, -0.1956,\n",
       "                      -0.1901, -0.0623, -0.0075, -0.1496, -0.1553, -0.1079, -0.0378, -0.0393,\n",
       "                       0.0507, -0.1130, -0.1221, -0.0767, -0.0999,  0.0065, -0.1820, -0.1094,\n",
       "                       0.0041, -0.1293, -0.1903, -0.0704, -0.2076, -0.0464, -0.0430, -0.1024,\n",
       "                       0.0063,  0.0663, -0.0311, -0.0079, -0.1798,  0.0109, -0.1962, -0.1974,\n",
       "                       0.0177, -0.0842, -0.0404, -0.1086, -0.1497, -0.0807,  0.0593, -0.0409],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ffwd.net.2.weight',\n",
       "              tensor([[-0.0295,  0.1124, -0.3116,  ...,  0.0818,  0.0402, -0.1095],\n",
       "                      [ 0.0354,  0.0396,  0.1087,  ...,  0.0829,  0.0332, -0.0323],\n",
       "                      [ 0.0751, -0.0336, -0.0356,  ..., -0.1133,  0.1755,  0.0822],\n",
       "                      ...,\n",
       "                      [-0.0221,  0.0881, -0.1165,  ..., -0.1233, -0.0186,  0.0063],\n",
       "                      [-0.0352, -0.1340,  0.1021,  ...,  0.1442, -0.0906, -0.1814],\n",
       "                      [-0.0623, -0.0326,  0.1251,  ..., -0.0444,  0.1636,  0.0517]],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ffwd.net.2.bias',\n",
       "              tensor([ 0.0740, -0.0307,  0.0229,  0.0373, -0.0244, -0.0317,  0.0695, -0.0852,\n",
       "                       0.0056,  0.0226, -0.0354, -0.0015, -0.0148,  0.0415,  0.0411,  0.0301,\n",
       "                      -0.0602,  0.0603,  0.0237,  0.0044, -0.0013, -0.0079,  0.0144, -0.0311,\n",
       "                      -0.0404,  0.0357,  0.0391,  0.0523,  0.0235, -0.0021,  0.0492, -0.0163,\n",
       "                       0.0140, -0.0136,  0.0698,  0.0386,  0.0069, -0.0232,  0.0066, -0.0605,\n",
       "                       0.0494,  0.0007,  0.0066, -0.0147,  0.0438, -0.0186,  0.0172,  0.0342,\n",
       "                       0.0202, -0.0599,  0.0511,  0.0167, -0.0378,  0.0739, -0.0382, -0.0165,\n",
       "                       0.0176,  0.0127, -0.0252,  0.0246, -0.0144, -0.0159, -0.0187,  0.0159],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ln1.weight',\n",
       "              tensor([1.0607, 1.2759, 1.1699, 1.2172, 1.0283, 1.0113, 1.1277, 1.1862, 1.0251,\n",
       "                      1.0898, 1.2633, 1.3000, 1.1168, 1.0761, 1.1595, 1.1021, 1.1604, 1.1803,\n",
       "                      1.2221, 1.1897, 1.0340, 1.0955, 1.1022, 1.0588, 1.1297, 1.0191, 1.1627,\n",
       "                      1.1272, 1.0663, 1.0946, 1.0528, 1.1338, 1.3485, 1.2218, 1.1621, 1.0414,\n",
       "                      1.1214, 1.0355, 1.0568, 1.2014, 1.2004, 1.1547, 1.2662, 1.1805, 1.3556,\n",
       "                      1.1186, 1.1861, 1.1668, 1.0421, 1.0990, 1.0356, 1.2476, 1.2279, 1.0308,\n",
       "                      1.0374, 1.1719, 1.2051, 1.0481, 1.1549, 1.1132, 1.2004, 1.2193, 1.3090,\n",
       "                      1.1486], device='cuda:0')),\n",
       "             ('blocks.3.ln1.bias',\n",
       "              tensor([-0.0274,  0.0123,  0.0145,  0.1289, -0.0012,  0.0187, -0.0261, -0.0189,\n",
       "                      -0.0246,  0.0281, -0.0084, -0.0537,  0.0562,  0.0147,  0.0362,  0.0571,\n",
       "                      -0.0611,  0.0793,  0.1301, -0.0066,  0.0541,  0.0144, -0.0751, -0.0278,\n",
       "                      -0.0053, -0.0103,  0.0459, -0.0448, -0.0512,  0.0346, -0.0807, -0.0739,\n",
       "                      -0.0463,  0.0520, -0.0766, -0.0236,  0.0172, -0.0504, -0.0498,  0.1465,\n",
       "                      -0.1717,  0.1374, -0.0227, -0.0683, -0.0208,  0.0448,  0.0211,  0.0427,\n",
       "                       0.0639,  0.0533, -0.0876,  0.0078, -0.0364, -0.0649, -0.0562, -0.0193,\n",
       "                      -0.0725, -0.0425, -0.0234,  0.0194, -0.0201,  0.0099,  0.0329,  0.0086],\n",
       "                     device='cuda:0')),\n",
       "             ('blocks.3.ln2.weight',\n",
       "              tensor([1.2115, 1.3674, 1.1154, 1.1472, 1.2564, 1.2863, 1.2175, 1.2879, 1.0905,\n",
       "                      1.1243, 1.3189, 1.3282, 1.2016, 1.0860, 1.2430, 1.2502, 1.2162, 1.2689,\n",
       "                      1.1539, 1.2757, 1.0659, 1.2136, 1.2073, 1.1708, 1.3151, 1.1160, 1.3000,\n",
       "                      1.2472, 1.2421, 1.4328, 1.2546, 1.2630, 1.2687, 1.3370, 1.1648, 1.3066,\n",
       "                      1.2930, 1.0973, 1.3173, 1.3547, 1.2151, 1.1933, 1.1794, 1.2666, 1.2693,\n",
       "                      1.1914, 1.2716, 1.2033, 1.2434, 1.1121, 1.0871, 1.3921, 1.2299, 1.1486,\n",
       "                      1.2976, 1.2767, 1.2366, 1.3345, 1.2251, 1.3561, 1.2106, 1.2924, 1.3311,\n",
       "                      1.1849], device='cuda:0')),\n",
       "             ('blocks.3.ln2.bias',\n",
       "              tensor([-0.0359, -0.0293,  0.0456, -0.0270, -0.1713,  0.0552, -0.0700,  0.1333,\n",
       "                       0.0106,  0.0604,  0.0357, -0.0248, -0.0737,  0.0378,  0.0113,  0.0992,\n",
       "                       0.0829, -0.1365,  0.0438, -0.0266, -0.0425,  0.0615,  0.0519, -0.0338,\n",
       "                       0.1213, -0.0065, -0.1631, -0.0424,  0.0555, -0.2156, -0.0152, -0.0039,\n",
       "                       0.1114,  0.0617, -0.0370,  0.1840, -0.0644,  0.0736,  0.0522,  0.0179,\n",
       "                      -0.0254,  0.0646,  0.1408, -0.0486, -0.0008,  0.0337, -0.0112,  0.0519,\n",
       "                      -0.0122, -0.0153,  0.0147,  0.0864, -0.0040, -0.0191,  0.1114,  0.1352,\n",
       "                       0.0136, -0.1047, -0.1518,  0.0174, -0.0003,  0.1887,  0.0223, -0.1047],\n",
       "                     device='cuda:0')),\n",
       "             ('ln_f.weight',\n",
       "              tensor([1.2814, 1.5119, 1.5145, 1.4594, 1.4427, 1.4494, 1.3963, 1.2673, 1.2369,\n",
       "                      1.3564, 1.4890, 1.3581, 1.3494, 1.3837, 1.5634, 1.4232, 1.4715, 1.4054,\n",
       "                      1.3309, 1.4961, 1.2388, 1.4809, 1.3024, 1.4153, 1.5513, 1.4247, 1.4213,\n",
       "                      1.4209, 1.4427, 1.4158, 1.4472, 1.3591, 1.3392, 1.5951, 1.3218, 1.4343,\n",
       "                      1.4377, 1.4174, 1.4776, 1.4313, 1.4057, 1.1496, 1.4891, 1.6565, 1.4777,\n",
       "                      1.5320, 1.3242, 1.5235, 1.3763, 1.4272, 1.2660, 1.5473, 1.4636, 1.4224,\n",
       "                      1.3235, 1.4419, 1.4632, 1.3713, 1.4112, 1.4200, 1.4025, 1.3645, 1.2655,\n",
       "                      1.4068], device='cuda:0')),\n",
       "             ('ln_f.bias',\n",
       "              tensor([-0.1318, -0.0492,  0.1019,  0.2120,  0.0637,  0.1414,  0.0894,  0.1048,\n",
       "                      -0.1329, -0.0905, -0.1119, -0.0386,  0.1088,  0.1209, -0.0633,  0.1856,\n",
       "                       0.0247, -0.1377,  0.1959, -0.1392, -0.1006, -0.1063, -0.1599, -0.1813,\n",
       "                       0.0004,  0.0839, -0.0934, -0.0333, -0.1085, -0.0931,  0.0885, -0.2112,\n",
       "                      -0.1442,  0.1759, -0.1849,  0.0806, -0.1184,  0.1058,  0.0984, -0.0979,\n",
       "                      -0.2183,  0.0829,  0.1487,  0.0122,  0.0030,  0.0898,  0.1509, -0.0562,\n",
       "                      -0.0447,  0.1473,  0.0769,  0.0863, -0.1896, -0.1532,  0.0859,  0.2027,\n",
       "                      -0.1744,  0.0255, -0.1848, -0.2269, -0.1286,  0.1611,  0.0896, -0.0814],\n",
       "                     device='cuda:0')),\n",
       "             ('lm_head.weight',\n",
       "              tensor([[ 0.1969,  0.1499, -0.0244,  ..., -0.2392,  0.1871,  0.1653],\n",
       "                      [-0.1503,  0.0473,  0.1241,  ...,  0.1259,  0.0500,  0.2240],\n",
       "                      [ 0.0575,  0.1289,  0.0969,  ..., -0.0038, -0.0511,  0.1306],\n",
       "                      ...,\n",
       "                      [ 0.1791, -0.1425, -0.0724,  ..., -0.4661, -0.0413,  0.0975],\n",
       "                      [ 0.2097, -0.1515, -0.0187,  ...,  0.0696,  0.0105, -0.1703],\n",
       "                      [ 0.1654, -0.2608, -0.1704,  ..., -0.2462, -0.1195,  0.0672]],\n",
       "                     device='cuda:0')),\n",
       "             ('lm_head.bias',\n",
       "              tensor([ 0.0718, -0.0047, -0.0497, -0.2349, -0.1882,  0.0532, -0.0648, -0.2120,\n",
       "                      -0.0617, -0.3245,  0.0350, -0.0859, -0.0123, -0.1096, -0.1499, -0.0525,\n",
       "                      -0.1319,  0.0520, -0.1144, -0.0723, -0.2085,  0.0868, -0.1069, -0.0848,\n",
       "                       0.0176,  0.0499, -0.0007, -0.1027, -0.1431, -0.2198,  0.1205, -0.0720,\n",
       "                      -0.1283, -0.1668, -0.2376, -0.0419, -0.1732, -0.1163, -0.2580,  0.0628,\n",
       "                      -0.1374,  0.0255,  0.0307,  0.0334,  0.0573,  0.0937, -0.0725,  0.1922,\n",
       "                      -0.1502,  0.0218, -0.0360,  0.0432, -0.0495,  0.1474, -0.0363, -0.2259,\n",
       "                       0.0526,  0.1633,  0.0302,  0.1577, -0.0552, -0.1259, -0.2759, -0.0610,\n",
       "                      -0.1717], device='cuda:0'))])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c1dab6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pathlib\n",
    "from pathlib import Path\n",
    "\n",
    "MODEL_PATH=Path(\"models\")\n",
    "MODEL_PATH.mkdir(parents=True,exist_ok=True)\n",
    "MODEL_PATH_Name=\"gpt_1_shakespear.pth\"\n",
    "Model_save_path=\n",
    "torch.save(obj=model.state_dict(),f="
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}