![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ImSfwIQxXDmD",
"outputId": "d7b6c771-5fe8-4df4-9727-98373f3c92c1"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Mounted at /content/drive\n"
]
}
],
"source": [
"from google.colab import drive\n",
"drive.mount('/content/drive')"
]
},
{
"cell_type": "code",
"source": [
"import os\n",
"\n",
"# Directory where the dataset will be stored\n",
"dataset_dir = \"TinyStories4\"\n",
"\n",
"# Create the TinyStories directory if it doesn't exist\n",
"os.makedirs(dataset_dir, exist_ok=True)\n",
"\n",
"# URL of the dataset archive\n",
"dataset_url = \"https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories_all_data.tar.gz\"\n",
"# Name of the file to save the downloaded dataset archive\n",
"dataset_archive_path = os.path.join(dataset_dir, \"TinyStories_all_data.tar.gz\")\n",
"\n",
"# Download the archive\n",
"if not os.path.exists(dataset_archive_path):\n",
" os.system(f\"wget {dataset_url} -O {dataset_archive_path}\")\n",
"\n",
"# Extract the archive into the TinyStories directory\n",
"os.system(f\"tar -xzf {dataset_archive_path} -C {dataset_dir}\")\n"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "L9gB6U-KXHZY",
"outputId": "8c46b020-3881-4436-bf82-9e951a1a15be"
},
"execution_count": 2,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"0"
]
},
"metadata": {},
"execution_count": 2
}
]
},
{
"cell_type": "code",
"source": [
"import os\n",
"import glob\n",
"import json\n",
"\n",
"shard_filenames = sorted(glob.glob(os.path.join('TinyStories4', \"*.json\")))\n",
"with open(shard_filenames[0], \"r\") as f:\n",
" data = json.load(f)\n",
"\n",
"stories = [x['story'] for x in data]\n",
"text = \"\\n\".join(stories)"
],
"metadata": {
"id": "uXrqxXeyXRwR"
},
"execution_count": 3,
"outputs": []
},
{
"cell_type": "code",
"source": [
"len(text)"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "3Isc2XRfYGW6",
"outputId": "d522f4c4-3e81-4714-cafe-a8dc393ae512"
},
"execution_count": 4,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"77586884"
]
},
"metadata": {},
"execution_count": 4
}
]
},
{
"cell_type": "code",
"source": [
"import string\n",
"\n",
"# Define the set of characters to remove\n",
"remove_chars = \"\\$%&*+-/;`|~ éñ–—…\"\n",
"\n",
"# Create a translation table that maps the characters to be removed to None\n",
"trans_table = str.maketrans(remove_chars, ' ' * len(remove_chars))\n",
"\n",
"# Remove the characters from the string using the translation table\n",
"text = text.translate(trans_table)"
],
"metadata": {
"id": "7PlaC0WIYeD-"
},
"execution_count": 5,
"outputs": []
},
{
"cell_type": "code",
"source": [
"### vocab size and characters used\n",
"\n",
"chars = sorted(list(set(text)))\n",
"vocab_size = len(chars)\n",
"print(''.join(chars))\n",
"print(vocab_size)"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "DpPYXMqJZkU7",
"outputId": "66cbd95c-4342-48ff-ce7f-de321d7d0501"
},
"execution_count": 6,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\t\n",
" 0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz \n",
"66\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"import torch\n",
"import torch.nn as nn\n",
"from torch.nn import functional as F\n",
"torch.manual_seed(1137)\n",
"\n",
"# hyperparameters\n",
"batch_size = 64 # how many independent sequences will we process in parallel?\n",
"block_size = 256 # what is the maximum context length for predictions?\n",
"max_iters = 6000\n",
"eval_interval = 100\n",
"learning_rate = 1e-3\n",
"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
"eval_iters = 200\n",
"n_embed = 384\n",
"n_head = 6\n",
"n_layer = 6\n",
"dropout = 0.0\n",
"# ------------"
],
"metadata": {
"id": "FQr01U7Hahth"
},
"execution_count": 7,
"outputs": []
},
{
"cell_type": "code",
"source": [
"### character encoding and decoding function, stoi : string to integer, itos: integer to string\n",
"\n",
"import torch\n",
"\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]\n",
"decode = lambda l: \"\".join([itos[x] for x in l])\n",
"\n",
"data = torch.tensor(encode(text), dtype = torch.long)"
],
"metadata": {
"id": "lQcB-pEzZvoh"
},
"execution_count": 8,
"outputs": []
},
{
"cell_type": "code",
"source": [
"### creating the train and text splits\n",
"\n",
"n = int(0.9*len(data))\n",
"train_data = data[:n]\n",
"val_data = data[n:]"
],
"metadata": {
"id": "u6wCdZ5waJNy"
},
"execution_count": 9,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# data loading\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"
],
"metadata": {
"id": "F86Zm98WaS8a"
},
"execution_count": 10,
"outputs": []
},
{
"cell_type": "code",
"source": [
"class RMSNorm(nn.Module): #### x * (sqrt(mean(x**2))\n",
" def __init__(self, n_embed, eps = 1e-6):\n",
" super().__init__()\n",
" self.eps = eps\n",
" self.weight = nn.Parameter(torch.ones(n_embed))\n",
"\n",
" def _norm(self, x):\n",
" return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n",
"\n",
" def forward(self, x):\n",
" output = self._norm(x.float()).type_as(x)\n",
" return output * self.weight"
],
"metadata": {
"id": "_93TwyTwaknl"
},
"execution_count": 11,
"outputs": []
},
{
"cell_type": "code",
"source": [
"class MoeLayer(nn.Module):\n",
" def __init__(self, experts, gate, k=2):\n",
" super().__init__()\n",
" assert len(experts) > 0\n",
" self.experts = nn.ModuleList(experts)\n",
" self.gate = gate\n",
" self.k = k\n",
"\n",
" def forward(self, inputs: torch.Tensor):\n",
" inputs_squashed = inputs.view(-1, inputs.shape[-1])\n",
" gate_logits = self.gate(inputs_squashed)\n",
" weights, selected_experts = torch.topk(\n",
" gate_logits, self.k\n",
" )\n",
" weights = nn.functional.softmax(\n",
" weights,\n",
" dim=1,\n",
" dtype=torch.float,\n",
" ).type_as(inputs)\n",
" results = torch.zeros_like(inputs_squashed)\n",
" for i, expert in enumerate(self.experts):\n",
" batch_idx, nth_expert = torch.where(selected_experts == i)\n",
" results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(\n",
" inputs_squashed[batch_idx]\n",
" )\n",
" return results.view_as(inputs)"
],
"metadata": {
"id": "KYiPeTfua3NM"
},
"execution_count": 12,
"outputs": []
},
{
"cell_type": "code",
"source": [
"class Head(nn.Module):\n",
" def __init__(self, head_size):\n",
" super().__init__()\n",
" self.key = nn.Linear(n_embed, head_size, bias = False)\n",
" self.query = nn.Linear(n_embed, head_size, bias = False)\n",
" self.value = nn.Linear(n_embed, head_size, bias = False)\n",
" self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
" self.dropout = nn.Dropout(dropout)\n",
"\n",
" def forward(self, x):\n",
" B, T, C = x.shape\n",
" k = self.key(x)\n",
" q = self.query(x)\n",
" wei = q @ k.transpose(-2, -1) * C**-0.5\n",
" wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
" wei = F.softmax(wei, dim=-1)\n",
" wei = self.dropout(wei)\n",
" v = self.value(x)\n",
" out = wei @ v\n",
" return out\n",
"\n",
"class MulitHeadAttention(nn.Module):\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_embed, n_embed)\n",
" self.dropout = nn.Dropout(dropout)\n",
"\n",
" def forward(self, x):\n",
" x = torch.cat([head(x) for head in self.heads], dim=-1)\n",
" out = self.dropout(self.proj(x))\n",
" return out\n",
"\n",
"\n",
"class Expert(nn.Module):\n",
" def __init__(self, n_embed):\n",
" super().__init__()\n",
" self.net = nn.Sequential(\n",
" nn.Linear(n_embed, 4* n_embed),\n",
" nn.SiLU(),\n",
" nn.Linear(4 * n_embed, n_embed),\n",
" nn.Dropout(dropout))\n",
"\n",
" def forward(self, x):\n",
" return self.net(x)\n",
"\n",
"class Block(nn.Module):\n",
" def __init__(self, n_embed, n_head, num_experts=4):\n",
" super().__init__()\n",
" self.sa_head= MulitHeadAttention(n_head, n_embed//n_head)\n",
" self.ffw = MoeLayer(\n",
" experts=[Expert(n_embed) for _ in range(num_experts)],\n",
" gate=nn.Linear(n_embed, num_experts, bias=False),\n",
" )\n",
"\n",
"# self.ffw= FeedForward(n_embed)\n",
" self.ln1 = RMSNorm(n_embed)\n",
" self.ln2 = RMSNorm(n_embed)\n",
"\n",
" def forward(self, x):\n",
" x = x + self.sa_head(self.ln1(x))\n",
" x = x+self.ffw(self.ln2(x))\n",
" return x\n",
"\n",
"\n",
"class Transformer(nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
"\n",
" self.token_embedding_table = nn.Embedding(vocab_size, n_embed, device=device)\n",
" self.position_embedding_table = nn.Embedding(block_size, n_embed, device=device)\n",
" self.blocks = nn.Sequential(*[Block(n_embed, n_head=n_head) for _ in range(n_layer)])\n",
" self.lm_head = nn.Linear(n_embed, vocab_size)\n",
"\n",
"\n",
" def forward(self, idx, targets=None):\n",
" B, T = idx.shape\n",
"\n",
" token_emb = self.token_embedding_table(idx)\n",
" pos_emb = self.position_embedding_table(torch.arange(T).to(device))\n",
" x = token_emb + pos_emb\n",
" x = self.blocks(x)\n",
" logits = self.lm_head(x)\n",
" if targets == 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",
" return logits, loss\n",
"\n",
" def generate(self, idx, max_new_tokes):\n",
" for _ in range(max_new_tokes):\n",
" idx_cond = idx[:, -block_size:]\n",
" logits, loss = self(idx_cond)\n",
" logits = logits[:, -1, :]\n",
" probs = F.softmax(logits, dim = -1)\n",
" idx_next = torch.multinomial(probs, num_samples = 1)\n",
" idx = torch.cat((idx, idx_next), dim = 1)\n",
" return idx"
],
"metadata": {
"id": "lJi2xz0obLHn"
},
"execution_count": 13,
"outputs": []
},
{
"cell_type": "code",
"source": [
"model = Transformer()\n",
"optimizer = torch.optim.AdamW(model.parameters(),lr=learning_rate)\n",
"print(sum(p.numel() for p in model.parameters()), 'total parameters')"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "M-R7m5FZbkbl",
"outputId": "96c6cfe5-dede-4c9d-8c20-9411a5e3e966"
},
"execution_count": 14,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"32061762 total parameters\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"model.eval()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "vmZb3UxWboZ9",
"outputId": "25f33004-7034-4fe8-af18-33477e566804"
},
"execution_count": 15,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"Transformer(\n",
" (token_embedding_table): Embedding(66, 384)\n",
" (position_embedding_table): Embedding(256, 384)\n",
" (blocks): Sequential(\n",
" (0): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" (1): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" (2): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" (3): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" (4): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" (5): Block(\n",
" (sa_head): MulitHeadAttention(\n",
" (heads): ModuleList(\n",
" (0-5): 6 x Head(\n",
" (key): Linear(in_features=384, out_features=64, bias=False)\n",
" (query): Linear(in_features=384, out_features=64, bias=False)\n",
" (value): Linear(in_features=384, out_features=64, bias=False)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" (proj): Linear(in_features=384, out_features=384, bias=True)\n",
" (dropout): Dropout(p=0.0, inplace=False)\n",
" )\n",
" (ffw): MoeLayer(\n",
" (experts): ModuleList(\n",
" (0-3): 4 x Expert(\n",
" (net): Sequential(\n",
" (0): Linear(in_features=384, out_features=1536, bias=True)\n",
" (1): SiLU()\n",
" (2): Linear(in_features=1536, out_features=384, bias=True)\n",
" (3): Dropout(p=0.0, inplace=False)\n",
" )\n",
" )\n",
" )\n",
" (gate): Linear(in_features=384, out_features=4, bias=False)\n",
" )\n",
" (ln1): RMSNorm()\n",
" (ln2): RMSNorm()\n",
" )\n",
" )\n",
" (lm_head): Linear(in_features=384, out_features=66, bias=True)\n",
")"
]
},
"metadata": {},
"execution_count": 15
}
]
},
{
"cell_type": "code",
"source": [
"!pip install wandb"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "RroEewThbtDy",
"outputId": "1974cda2-feee-49c5-e27d-c400255a822c"
},
"execution_count": 16,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Collecting wandb\n",
" Downloading wandb-0.16.4-py3-none-any.whl (2.2 MB)\n",
"\u001b[?25l \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/2.2 MB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K \u001b[91m━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.3/2.2 MB\u001b[0m \u001b[31m7.4 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m35.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hRequirement already satisfied: Click!=8.0.0,>=7.1 in /usr/local/lib/python3.10/dist-packages (from wandb) (8.1.7)\n",
"Collecting GitPython!=3.1.29,>=1.0.0 (from wandb)\n",
" Downloading GitPython-3.1.42-py3-none-any.whl (195 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m195.4/195.4 kB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hRequirement already satisfied: requests<3,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (2.31.0)\n",
"Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (5.9.5)\n",
"Collecting sentry-sdk>=1.0.0 (from wandb)\n",
" Downloading sentry_sdk-1.43.0-py2.py3-none-any.whl (264 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m264.6/264.6 kB\u001b[0m \u001b[31m29.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hCollecting docker-pycreds>=0.4.0 (from wandb)\n",
" Downloading docker_pycreds-0.4.0-py2.py3-none-any.whl (9.0 kB)\n",
"Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from wandb) (6.0.1)\n",
"Collecting setproctitle (from wandb)\n",
" Downloading setproctitle-1.3.3-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (30 kB)\n",
"Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from wandb) (67.7.2)\n",
"Requirement already satisfied: appdirs>=1.4.3 in /usr/local/lib/python3.10/dist-packages (from wandb) (1.4.4)\n",
"Requirement already satisfied: protobuf!=4.21.0,<5,>=3.19.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (3.20.3)\n",
"Requirement already satisfied: six>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from docker-pycreds>=0.4.0->wandb) (1.16.0)\n",
"Collecting gitdb<5,>=4.0.1 (from GitPython!=3.1.29,>=1.0.0->wandb)\n",
" Downloading gitdb-4.0.11-py3-none-any.whl (62 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m8.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hRequirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (3.3.2)\n",
"Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (3.6)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (2.0.7)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.0.0->wandb) (2024.2.2)\n",
"Collecting smmap<6,>=3.0.1 (from gitdb<5,>=4.0.1->GitPython!=3.1.29,>=1.0.0->wandb)\n",
" Downloading smmap-5.0.1-py3-none-any.whl (24 kB)\n",
"Installing collected packages: smmap, setproctitle, sentry-sdk, docker-pycreds, gitdb, GitPython, wandb\n",
"Successfully installed GitPython-3.1.42 docker-pycreds-0.4.0 gitdb-4.0.11 sentry-sdk-1.43.0 setproctitle-1.3.3 smmap-5.0.1 wandb-0.16.4\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"import wandb\n",
"!wandb login"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "E-rS2QXCcdpE",
"outputId": "c79ac3b0-b729-45e4-ae1d-ed682cc1684f"
},
"execution_count": 17,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://wandb.ai/authorize\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Paste an API key from your profile and hit enter, or press ctrl+c to quit: \n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"@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",
" X = X.to(device)\n",
" Y = Y.to(device)\n",
" logits, loss = model(X, Y)\n",
" losses[k] = loss.item()\n",
" out[split] = losses.mean()\n",
" model.train()\n",
" return out"
],
"metadata": {
"id": "oNzq5Vkccmk0"
},
"execution_count": 18,
"outputs": []
},
{
"cell_type": "code",
"source": [
"model = model.to(device)\n",
"optimizer = torch.optim.AdamW(model.parameters(),lr=1e-4)\n",
"\n",
"wandb.init(\n",
" # set the wandb project where this run will be logged\n",
" project=\"mixture of experts\",\n",
"\n",
" # track hyperparameters and run metadata\n",
" config={\n",
" 'batch_size': batch_size,\n",
" 'block_size': block_size,\n",
" 'max_iters': max_iters,\n",
" 'eval_interval': eval_interval,\n",
" 'learning_rate': learning_rate,\n",
" 'device': device,\n",
" 'eval_iters': eval_iters,\n",
" 'n_embed': n_embed,\n",
" 'n_head': n_head,\n",
" 'n_layer': n_layer,\n",
" 'dropout': dropout\n",
" }\n",
")\n",
"\n",
"wandb.watch(model)\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 % 100 == 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",
" wandb.log({'train_loss': losses['train'], 'val_loss': losses['val']}, step=iter)\n",
"\n",
" # sample a batch of data\n",
" xb, yb = get_batch('train')\n",
" xb = xb.to(device)\n",
" yb = yb.to(device)\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",
"model_dir = 'Mixture of experts models'\n",
"os.makedirs(model_dir, exist_ok=True)\n",
"final_model_path = os.path.join(model_dir, 'moe2_model.pth')\n",
"torch.save(model.state_dict(), final_model_path)\n",
"print('Final trained model saved!')\n",
"\n",
"wandb.finish()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000,
"referenced_widgets": [
"6a9a993a2bde4bc7acdd44e35ee8a423",
"63b857b0db6f4b5bbb55230dd1e0f66e",
"ae1f1250131c4283b929d3108e581926",
"e8190f4f922d4a1c9f8c730604c7e806",
"2b5adaa739be47f088c2a829508a1d43",
"fa697fde7f024d07aee37a56b55fbf59",
"9dbb0157705f4efb90649e4468708109",
"36f393f4de7c434c8480fbf195d0ac8b"
]
},
"id": "J-PJuS8qdCSW",
"outputId": "4b4a6adb-6df8-4b62-e082-fb3dbd6c0a0b"
},
"execution_count": 19,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mdruvithlgowda00\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
]
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"/content/wandb/run-20240323_052806-r801x9fx"
]
},
"metadata": {}
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"| train_loss | █▄▄▄▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ |
| val_loss | █▄▄▄▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ |
| train_loss | 0.84435 |
| val_loss | 0.84581 |
./wandb/run-20240323_052806-r801x9fx/logs"
]
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"# generation\n",
"d = 'once'\n",
"context = torch.tensor(encode(d), dtype=torch.long, device=device).unsqueeze(0)\n",
"print(decode(model.generate(context, max_new_tokes=500)[0].tolist()))\n"
],
"metadata": {
"id": "FyJODzaHezJa",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "318940c5-e203-4f4f-bfde-a114edaff372"
},
"execution_count": 26,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"once to go home \n",
"Once upon a time there was a little girl named Lily She loved to play with her together in the tracks One day her mom made a little here with a big sea She had answed \n",
" Let s play as play white The tree was crying and continued but then we down the sea tost inside the spick replied That dog for it is thought about \n",
"Then it s house The boy and the man fell down about the pole and excited away The wind was go away to and Tom again They excited to put the circle \n"
]
}
]
},
{
"cell_type": "code",
"source": [
"#push the weights to hugging_face.hub\n",
"!pip install transformers huggingface_hub"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Ie3l0Ks79VWV",
"outputId": "161fb7ba-9b18-43dc-e2cb-8577484ad939"
},
"execution_count": 27,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.38.2)\n",
"Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.10/dist-packages (0.20.3)\n",
"Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.13.1)\n",
"Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.25.2)\n",
"Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (24.0)\n",
"Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.1)\n",
"Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2023.12.25)\n",
"Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.31.0)\n",
"Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.15.2)\n",
"Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.2)\n",
"Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.66.2)\n",
"Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub) (2023.6.0)\n",
"Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub) (4.10.0)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.3.2)\n",
"Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.6)\n",
"Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.7)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2024.2.2)\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"from huggingface_hub import notebook_login\n",
"\n",
"notebook_login()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 157,
"referenced_widgets": [
"0a85aa309bda4b9abf9240eeb8f403f0",
"38c9ff2937f142fdaf2fdf488fda9656",
"1e8f3c64ce81450994fc21b78e3aba76",
"c80588db8c264ec0a4c17c890c213963",
"5aa58b3190a148b1b9bbe28acc811192",
"12a46681a6004746b6869569c65b71a2",
"fb800a6d97954553b45cfb9c712dab9c",
"9c102c35b48a454d9e83f61507097b4c",
"aa87be83bcc24e8aad06ffa648d1a057",
"a3eee979a8d544e8bdb76682d4e1a89f",
"7089322af59a42d4977b3251f0b7e2cf",
"057290b3b190427c8384d5a271bac2b5",
"53c0c50ce74d425d8125d70aeead96c8",
"5bd49cf339f64cf2b46079cf75bda5fd",
"c78516a5de83466f884572700ce225d1",
"9444ddb29edd421da3d74baf0cece75d",
"20cdbaaf3958455c862a2756e92e8f92",
"eb83010cb46043268a8b36f8f25c35e7",
"8247f0e0fe2d48bdb1cd9627f6e53417",
"ffb23aea10c94126bd8bdc54b265908d",
"9df202bad1e84ef8bc68982042d8831f",
"7681267bc11d410ebc3954064d1cd416",
"eaf3299530ba4f68a75b107d54bcd027",
"2b4660cfb7c24a269c8dda2a89c7fcef",
"075428eaefb44bcbb96d96d2019deeef",
"c75f2dcbcd724ba0aed77744c6943faf",
"6dba6d0fdf464e84a1319e793974a455",
"d1b60ca0cb6947a7afb1007de23f59aa",
"3f86a6ed43c14dbeb965535c44a114bd",
"f458884d36c64771ba0011bc94017dcf",
"e902c5cb7e884850ad1ba65fe7300a7a",
"cd97d20ae2e742ed9476c0829f1ee586"
]
},
"id": "7EUxRlt0ARds",
"outputId": "c18522ee-f346-460a-925f-d421ac92be58"
},
"execution_count": 28,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"VBox(children=(HTML(value='