Task3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Task 3: SGPU, DDP, FSDP"
      ],
      "metadata": {
        "id": "6tB4nu0rgg0x"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D9mAD8bPRiNH"
      },
      "outputs": [],
      "source": [
        "from dataclasses import dataclass\n",
        "import copy\n",
        "import torch\n",
        "import math\n",
        "import torch.nn as nn\n",
        "from torch.nn.parameter import Parameter\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "# from torch.nn.modules.normalization import LayerNorm\n",
        "from torch.nn.modules import ModuleList\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ywDaj8xQhB1k"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7ZnOiePxoC1"
      },
      "source": [
        "#Model Architecture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "enLLYvtdWiuh"
      },
      "outputs": [],
      "source": [
        "class LayerNorm(nn.Module):\n",
        "\n",
        "    def __init__(self, hidden_size, epsilon=1e-12):\n",
        "      \"\"\"\n",
        "      Initialize LayerNorm module.\n",
        "      \"\"\"\n",
        "      super().__init__()\n",
        "\n",
        "      # Learnable weight parameter for scaling.\n",
        "      self.weight = nn.Parameter(torch.ones(hidden_size))\n",
        "\n",
        "      # Learnable bias parameter for shifting.\n",
        "      self.bias = nn.Parameter(torch.zeros(hidden_size))\n",
        "\n",
        "      # Small value to avoid division by zero in normalization.\n",
        "      self.epsilon = epsilon\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "      # Compute mean and variance along the last dimension.\n",
        "      u = x.mean(-1, keepdim=True)\n",
        "      s = (x - u).pow(2).mean(-1, keepdim=True)\n",
        "\n",
        "      # Normalize the input tensor.\n",
        "      x = (x - u) / torch.sqrt(s + self.epsilon)\n",
        "\n",
        "      # Scale and shift using learnable parameters.\n",
        "      return self.weight * x + self.bias"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IA2GLIrnXC3M"
      },
      "outputs": [],
      "source": [
        "class Conv1D(nn.Module):\n",
        "    def __init__(self, nx, nf):\n",
        "        '''\n",
        "        nx: Number of input features.\n",
        "        nf: Number of filters (output channels).\n",
        "        '''\n",
        "        super().__init__()\n",
        "        self.nf = nf\n",
        "        #intialising an empty matrix as weights for size of (nx)X(nf)\n",
        "        w = torch.empty(nx, nf)\n",
        "        #initialising these weights as normal distribution\n",
        "        nn.init.normal_(w, std=0.02)\n",
        "        #calculating the weights and biases by encoding them using nn.Parameter\n",
        "        self.weight = nn.Parameter(w)\n",
        "        self.bias = nn.Parameter(torch.zeros(nf))\n",
        "\n",
        "    def forward(self, x):\n",
        "        '''x:The input tensor'''\n",
        "        #this size output is summation of x second dimension and the nf dimension\n",
        "        size_out = x.size()[:-1] + (self.nf,)\n",
        "        # dot multiplying Q,K(transpose) and V\n",
        "        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)#x.view helps in taking the transpose out\n",
        "        x = x.view(*size_out)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LUbQPBgRhHus",
        "outputId": "063a9db5-1b83-4e9a-864a-431092f11f19"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 4, 2304])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "d_model = 768\n",
        "conv1d  = Conv1D(d_model, d_model*3)\n",
        "x = torch.rand(1,4,d_model) # A random input tensor for testing out our model\n",
        "x = conv1d(x)\n",
        "x.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8CL-y_u6hIzH",
        "outputId": "de8e574f-1d33-42a5-921f-28474d6e90a0"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([1, 4, 768]), torch.Size([1, 4, 768]), torch.Size([1, 4, 768]))"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "query, key, value = x.split(d_model, dim=-1)# here we can split our output attention into query, key and value\n",
        "\n",
        "query.shape, key.shape, value.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "emM2PLrfm_cz"
      },
      "outputs": [],
      "source": [
        "class FeedForward(nn.Module):\n",
        "    def __init__(self, dropout, d_model=768, nx=768*4):\n",
        "        super().__init__()\n",
        "        self.c_fc    = Conv1D(d_model, nx)\n",
        "        self.c_proj  = Conv1D(nx, d_model)\n",
        "        self.act     = F.gelu\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.dropout(self.c_proj(self.act(self.c_fc(x))))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "05jSf9YYnJdY"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, d_model=768, n_head=12, n_ctx=1024, d_head=64, bias=True, scale=False):\n",
        "        '''Constructor funtion\n",
        "        Params:\n",
        "        d_model:The dimension that needs to be feed into our model\n",
        "        n_head:The number of heads for attention\n",
        "        n_ctx:a parameters for buffer registry for bias\n",
        "        d_head:the dimension head output\n",
        "        bias:A bool for including or not including bias\n",
        "        scale: Whether to scale the attention scores by the square root of the dimension of the queries(use sqrt(dk) or not) \"\n",
        "        '''\n",
        "        super().__init__()\n",
        "        self.n_head  = n_head\n",
        "        self.d_model = d_model\n",
        "        self.c_attn  = Conv1D(d_model, d_model*3)\n",
        "        self.scale   = scale\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "        self.register_buffer(\"bias\", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "        self.c_proj  = Conv1D(d_model, d_model)\n",
        "\n",
        "    def split_heads(self, x):\n",
        "        \"\"\"\n",
        "        spliting inyo given number of heads and then returning\n",
        "        return shape [`batch`, `head`, `sequence`, `features`]\n",
        "        \"\"\"\n",
        "        new_shape = x.size()[:-1] + (self.n_head, x.size(-1)//self.n_head)\n",
        "        x = x.view(*new_shape)\n",
        "        return x.permute(0, 2, 1, 3)\n",
        "\n",
        "    def _attn(self, q, k, v, attn_mask=None):\n",
        "        \"\"\"The main attention function\n",
        "        That claculates the attention using our dot product formula\"\"\"\n",
        "        scores  = torch.matmul(q, k.transpose(-2, -1))# dot multiplication between q and k transpose\n",
        "        if self.scale: scores = scores/math.sqrt(v.size(-1))# scaling it by dividing by sqrt(dk)\n",
        "        nd, ns  = scores.size(-2), scores.size(-1)\n",
        "        if attn_mask is not None: scores = scores + attn_mask# adding scores with attention mask values\n",
        "        scores  = self.softmax(scores)# adding softmax values\n",
        "        scores  = self.dropout(scores) #dropout of 0.1 as mentioned\n",
        "        outputs = torch.matmul(scores, v) # now the final matrix multiplication between score and V\n",
        "        return outputs\n",
        "\n",
        "    def merge_heads(self, x):\n",
        "        # merging the attention heads into one\n",
        "        x = x.permute(0, 2, 1, 3).contiguous()\n",
        "        new_shape = x.size()[:-2] + (x.size(-2)*x.size(-1),)\n",
        "        return x.view(*new_shape)\n",
        "\n",
        "    def forward(self, x):\n",
        "        '''The feed forward function that calculates the attention, split the heads, make attention, merge heads and project out the output '''\n",
        "        x        = self.c_attn(x) #new `x` shape - `[1,3,2304]`\n",
        "        q, k, v  = x.split(self.d_model, dim=2)\n",
        "        q, k, v  = self.split_heads(q), self.split_heads(k), self.split_heads(v)\n",
        "        out      = self._attn(q, k, v)\n",
        "        out      = self.merge_heads(out)\n",
        "        out      = self.c_proj(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BA8VXNSEnL52"
      },
      "outputs": [],
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, d_model=768, n_head=12, dropout=0.1):\n",
        "        super(TransformerBlock, self).__init__()\n",
        "        self.attn        = MultiHeadAttention(d_model=768, n_head=12, d_head=64, n_ctx=1024, bias=True, scale=False)\n",
        "        self.feedforward = FeedForward(dropout=0.1, d_model=768, nx=768*4)\n",
        "        self.ln_1        = LayerNorm(d_model)\n",
        "        self.ln_2        = LayerNorm(d_model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.attn(self.ln_1(x))\n",
        "        x = x + self.feedforward(self.ln_2(x))\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GKnGdiUMeHC1"
      },
      "outputs": [],
      "source": [
        "def _get_clones(module, n):\n",
        "    '''Here we can make certain copies of transformers'''\n",
        "    return ModuleList([copy.deepcopy(module) for i in range(n)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fH_GtXHSUkmG"
      },
      "outputs": [],
      "source": [
        "class GPT2(nn.Module):\n",
        "    def __init__(self, nlayers=12, n_ctx=1024, d_model=768, vcb_sz=50257):\n",
        "        '''nlayer:The number of times the tarnsformer needs to get cloned\n",
        "        n_ctx:The highest length that can be these to get teh string positional embeddings\n",
        "        d_model:The dimenionalities for model\n",
        "        vcb_sz:The vocablury size which can be later altered while training'''\n",
        "        super(GPT2, self).__init__()\n",
        "        self.nlayers = nlayers\n",
        "        block        = TransformerBlock(d_model=768, n_head=12, dropout=0.1)\n",
        "        self.h       = _get_clones(block, 12)\n",
        "        self.wte     = nn.Embedding(vcb_sz, d_model)\n",
        "        self.wpe     = nn.Embedding(n_ctx, d_model)\n",
        "        self.drop    = nn.Dropout(0.1)\n",
        "        self.ln_f    = LayerNorm(d_model)\n",
        "        self.out     = nn.Linear(d_model, vcb_sz, bias=False)\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "        self.init_weights()\n",
        "\n",
        "    def init_weights(self):\n",
        "        '''Initialization of weights'''\n",
        "        self.out.weight = self.wte.weight\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        '''If the Linear, Embedding and Conv1D then nomrally initializing with mean and S.D'''\n",
        "        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):\n",
        "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
        "            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:\n",
        "                '''Data Bias zero'''\n",
        "                module.bias.data.zero_()\n",
        "        elif isinstance(module, nn.LayerNorm):\n",
        "            module.bias.data.zero_()\n",
        "            module.weight.data.fill_(1.0)\n",
        "\n",
        "    def forward(self, src, labels=None, pos_ids=None):\n",
        "        '''Adding the positional embedding, dropping, then adding inputs, logits and outputs which are being used for loss function and then adding outputs and loss'''\n",
        "        if pos_ids is None:\n",
        "            pos_ids = torch.arange(0, src.size(-1)).unsqueeze(0)\n",
        "        pos_ids = pos_ids.to(src.device)  # Ensure pos_ids is on the same device as src\n",
        "        inp = self.drop((self.wte(src) + self.wpe(pos_ids)))\n",
        "        for i in range(self.nlayers): inp = self.h[i](inp)\n",
        "        inp     = self.ln_f(inp)\n",
        "        logits  = self.out(inp)\n",
        "        outputs = (logits,) + (inp,)\n",
        "\n",
        "        if labels is not None:\n",
        "            shift_logits = logits[..., :-1, :].contiguous()\n",
        "            shift_labels = labels[..., 1:].contiguous()\n",
        "            loss = self.loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n",
        "            outputs = (loss,) + outputs\n",
        "            return loss.mean()\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WVQCPgaJpg43"
      },
      "source": [
        "# Single GPU Training Loop\n",
        "This is a code for training our GPT-2 on a single GPU\n",
        "For these purpose I would be using tinyshakespeare/input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tnbPIZ7fyWYc"
      },
      "outputs": [],
      "source": [
        "class TrainerConfig:\n",
        "    '''Class for setting the Training for training configurations'''\n",
        "    # optimization parameters\n",
        "    max_epochs = 10\n",
        "    batch_size = 8\n",
        "    learning_rate = 3e-4\n",
        "    betas = (0.9, 0.95)\n",
        "    grad_norm_clip = 1.0\n",
        "    weight_decay = 0.1  # only applied on matmul weights\n",
        "    # learning rate decay params: linear warmup followed by cosine decay to 10% of original\n",
        "    lr_decay = False\n",
        "    warmup_tokens = 375e6  # these two numbers come from the GPT-3 paper, but may not be good defaults elsewhere\n",
        "    final_tokens = 260e9  # (at what we reach 10% of original LR)\n",
        "    # checkpoint settings\n",
        "    ckpt_path = None\n",
        "    num_workers = 0  # for DataLoader\n",
        "\n",
        "    def __init__(self, **kwargs):\n",
        "        for k, v in kwargs.items():\n",
        "            setattr(self, k, v)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6VSBiDMYVc4n"
      },
      "outputs": [],
      "source": [
        "class Trainer:\n",
        "    def __init__(self, model, train_dataset, test_dataset, config):\n",
        "        self.model = model\n",
        "        self.train_dataset = train_dataset\n",
        "        self.test_dataset = test_dataset\n",
        "        self.config = config\n",
        "\n",
        "        # take over whatever gpus are on the system\n",
        "        self.device = \"cpu\"\n",
        "        if torch.cuda.is_available():\n",
        "            self.device = torch.cuda.current_device()\n",
        "            self.model = torch.nn.DataParallel(self.model).to(self.device)\n",
        "\n",
        "    def save_checkpoint(self):\n",
        "        # DataParallel wrappers keep raw model object in .module attribute\n",
        "        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n",
        "        logger.info(f\"saving {self.config.ckpt_path}\")\n",
        "        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n",
        "\n",
        "    def train(self):\n",
        "        model, config = self.model, self.config#getting the model and training configurations\n",
        "        raw_model = model.module if hasattr(self.model, \"module\") else model\n",
        "        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=config.learning_rate)\n",
        "\n",
        "        def run_epoch(split):\n",
        "            is_train = split == \"train\"\n",
        "            model.train(is_train)\n",
        "            data = self.train_dataset if is_train else self.test_dataset\n",
        "            loader = DataLoader(\n",
        "                data,\n",
        "                shuffle=True,\n",
        "                pin_memory=True,\n",
        "                batch_size=config.batch_size,\n",
        "                num_workers=config.num_workers,\n",
        "            )\n",
        "\n",
        "            losses = []\n",
        "            pbar = (\n",
        "                tqdm(enumerate(loader), total=len(loader))\n",
        "                if is_train\n",
        "                else enumerate(loader)\n",
        "            )\n",
        "            for it, (x, y) in pbar:\n",
        "\n",
        "                # place data on the correct device\n",
        "                x = x.to(self.device)\n",
        "                y = y.to(self.device)\n",
        "\n",
        "                # forward the model\n",
        "                with torch.set_grad_enabled(is_train):\n",
        "                    loss = model(x, y)  # The forward method returns the mean of the loss directly\n",
        "                    logits = model(x)\n",
        "                    losses.append(loss.item())\n",
        "\n",
        "                if is_train:\n",
        "\n",
        "                    # backprop and update the parameters\n",
        "                    model.zero_grad()\n",
        "                    loss.backward(retain_graph=True)\n",
        "                    torch.nn.utils.clip_grad_norm_(\n",
        "                        model.parameters(), config.grad_norm_clip\n",
        "                    )\n",
        "                    self.optimizer.step()\n",
        "\n",
        "                    # decay the learning rate based on our progress\n",
        "                    if config.lr_decay:\n",
        "                        self.tokens += (\n",
        "                            y >= 0\n",
        "                        ).sum()  # number of tokens processed this step (i.e. label is not -100)\n",
        "                        if self.tokens < config.warmup_tokens:\n",
        "                            # linear warmup\n",
        "                            lr_mult = float(self.tokens) / float(\n",
        "                                max(1, config.warmup_tokens)\n",
        "                            )\n",
        "                        else:\n",
        "                            # cosine learning rate decay\n",
        "                            progress = float(\n",
        "                                self.tokens - config.warmup_tokens\n",
        "                            ) / float(\n",
        "                                max(1, config.final_tokens - config.warmup_tokens)\n",
        "                            )\n",
        "                            lr_mult = max(\n",
        "                                0.1, 0.5 * (1.0 + math.cos(math.pi * progress))\n",
        "                            )\n",
        "                        lr = config.learning_rate * lr_mult\n",
        "                        for param_group in self.optimizer.param_groups:\n",
        "                            param_group[\"lr\"] = lr\n",
        "                    else:\n",
        "                        lr = config.learning_rate\n",
        "\n",
        "                    # repeat progress\n",
        "                    pbar.set_description(\n",
        "                        f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}, lr {lr:e}\"\n",
        "                    )\n",
        "\n",
        "            if not is_train:\n",
        "                test_loss = float(np.mean(losses))\n",
        "                logger.info(f\"test loss: {test_loss}\")\n",
        "                return test_loss\n",
        "\n",
        "        best_loss = float(\"inf\")\n",
        "        self.tokens = 0  # counter used for learning rate decay\n",
        "        for epoch in range(config.max_epochs):\n",
        "\n",
        "            run_epoch(\"train\")\n",
        "            if self.test_dataset is not None:\n",
        "                test_loss = run_epoch(\"test\")\n",
        "\n",
        "            # supports early stopping based on the test loss, or just save always is no test set is provided\n",
        "            good_model = self.test_dataset is None or test_loss < best_loss\n",
        "            if self.config.ckpt_path is not None and good_model:\n",
        "                best_loss = test_loss if self.test_dataset is not None else float(\"inf\")\n",
        "                self.save_checkpoint()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QUEMso6wWCKG"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import Dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UU7hvyBzVhG9"
      },
      "outputs": [],
      "source": [
        "class CharDataset(Dataset):\n",
        "\n",
        "    def __init__(self, data, block_size):\n",
        "        chars = sorted(list(set(data)))\n",
        "        data_size, vocab_size = len(data), len(chars)\n",
        "        print(f\"data has {data_size:d} characters, {vocab_size:d} unique.\")\n",
        "\n",
        "        self.stoi = { ch: i for i, ch in enumerate(chars) }\n",
        "        self.itos = { i: ch for i, ch in enumerate(chars) }\n",
        "        self.block_size = block_size\n",
        "        self.vocab_size = vocab_size\n",
        "        self.data = data\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data) - self.block_size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        # grab a chunk of (block_size + 1) characters from the data\n",
        "        chunk = self.data[idx:idx+self.block_size+1]\n",
        "        # encode every character to an integer\n",
        "        dix = [self.stoi[s] for s in chunk]\n",
        "        x = torch.tensor(dix[:-1], dtype=torch.long)\n",
        "        y = torch.tensor(dix[1:], dtype=torch.long)\n",
        "        return x, y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GEjzIKZZVlqX"
      },
      "outputs": [],
      "source": [
        "block_size=25"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y8F1NfERXRKB"
      },
      "source": [
        "Reference link for dataset -- https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0e1Ihc1sWKm5",
        "outputId": "9b3cc7a1-925c-4667-e535-409ddb5cdcd2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-12-17 15:57:50--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1115394 (1.1M) [text/plain]\n",
            "Saving to: ‘input.txt.8’\n",
            "\n",
            "input.txt.8         100%[===================>]   1.06M  6.04MB/s    in 0.2s    \n",
            "\n",
            "2023-12-17 15:57:51 (6.04 MB/s) - ‘input.txt.8’ saved [1115394/1115394]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "neWHlsO4WMgG",
        "outputId": "5849da1a-61e3-4e43-9ad0-53b55a7cdfc6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "data has 1115394 characters, 65 unique.\n"
          ]
        }
      ],
      "source": [
        "text = open('input.txt', 'r').read()\n",
        "train_dataset = CharDataset(text, block_size = 25) # 25 is for context storing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O_wiunBDWRED"
      },
      "outputs": [],
      "source": [
        "\n",
        "trainable_model = GPT2()\n",
        "trainable_model=trainable_model.to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "62DjMwqwXa_E"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data.dataloader import DataLoader\n",
        "from tqdm import tqdm\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 427
        },
        "id": "kiuYM78CW_rs",
        "outputId": "9fdbc9e3-d920-405d-b0cd-d2cb79020759"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:557: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
            "  warnings.warn(_create_warning_msg(\n",
            "  0%|          | 0/69711 [00:02<?, ?it/s]\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-21-6f4d70cd4a74>\u001b[0m in \u001b[0;36m<cell line: 11>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m )\n\u001b[1;32m     10\u001b[0m \u001b[0mtrainer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTrainer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrainable_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_dataset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtconf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-13-a77037aea5c5>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    103\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax_epochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 105\u001b[0;31m             \u001b[0mrun_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"train\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    106\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest_dataset\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    107\u001b[0m                 \u001b[0mtest_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrun_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"test\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-13-a77037aea5c5>\u001b[0m in \u001b[0;36mrun_epoch\u001b[0;34m(split)\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0;31m# forward the model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_grad_enabled\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mis_train\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 51\u001b[0;31m                     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# The forward method returns the mean of the loss directly\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     52\u001b[0m                     \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m                     \u001b[0mlosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-11-000576ec099f>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, src, labels, pos_ids)\u001b[0m\n\u001b[1;32m     41\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnlayers\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0minp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     42\u001b[0m         \u001b[0minp\u001b[0m     \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mln_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 43\u001b[0;31m         \u001b[0mlogits\u001b[0m  \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     44\u001b[0m         \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mlogits\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0minp\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    113\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 114\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    116\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "tconf = TrainerConfig(\n",
        "    max_epochs=1,\n",
        "    batch_size=16,\n",
        "    learning_rate=6e-4,\n",
        "    lr_decay=True,\n",
        "    warmup_tokens=512,\n",
        "    final_tokens=2*len(train_dataset)*block_size,\n",
        "    num_workers=4,\n",
        ")\n",
        "trainer = Trainer(trainable_model, train_dataset, None, tconf)\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hUIr-XV3aXRE"
      },
      "source": [
        "![image.png](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABx4AAABLCAYAAABOSCJlAAAAAXNSR0IArs4c6QAAAARnQU1BAACxjwv8YQUAAAAJcEhZcwAAEnQAABJ0Ad5mH3gAABlySURBVHhe7d1LqCxX2Qbgzo/ibSSigoKSIE4iomSgYoII4kC8RRARR95QxEsGXnCi6EQjDoKKGMQ4CqKCgooDEUSDoAMxCE5UIhEURdSRiTfYv2+5v5N11qmqru7qvc++PA9UTnZ3dV1WrVpV9X1rdd/w8MMPH20AAAAAAAAAVvi/438BAAAAAAAA9ibxCAAAAAAAAKwm8QgAAAAAAACsdsPRfx3/PwAAAAAAAMBejHgEAAAAAAAAVpN4BAAAAAAAAFaTeAQAAAAAAABWk3gEAAAAAAAAVpN4BAAAAAAAAFaTeAQAAAAAAABWk3gEAAAAAAAAVpN4BAAAAAAAAFaTeAQAAAAAAABWk3gEAAAAAAAAVpN4BAAAAAAAAFaTeAQAOEU//elPN7fddtvmN7/5zfErnEVf//rXN6973es2f/vb345fAQAAAGAbicfrJEGsBLMS1DqrEhBNYFTQ7WxLAPumm27a3HHHHZuHH374+NWTUeuq6VOf+tTxO1x2qQtt3Uhdgcugruep930ysa6jzgd2let5ruuX9Tqb++P2mnIa9zjA7tJGjZ2fh3qOrOVUW9BfZwEAgLNJ4pFVjAZ4xHlIJq/1ghe8YPPAAw9sfvnLX25e/epXH78Km80HP/jBoW5873vf2zz96U8/fhUuvi9+8YubF77whUP9f9WrXrX5xje+MbyeIOznPve54bW0nbCL3//+95sHH3xwuK+4rJ73vOdtfvaznw3n1l133bV53OMed/wOcBbk2ecnP/nJ5o1vfONe5+eSzjnPetazNvfdd9/QDnzlK185fhUAADjrJB4vmBqRdohkYD3oJYj6xCc+8fjV86GSgJehd2wlAwXltutHUGRaM5qk74Xd9/ju3++ndt01uqV9f2qERy13W72uUYiXYcRM7evaUTF9Hek7Eowd03aepce8b6PaqQ/ALW3P+m0/jeO+yzrb/ejL9ayr+lVTbX/q2h/+8IfNjTfeOPydf/N3Xv/Od76z+d3vfrd5+9vfPrzH5VZ1aOl5mXuvZz7zmaMdOfp2Zqzd6+ts267MtT+Z5pY3t/35TF3HTuMcX9o29uWVz+Sz+6rlTa1vrHzXjHruj2V/DNpy76f+OCwts23G9nFqWf11Yu11ulS5tMvrj3U/zdXfQ9j2HNZv377lX/q6sa2ejZXZGt///veHf5/97GcP/7bO83MkAACwnsTjBVEBgB/+8Iebd77zncevXk55gM9ovBqBUqNQ3vzmN696uOf8e/3rX3+lTmTK6Lxvf/vbewWiEtx5+ctfvvn0pz99ZXlPe9rTNm9605uuBJvaXtrtVKMCX/KSlwzzZf58Lp+veTLKI0mLD3/4w1eCQxVczCiqt771rcNrYyrw9ZznPGcYMXLRZX+/8IUvHP+1n2pDf/zjHw8jeus4pM6UsWOe3vcf+tCHrgR3lx7zcuedd14zbzs6bml7ljr8mc98ZlhP5llTt5fadZ0ZGXj//fcf/3U+1DmXER018qo95unskfP2t7/97TB//s3fGa2Wssl8Aq6X2z7tcdqj1LkXv/jF13QoyvmV8//LX/7ylTah7XhUdTYJ8GrLcm6+//3vv9JOpU4mGVCfryl1PNvYrnfp9mfZaR8zncao96VtY7Xb733ve4d5UibPeMYzhmtoXauXWnINzjLzXs792q608xkRlm2Zk+2+/fbbr2nbo5Y1187m/r/mq6m9hi0tsyXG6tDYslIv2uvq2L3NPlKWY9f9Xa/Bh1L3EHPPYSmXD3zgA1edu/uW/5LzvDdVZvvKNuTYpj65zgEAAL1LkXish7METjL1vUvrYTEPZHlYq/mmeqvmYb/myTT1gNcua+l8U+uck/372Mc+NgQ58nWHa6Vs2t64Kbs2OFDllfeyzgSSb7nlltl92FYW+TufywN0e6z6dS+RQEi25z3vec/xK5vh//Na3uvV/k4Fy7epfcvyUxYpk9r+TLWvVQ+znr6M+/Lo6+zYPKmv7ftj279r3T6kJedJW5cyTfX8brc9077HqpcAVYI+NUJpqWx3AomZ2iRR9uXPf/7zlR7gU1IPn/zkJ1/pIf6rX/1q+Fw+XxLESaAyAbp//OMfw2uf/exnh4BwgsxPeMIThtd6Kb+77757CO696EUvOn71bEg9rGPc1o819THHLYmgBPrWdLpIUizB6E984hOTI4cTUEzQ9rnPfe7xK5vh//NaJZ6m9Md8qSXtWcozgegkRFOnI//m77zen1P9+TR2bm6z6zpr/gTs55IXmW9Ne3xov/jFL4ZzKdtTgdWc8zn3U++qDcu+pSyTLMoIxxybtC2+YvVyS33epz2ua8jLXvay4d+SNjR1LYmLOu96NdL2ox/96JW2LPOmPladnZL15lr0/Oc/f/h76fZnu5KASOJj1zZuX0vaxvb6UAm4lEnKJsau1TnXp+5FllyDKyHXnvs5jmn3cg2ZkjJMgrQf5Zr7+vbevo5l2ppdr5tLyqy1a3tc9zB/+ctfhn8j5d4mxuveJnUqHTR6dX1KeUzZ57qf/dvnGrxEtmfJc1iO3Te/+c2rzt0qs5///OfDv0vtep4vKbM63nVvsO3erK6PtQ+lX87Ys1zdA6bOpx6kTtT8meaOPwAAcD5c+MRjHpj6kUTp9TzWuzQPPdUrt3pE50GyHpayrDxc9b1LM6qhfSjP/HnIakeCZEqQol9nHlL7dSYAvos8cOah/lABzjy4Vm/hsYfTCqrk/QRfE0zJg2ftZ96rAG2kbPKwW/NUmfXB7iTtbr311iGok/kyfx6qE+hZKsdobJRAHo6/9a1vjQZq8rCfh959gjhRo+iyvSmLlEmVRaa2t3mkt3HqX/V4TtmkHrQP2Xnv3e9+95Vl1DxtmeV4573UmyQ+5qRuJzFS8/d1+5CWnieZL6MC2p73OUapP+12ZZ/7cylfZ3g9gxIJqiU4m3JsJViYAN5cEiptQALHOSbteZI62AbrIsvJOh772McOfyeg1denXs7fe+6556plnyXZzwSaIsdybX08xNdZtsekbTfGtIngyP/ntfqqzTFTx3ybpe1Z2rA+oJqyTLuR8m4Dmkvb4212WWder9863JZ8WdseH1oSBTmn20RAtiv7mWtWOg2018y0X3ktx7sPxnL57NMe53wZG0VU51fOoyx3Sn/dKBnpletWf50pVa/b5S/d/tyPtMmlk7a0bUxbkrauH+WWhGPO35RzyrXUcvs2rCy5Bu8j7W+uD7nnPqlyXFpmrdNuj6vux1ySdtfr/r7X4KVSnod8Dlti1/N8W5mljPJsUiODM9X9wpSx62O018SpJGfOpboHyeeznlpvptMsSwAA4GTsnXg8OjoaAsY/+tGPNn/84x+PX31EHlDz8JiH1n/961/Hr/7Pts8eUvVmbh+0XvnKV4727k3ypka7ZMoDagIWeeiOWlbfuzQPaXmgzUNbVA/Qvkf6K17xiqv+jn6dCQic1gP+aaiH/STN6mG/yiwPmf1+5rUK6mT+BN52GY02lhBKQCfJwAS+xx7G07M/D72n9VVBSU4myFN1IYH7vJaAQHnXu951VV2pUVV9kGypfLZ6uFc96xMoh7L0PKlj1QYEU/4JRtTnIsGVPrCROnKIoESSl0kEL0k4tZ70pCcNSZf2mLXm6mzanXy2HcmSfUlwJttRCaAEfFJeSUDvsm3nQc7HGhWwpj6mLiVp1rYv+6g2If/O9fav60jOp6w77yV5nvYm15UpY8e8ZNunevkvbc/6AGS2Kx1u0q5kW/N+7Noez1m6ztglSHza7fGcnMM5l9Nxqc7BlGH2721ve9s17XakHHJ8UqZtGw5L5Z7zwQcfvCZxXZ0cnvKUp1zzjQhtu5FOEFPtaZZd7V0v1+60KechYb60bax9zTW75Nqae6mPfOQj15RT2py0PWmDatTnIdT2jnVQyfakTU4bXNfFOdUute1vyf1MWy/qfiKWlllrl/Y425Xl5H6t/VaAMbk+5JrYHpeo+4GY+krUfa77c9fg6y1lnvOyPS5L7HKeLymzmr+t97mG5XlkTK51eV7d9f4ZAAC4PPZKPP7nP//ZfP7znx8Caxml8pa3vGV4+Ch5wHnHO94xPNTntyY+/vGPbx566KHhvW2fPaQ8BCe40D8w5wEpgcQ+OdD3Aq4H4jyM1bLGHvTrIa0e2qZ6gI7p13nRJPGch/3+q41Sjn2QI+XVByEShNm393eOWYJzCXAk6NAvu+TBOj1zlwR8DqGvQ6mb2b65XvRVZ/d1WvVsl/MkxyN1I0GLudGLCa4kUbv294BK1lVBuaw7Qbddk5g5ZmlXEshJe1eSYMm2Tsm8SfxkvX3wJ/UvgcdKRKVdbBPUF0kfYEvdzzkwFRCbks8sCXJuk+BdAnXpeV8ja2skZvs7YHWuZkRQRm1m3akHc23U1DGvZbU9/FMXM18bLC5L27PU7yTGEsx+wxvecPzq/+zSHu9ibp0pu9TrsTo/5rTb413kuGS7cl5OjdxM8iblOZeIhjk5v/uv24ycnzlPv/SlL131jQjVbtR1NNfazNd+e0YlxKfUebptNOVZtLRtzD4mqZp7irTZU1+VmnM8bdChyiHbl4RctEmd2u5cH3LdWXofUvcZ7X1d/s0+tdeTup9Iu9xbWmbb2uNaTu5Zbr755uFa2nY6G5N6mgTpVDIz9wPZ/qnyyPbuct2fu++63uq8TCeW/rq8zS7n+ZIyqzqw9Pcmq5PhrtsNAABcHnslHv/0pz8ND6Lln//85/DQ/O9//3v4+wc/+MFVAcyM/quHmG2fPaTqHd73AM6U1/bRjnzoZV15CO9HSFxmCWr0vwGZKQGAk5T6l9E3CcycxQD2Njlf2pFXmfats9fDtvMkEgC69957h4RF6kP2MfvcBzwShEoALYG2BLYy31ggbakEsyowl+RSEqX9yLYlUq8qAVXH6O9///uwP1P7n+DPVK/77FPKIftaSa+cN2NJKP4XvEww8VAjQhPkb5eVf/N3Anv56syowHWSwrmu5dqVbRirt2XumPeSrEr9SZ3MtaQsbc+++tWvDkHHqYT10va4DSa3UyU2WtvWWb+ded4TcTk/c1zSZk0Fr1MH0hkh5ZlvPqhyW9Necbmkjdk2iqgfTVvtRn01Zd7LN26kbao6mA4UGaU71sErztNox9bStjFteMooibi5Tl6HlrY0nabSPra/hxv1Fae7yP1A9iHfkLBtP+qbFFKf2vubQ94fp462Cc9s29x9S64hqds5Fu3vTC61z3V/l2vwacoxyXmZ8y7XiF2TokvP86VlluWlPqas6r526t449TrXwxzLs5bMBQAAzo4L/RuPGXGVgGcevOuhuJ32HUk3JevK8pJ0mPuqxcskPcvTkzdB+r78D9mjvORBOw/NecDuA0xJeOW9qZ7dZ0WC1+lxnKRWW15Tv5Ny3uQ8KW3QKkmcGOtt3SYLUw5Jwh4imJ/198mlXSRoV9uV6TWvec1Qz8a+Ti37lODPWKAmQbrsU5KO2dcql+xrP6qSR3r1H+rrLNs6OSXteUbpR0Yc5Ri2gbqMaOnb/LljPqauH6lD6TiztD2rUcFJKraJsep8U/VxaXvcnpft1I5AWbLOQyeHT1sdj+r00d4zJHifdqPqTo596kDa7RyTT37yk0OblrJO4H8saQu9uVFE1R4skXM553Sdu0m+pGPM2D1Q2tPzNtpxadtY+5qfNEiSpm3D0hEj52//7QyHlN8oTzuZ61W77qj7murAtK2TUdqQ7OsuSbu0w2mn0l4tLbM1sk+5b+k7z0Suh/kWnFyDto2KHLPPdX/Xa/BpqfuJXD/7hPQutp3nu5ZZyijLyLLSJmT72m99KPWTIu0IXgAAgN5eicenPvWpwwiP8pjHPGZ4aH70ox89/P3Sl770qofX9H6tB55tnz2kChrumwRsv5Yuy0rv4L7ncPRfX5cH/TyQ5av7LroEbSqoMWbb+4eWh+Z8fVMCM+1XCuX4JxAy9tVOCUykTq5NZFWiO8GsNVKf4ryNPIhdzpNe2ogEBmOuviTZl8DWaST3UyfS63vpqMNKXo71rJ/rdZ86MzYSJW3J3G9yHULKsEa3nZfkSMo5AbEET6unf6YkiBLkzcjY7FNbP7Jvmad/Parc+3Kuv/N+JdT6QHWd93kv87R2HWmRcybnTrVTS9uzBP9SfxJcbYO5Kae0vxUcPGR7vGSdGYGV+tuOCs79QHvs+nPrUO3xoeR3xrKffXvct2ft71imfG+44Yah3rT1A+b053Vv6h6j72AwppY9loA5j6Mdl7aNOXfT5vRJ1Wprc7/Stl+RtmduFPtSWU6uSUn8zI1OzD1N5kmbONXu5fqVY5f9rd+F36bKIvufclhaZq1DtcdZTjqVpc2sjjtTck3ItaG/H9nnur/0Gpz9G7senYRsX42Crc5mU6r8p0Ye9upY1nm+T5mVbFe2b+yeYZefFNkm18kcI9dIAAC4ePZKPD7qUY8afmw+I3HSa/Wee+4ZHlhLHu7vvvvu4bcd05Mzv/X4+Mc/fnhv22cPLQ9rSQKm1/Eu8rBXX5lWD8j14Nr+nsbUfHmISjChfVD87ne/O8x/kVTAPg/3YxLcyMNpXxYnqY55AsElx7+Cwr0EkBMcH0uW7SKBoCS607t6zXFOYDHbUwnIBAQSGKhRN2fd0vMkQaU+yFP7XPUq+56AUHtc8v85VinrJcG3OVlW6mbqaBuIi1pPJJAzFphpZTvTmz//9kG17P9cr/skN3LM2/Mo60/QZ5/f/tlFBa0j69u2n2dBjRTppySkE1TNV9X2I9rrawjHOoXkmpXAdIK/Vdfyb/6ugHWOW65VCdalt3+pr8vrg7bbjnkv5V4jKtt2akl7Nrb9WX/qYzvS4ZDt8ZJ1Zj39Mcq+pE7nc/m7D8gfqj0+lCqzdkRr2q4EcjPl2GY7c+7U32m/jo6OhmBt9iX7nHYd5tQooqkEYNqztCe5F2iTMnWtnUqu5LxM8j/XzP58q7rbJ+bOgyVt41iZVVubc7L/CuiUR9qe9h5sH7kPyDrTzs0lHUvmyXHIZ8Y6zWQfdkk6Rsoi5dOOOD/p++Pc02Ufsr21ztS/pUnH7Hfut6Ku2WXX6/7Sa3DWmY5skWNwktedrGtp0jGq/JM8TBJxzth5vkuZ5di17UrkGOS41T15tOW6tC7OqQ4VJ132AADAdXB0Cfz1r389uv32249uvPHGq6b/Pkgver81Nu/YfA899NDR+973vqvm+9rXvnb87iPLaV+L/J3X8/4u7rzzzqvW1U55bxfZhrHlZOq3N7L/7Txj2z+2fe12Zbm33nrr0a9//evjV9bJcrK8Wtdcmda8u5bTmLnjXu9lyv/P6Y9Byjjb1352rExr6ut2f9zy96HrWXse1Hqn3i/9fo5t09iy+v1Zamz7x7ar1Pxj6+u3PfNOyXvb6ne2o11epr6u9Otsp5p3rLzaae44LKmbY+bW2R7T2se5Ml8j5Ty1D7XuuX3s68fYce/nmZtv7pj3bVSmfGbM0vZsyfbH2D5MrXubpessVVem5jtke3woqS+pN+1+Vh2u9/p6lf2bK5N8flubwPWXYzd1vi0x1zZmatvC1Pm59qlUW1ZT/5l+nXP1rOppux2tpds/1p7VtG89X1L2/Xqn5t9WZq0ch6ltbs/rfqplzpVFpqmyjnz2ta997ZV1Z3lZ7thyMlU7OXacpvZxaZlFzTvWHo/t59g689l2nnYam39bnexNnTdzx7FX9WNbfZszt59Vfn09bKexba0yHtuu/pgv3deYK7N2m6aOz7ZymjtP8l6v35dMU8c/r++yrwAAwPVzQ/5znIO8tNLDMr9hkd6bS3omA1wkNSIho14ySg04WRlZktGh+Xrp8zbS7DLJKKCMxNk2Umutug/NqNlto6Aui9Mqe6jzLyPvdhlVetnUiM18RfH1el527QQAgPNjr69aBeBiSCApXyOZr9Ma+6o1AE5WfmcxTvJrtYFx+bri/JZh+7W0XKu+Drp+NxoAAGCOxCPAJZURJTfffPPw//fee69RJQDXQUYP5Td+tcFwejJ67qabbhp+wzK/u2gE3byMxr7vvvuUEwAAsIjEI8AllWD3Aw88sLnrrrv08gfgTLn//vs3t9xyy5AcuuOOO4YR+nAoSaTlHkjS/2zLzwHcdtttQzuQn0UBAADOB7/xCAAAAAAAAKxmxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACwmsQjAAAAAAAAsJrEIwAAAAAAALCaxCMAAAAAAACw0mbz/77u8dp++lRXAAAAAElFTkSuQmCC)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QUZsLBG8aaj4"
      },
      "source": [
        "It keeps going.. takes 3hrs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y43qpYhRkcNp"
      },
      "source": [
        "# DDP"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GH7ObrLAXNvL"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import torch.distributed as dist\n",
        "\n",
        "def setup(rank, world_size):\n",
        "    os.environ['MASTER_ADDR'] = 'localhost'\n",
        "    os.environ['MASTER_PORT'] = '12355'\n",
        "\n",
        "    # initialize the process group\n",
        "    dist.init_process_group(\"gloo\", rank=rank, world_size=world_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MK5FouO3kfSH"
      },
      "outputs": [],
      "source": [
        "def cleanup():\n",
        "    dist.destroy_process_group()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LNWAk2-IlCyP"
      },
      "outputs": [],
      "source": [
        "from torch.nn.parallel import DistributedDataParallel as DDP\n",
        "from torch.utils.data.distributed import DistributedSampler\n",
        "\n",
        "import logging\n",
        "logging.basicConfig(level = logging.INFO)\n",
        "logger = logging.getLogger()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CkX8ttXDkgye"
      },
      "outputs": [],
      "source": [
        "class Trainer_DDP:\n",
        "    def __init__(self, model, train_dataset, test_dataset, config,rank,world_size):\n",
        "        self.model = model\n",
        "        self.train_dataset = train_dataset\n",
        "        self.test_dataset = test_dataset\n",
        "        self.config = config\n",
        "        self.rank=rank\n",
        "        self.world_size=world_size\n",
        "\n",
        "        # take over whatever gpus are on the system\n",
        "        self.device = \"cpu\"\n",
        "        if torch.cuda.is_available():\n",
        "            self.device = torch.cuda.current_device()\n",
        "            self.model = torch.nn.DataParallel(self.model).to(self.device)\n",
        "\n",
        "    def save_checkpoint(self):\n",
        "        # DataParallel wrappers keep raw model object in .module attribute\n",
        "        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n",
        "        logger.info(f\"saving {self.config.ckpt_path}\")\n",
        "        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n",
        "\n",
        "    def train(self):\n",
        "        model, config = self.model, self.config#getting the model and training configurations\n",
        "        setup(self.rank, self.world_size)\n",
        "        model = model().to(self.rank)#moving the model to device\n",
        "        model = DDP(model, device_ids=[self.rank], output_device=self.rank, find_unused_parameters=True)#wrapping the model with DDP\n",
        "        raw_model = model.module if hasattr(self.model, \"module\") else model\n",
        "        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=config.learning_rate)\n",
        "\n",
        "        def run_epoch(self,split):\n",
        "            print(\"running\")\n",
        "            is_train = split == \"train\"\n",
        "            model.train(is_train)\n",
        "            train_sampler=DistributedSampler(train_dataset, num_replicas=world_size, rank=self.rank, shuffle=False, drop_last=False)#making a distributed sampler for data parallel\n",
        "            data = self.train_dataset if is_train else self.test_dataset\n",
        "\n",
        "            self.loader = DataLoader(\n",
        "                data,\n",
        "                shuffle=True,\n",
        "                pin_memory=True,\n",
        "                batch_size=config.batch_size,\n",
        "                num_workers=config.num_workers,\n",
        "                sampler=train_sampler # adding this as sampler fr distributed data\n",
        "            )\n",
        "\n",
        "            losses = []\n",
        "            pbar = (\n",
        "                tqdm(enumerate(self.loader), total=len(self.loader))\n",
        "                if is_train\n",
        "                else enumerate(self.loader)\n",
        "            )\n",
        "            for it, (x, y) in pbar:\n",
        "\n",
        "                # place data on the correct device\n",
        "                x = x.to(self.device)\n",
        "                y = y.to(self.device)\n",
        "\n",
        "                # forward the model\n",
        "                with torch.set_grad_enabled(is_train):\n",
        "                    loss = model(x, y)  # The forward method returns the mean of the loss directly\n",
        "                    logits = model(x)\n",
        "                    losses.append(loss.item())\n",
        "\n",
        "                if is_train:\n",
        "\n",
        "                    # backprop and update the parameters\n",
        "                    model.zero_grad()\n",
        "                    loss.backward(retain_graph=True)\n",
        "                    torch.nn.utils.clip_grad_norm_(\n",
        "                        model.parameters(), config.grad_norm_clip\n",
        "                    )\n",
        "                    self.optimizer.step()\n",
        "\n",
        "                    # decay the learning rate based on our progress\n",
        "                    if config.lr_decay:\n",
        "                        self.tokens += (\n",
        "                            y >= 0\n",
        "                        ).sum()  # number of tokens processed this step (i.e. label is not -100)\n",
        "                        if self.tokens < config.warmup_tokens:\n",
        "                            # linear warmup\n",
        "                            lr_mult = float(self.tokens) / float(\n",
        "                                max(1, config.warmup_tokens)\n",
        "                            )\n",
        "                        else:\n",
        "                            # cosine learning rate decay\n",
        "                            progress = float(\n",
        "                                self.tokens - config.warmup_tokens\n",
        "                            ) / float(\n",
        "                                max(1, config.final_tokens - config.warmup_tokens)\n",
        "                            )\n",
        "                            lr_mult = max(\n",
        "                                0.1, 0.5 * (1.0 + math.cos(math.pi * progress))\n",
        "                            )\n",
        "                        lr = config.learning_rate * lr_mult\n",
        "                        for param_group in self.optimizer.param_groups:\n",
        "                            param_group[\"lr\"] = lr\n",
        "                    else:\n",
        "                        lr = config.learning_rate\n",
        "\n",
        "                    # repeat progress\n",
        "                    pbar.set_description(\n",
        "                        f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}, lr {lr:e}\"\n",
        "                    )\n",
        "                    cleanup()#cleaning up all the the done files on all GPUs\n",
        "\n",
        "            if not is_train:\n",
        "                test_loss = float(np.mean(losses))\n",
        "                logger.info(f\"test loss: {test_loss}\")\n",
        "                return test_loss\n",
        "\n",
        "        best_loss = float(\"inf\")\n",
        "        self.tokens = 0  # counter used for learning rate decay\n",
        "        for epoch in range(config.max_epochs):\n",
        "            # if we are using DistributedSampler, we have to tell it which epoch this is\n",
        "            self.loader.sampler.set_epoch(epoch)\n",
        "            run_epoch(\"train\")\n",
        "            if self.test_dataset is not None:\n",
        "                test_loss = run_epoch(\"test\")\n",
        "\n",
        "            # supports early stopping based on the test loss, or just save always is no test set is provided\n",
        "            good_model = self.test_dataset is None or test_loss < best_loss\n",
        "            if self.config.ckpt_path is not None and good_model:\n",
        "                best_loss = test_loss if self.test_dataset is not None else float(\"inf\")\n",
        "                self.save_checkpoint()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L4mkK8PikjMd"
      },
      "outputs": [],
      "source": [
        "def init_process(rank,model,train_dataset,test_dataset,config, world_size):\n",
        "    torch.cuda.set_device(rank)\n",
        "    setup(rank, world_size)\n",
        "\n",
        "    # create model, datasets, etc.\n",
        "    trainer = Trainer_DDP(model, train_dataset, test_dataset, config, rank, world_size)\n",
        "    trainer.train()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nNtoEilYkpo5"
      },
      "outputs": [],
      "source": [
        "trainable_model = GPT2()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0KxXZM52nNZd"
      },
      "outputs": [],
      "source": [
        "tconf = TrainerConfig(\n",
        "    max_epochs=1,\n",
        "    batch_size=8,\n",
        "    learning_rate=6e-4,\n",
        "    lr_decay=True,\n",
        "    warmup_tokens=512,\n",
        "    final_tokens=2*len(train_dataset)*block_size,\n",
        "    num_workers=4,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Uw7p2QAuQcP",
        "outputId": "fac4c969-a2b6-4f90-fe6c-ef9ee9735de9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Error: This script requires at least 2 GPUs for multiprocessing.\n"
          ]
        }
      ],
      "source": [
        "import torch.multiprocessing as mp\n",
        "\n",
        "# Check the number of available GPUs\n",
        "num_gpus = torch.cuda.device_count()\n",
        "\n",
        "# Ensure that there are at least 2 GPUs available\n",
        "if num_gpus < 2:\n",
        "    print(\"Error: This script requires at least 2 GPUs for multiprocessing.\")\n",
        "    exit()\n",
        "\n",
        "# Set the world size based on the number of GPUs\n",
        "world_size = num_gpus\n",
        "\n",
        "mp.spawn(\n",
        "    init_process,\n",
        "    args=(trainable_model, train_dataset, None, tconf, world_size),\n",
        "    nprocs=world_size,\n",
        "    join=True\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P8-7nhacoMlB"
      },
      "source": [
        "# FSDP\n",
        "FSDP, or Fully Sharded Data Parallel, optimizes model training by combining model sharding and data parallelism. Unlike DDP, it allows efficient computation by wrapping model layers in a nested way, reducing the need to gather parameters across devices during forward and backward passes. This results in enhanced model efficiency and optimized training.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "abuZKQxQoYiq"
      },
      "outputs": [],
      "source": [
        "from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n",
        "from torch.distributed.fsdp.wrap import (\n",
        "    size_based_auto_wrap_policy,\n",
        "    enable_wrap,\n",
        "    wrap,\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BwoU-gyanWdC"
      },
      "outputs": [],
      "source": [
        "class Trainer_FSDP:\n",
        "    def __init__(self, model, train_dataset, test_dataset, config, rank, world_size):\n",
        "        self.model = model\n",
        "        self.train_dataset = train_dataset\n",
        "        self.test_dataset = test_dataset\n",
        "        self.config = config\n",
        "        self.rank = rank\n",
        "        self.world_size = world_size\n",
        "\n",
        "        # Use available GPUs if present\n",
        "        self.device = \"cpu\"\n",
        "        if torch.cuda.is_available():\n",
        "            self.device = torch.cuda.current_device()\n",
        "            self.model = torch.nn.DataParallel(self.model).to(self.device)\n",
        "\n",
        "    def save_checkpoint(self):\n",
        "        # DataParallel wrappers keep raw model object in .module attribute\n",
        "        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n",
        "        logger.info(f\"saving {self.config.ckpt_path}\")\n",
        "        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n",
        "\n",
        "    def setup_ddp(self):\n",
        "        # Distributed Data Parallel setup\n",
        "        setup(self.rank, self.world_size)\n",
        "        self.model = self.model().to(self.rank)\n",
        "        self.model = DDP(self.model, device_ids=[self.rank], output_device=self.rank, find_unused_parameters=True)\n",
        "\n",
        "    def initialize_optimizer(self):\n",
        "        # Initialize optimizer\n",
        "        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config.learning_rate)\n",
        "\n",
        "    def run_epoch(self, split):\n",
        "        print(\"running\")\n",
        "        is_train = split == \"train\"\n",
        "        model.train(is_train)\n",
        "        train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=self.rank, shuffle=False, drop_last=False)\n",
        "        data = self.train_dataset if is_train else self.test_dataset\n",
        "\n",
        "        self.loader = DataLoader(\n",
        "            data,\n",
        "            shuffle=True,\n",
        "            pin_memory=True,\n",
        "            batch_size=self.config.batch_size,\n",
        "            num_workers=self.config.num_workers,\n",
        "            sampler=train_sampler\n",
        "        )\n",
        "        # Set up auto wrap policy for our model\n",
        "        my_auto_wrap_policy = functools.partial(\n",
        "            size_based_auto_wrap_policy, min_num_params=100\n",
        "        )\n",
        "        torch.cuda.set_device(rank)\n",
        "\n",
        "        # Starting and ending for our cuda event denoting when to allocate and free the memory\n",
        "        init_start_event = torch.cuda.Event(enable_timing=True)\n",
        "        init_end_event = torch.cuda.Event(enable_timing=True)\n",
        "\n",
        "        # Setting up FSDP on our model\n",
        "        model = FSDP(self.model)\n",
        "\n",
        "        losses = []\n",
        "        pbar = tqdm(enumerate(self.loader), total=len(self.loader)) if is_train else enumerate(self.loader)\n",
        "        for it, (x, y) in pbar:\n",
        "            # Place data on the correct device\n",
        "            x = x.to(self.device)\n",
        "            y = y.to(self.device)\n",
        "\n",
        "            # Forward the model\n",
        "            with torch.set_grad_enabled(is_train):\n",
        "                loss = model(x, y)\n",
        "                logits = model(x)\n",
        "                losses.append(loss.item())\n",
        "\n",
        "            if is_train:\n",
        "                # Backprop and update the parameters\n",
        "                model.zero_grad()\n",
        "                loss.backward(retain_graph=True)\n",
        "                torch.nn.utils.clip_grad_norm_(model.parameters(), self.config.grad_norm_clip)\n",
        "                self.optimizer.step()\n",
        "\n",
        "                # Decay the learning rate based on our progress\n",
        "                if self.config.lr_decay:\n",
        "                    self.tokens += (y >= 0).sum()\n",
        "                    if self.tokens < self.config.warmup_tokens:\n",
        "                        lr_mult = float(self.tokens) / float(max(1, self.config.warmup_tokens))\n",
        "                    else:\n",
        "                        progress = float(self.tokens - self.config.warmup_tokens) / float(\n",
        "                            max(1, self.config.final_tokens - self.config.warmup_tokens)\n",
        "                        )\n",
        "                        lr_mult = max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))\n",
        "                    lr = self.config.learning_rate * lr_mult\n",
        "                    for param_group in self.optimizer.param_groups:\n",
        "                        param_group[\"lr\"] = lr\n",
        "                else:\n",
        "                    lr = self.config.learning_rate\n",
        "\n",
        "                pbar.set_description(f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}, lr {lr:e}\")\n",
        "                cleanup()\n",
        "\n",
        "        if not is_train:\n",
        "            test_loss = float(np.mean(losses))\n",
        "            logger.info(f\"test loss: {test_loss}\")\n",
        "            return test_loss\n",
        "\n",
        "    def train(self):\n",
        "        self.setup_ddp()\n",
        "        self.initialize_optimizer()\n",
        "        best_loss = float(\"inf\")\n",
        "        self.tokens = 0  # Counter used for learning rate decay\n",
        "        for epoch in range(self.config.max_epochs):\n",
        "            # If we are using DistributedSampler, we have to tell it which epoch this is\n",
        "            self.loader.sampler.set_epoch(epoch)\n",
        "            self.run_epoch(\"train\")\n",
        "            if self.test_dataset is not None:\n",
        "                test_loss = self.run_epoch(\"test\")\n",
        "\n",
        "            # Supports early stopping based on the test loss, or just save always if no test set is provided\n",
        "            good_model = self.test_dataset is None or test_loss < best_loss\n",
        "            if self.config.ckpt_path is not None and good_model:\n",
        "                best_loss = test_loss if self.test_dataset is not None else float(\"inf\")\n",
        "                self.save_checkpoint()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "swCn_0tWoPnD"
      },
      "outputs": [],
      "source": [
        "trainable_model = GPT2()#initialsing our model for training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "42eeLaTdoTVC"
      },
      "outputs": [],
      "source": [
        "tconf = TrainerConfig(\n",
        "    max_epochs=1,\n",
        "    batch_size=8,\n",
        "    learning_rate=6e-4,\n",
        "    lr_decay=True,\n",
        "    warmup_tokens=512,\n",
        "    final_tokens=2*len(train_dataset)*block_size,\n",
        "    num_workers=4,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xVnnNbEGoVNG",
        "outputId": "2682b759-5d0a-4e0a-aa61-d2c4fedc5f57"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Error: This script requires at least 2 GPUs for parallel training.\n"
          ]
        }
      ],
      "source": [
        "import torch.multiprocessing as mp\n",
        "\n",
        "# Check the number of available GPUs\n",
        "num_gpus = torch.cuda.device_count()\n",
        "\n",
        "# Ensure that there are at least 2 GPUs available\n",
        "if num_gpus < 2:\n",
        "    print(\"Error: This script requires at least 2 GPUs for parallel training.\")\n",
        "    exit()\n",
        "\n",
        "# Set the world size based on the number of GPUs\n",
        "world_size = num_gpus\n",
        "\n",
        "# Spawn processes for parallel training\n",
        "mp.spawn(\n",
        "    init_process,\n",
        "    args=(trainable_model, train_dataset, None, tconf, world_size),\n",
        "    nprocs=world_size,\n",
        "    join=True\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "QXJlRs9Z3Lxw"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}