Baby_Llama_361.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OnipH4BHgpu1",
        "outputId": "c76968c7-aff3-463a-e6e1-1e3d3b2daac9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'llama2.c'...\n",
            "remote: Enumerating objects: 1434, done.\u001b[K\n",
            "remote: Total 1434 (delta 0), reused 0 (delta 0), pack-reused 1434\u001b[K\n",
            "Receiving objects: 100% (1434/1434), 1.17 MiB | 3.00 MiB/s, done.\n",
            "Resolving deltas: 100% (871/871), done.\n"
          ]
        }
      ],
      "source": [
        "!git clone https://github.com/karpathy/llama2.c.git"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install sentencepiece"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zRHKP8PpgxtI",
        "outputId": "ae4eef27-6042-44a2-f818-16fd8e14bd27"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting sentencepiece\n",
            "  Downloading sentencepiece-0.1.99-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m8.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: sentencepiece\n",
            "Successfully installed sentencepiece-0.1.99\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# runtime 100s\n",
        "!cd llama2.c && python tinystories.py download"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mWzFVICBv1sY",
        "outputId": "b81d1b50-f2ce-46eb-bed4-43096238a3c3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories_all_data.tar.gz to data/TinyStories_all_data.tar.gz...\n",
            "data/TinyStories_all_data.tar.gz: 100% 1.50G/1.50G [00:22<00:00, 71.3MiB/s]\n",
            "Unpacking data/TinyStories_all_data.tar.gz...\n",
            "Download done.\n",
            "Number of shards: 50\n",
            "Example story:\n",
            "{'story': '\\n\\nLily and Ben are friends. They like to play in the park. One day, they see a big tree with a swing. Lily wants to try the swing. She runs to the tree and climbs on the swing.\\n\"Push me, Ben!\" she says. Ben pushes her gently. Lily feels happy. She swings higher and higher. She laughs and shouts.\\nBen watches Lily. He thinks she is cute. He wants to swing too. He waits for Lily to stop. But Lily does not stop. She swings faster and faster. She is having too much fun.\\n\"Can I swing too, Lily?\" Ben asks. Lily does not hear him. She is too busy swinging. Ben feels sad. He walks away.\\nLily swings so high that she loses her grip. She falls off the swing. She lands on the ground. She hurts her foot. She cries.\\n\"Ow, ow, ow!\" she says. She looks for Ben. She wants him to help her. But Ben is not there. He is gone.\\nLily feels sorry. She wishes she had shared the swing with Ben. She wishes he was there to hug her. She limps to the tree. She sees something hanging from a branch. It is Ben\\'s hat. He left it for her.\\nLily smiles. She thinks Ben is nice. She puts on his hat. She hopes he will come back. She wants to say sorry. She wants to be friends again.', 'instruction': {'prompt:': 'Write a short story (3-5 paragraphs) which only uses very simple words that a 3 year old child would understand. The story should use the verb \"hang\", the noun \"foot\" and the adjective \"cute\". The story has the following features: the story should contain at least one dialogue. Remember to only use simple words!\\n\\nPossible story:', 'words': ['hang', 'foot', 'cute'], 'features': ['Dialogue']}, 'summary': 'Lily and Ben play in the park and Lily gets too caught up in swinging, causing Ben to leave. Lily falls off the swing and hurts herself, but Ben leaves his hat for her as a kind gesture.', 'source': 'GPT-4'}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#!cd llama2.c && python tinystories.py train_vocab --vocab_size=361\n",
        "!cd ./llama2.c/data && wget https://huggingface.co/datasets/enio/TinyStories/resolve/main/tok361/tok361.vocab\n",
        "!cd ./llama2.c/data && wget https://huggingface.co/datasets/enio/TinyStories/resolve/main/tok361/tok361.model\n",
        "\n",
        "#!cd llama2.c && python tinystories.py pretokenize --vocab_size=361\n",
        "!cd ./llama2.c/data && wget https://huggingface.co/datasets/enio/TinyStories/resolve/main/tok361/tok361.tar.gz\n",
        "!cd ./llama2.c/data && tar -xf tok361.tar.gz"
      ],
      "metadata": {
        "id": "Tx04_T-pv1zf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# change bfloat16 to float16 on line 72\n",
        "!sed -i \"72s/.*/dtype = 'float16'  # float32|bfloat16|float16/\" ./llama2.c/train.py"
      ],
      "metadata": {
        "id": "tDQu756djGVD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# runtime ~ 2 mins\n",
        "!cd llama2.c && python train.py --vocab_source=custom --vocab_size=361 --compile=False \\\n",
        "  --dim=128 --n_layers=5 --n_heads=8 --n_kv_heads=4 --batch_size=32 \\\n",
        "  --always_save_checkpoint=True --eval_interval=1 --max_iters=1 #--init_from='resume'"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sdMK1MW8g4eM",
        "outputId": "07b008fe-b4d7-4383-a4b1-1896667a70bd"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overriding: vocab_source = custom\n",
            "Overriding: vocab_size = 361\n",
            "Overriding: compile = False\n",
            "Overriding: dim = 128\n",
            "Overriding: n_layers = 5\n",
            "Overriding: n_heads = 8\n",
            "Overriding: n_kv_heads = 4\n",
            "Overriding: batch_size = 32\n",
            "Overriding: always_save_checkpoint = True\n",
            "Overriding: eval_interval = 1\n",
            "Overriding: max_iters = 1\n",
            "tokens per iteration will be: 32,768\n",
            "breaks down as: 4 grad accum steps * 1 processes * 32 batch size * 256 max seq len\n",
            "Initializing a new model from scratch\n",
            "num decayed parameter tensors: 36, with 967,808 parameters\n",
            "num non-decayed parameter tensors: 11, with 1,408 parameters\n",
            "using fused AdamW: True\n",
            "Created a PretokDataset with rng seed 42\n",
            "Created a PretokDataset with rng seed 42\n",
            "Created a PretokDataset with rng seed 42\n",
            "step 0: train loss 5.9283, val loss 5.9287\n",
            "0 | loss 5.9330 | lr 0.000000e+00 | 3368.74ms | mfu -100.00%\n",
            "Created a PretokDataset with rng seed 42\n",
            "Created a PretokDataset with rng seed 42\n",
            "step 1: train loss 5.9283, val loss 5.9287\n",
            "saving checkpoint to out\n",
            "wrote out/model.bin\n",
            "1 | loss 5.9234 | lr 5.000000e-07 | 2376.90ms | mfu -100.00%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cd llama2.c && gcc -O3 -o run run.c -lm\n",
        "!cd llama2.c && python tokenizer.py --tokenizer-model=data/tok361.model"
      ],
      "metadata": {
        "id": "oUpbYitgNdMh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!cd ./llama2.c && ./run out/model.bin -z data/tok361.bin -t 0.8 -n 256 -i \"Once upon a time \""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1-ylANKlOfCr",
        "outputId": "395a6b54-c210-4ce3-85d7-71caea890146"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Once upon a time E<unk>II9M g——l-L-&‘ppe`6)ib wJ!Dn~ ‘!-<unk>L'.0rNKélVU\"\"â™cggmkzIEoQ\\?`LL4XXJ`7]11\\/L??!p9Eqs\n",
            "</s>\n",
            "%VV.;TEfDD90zz88aW5#S,\n",
            "achieved tok/s: 812.734082\n"
          ]
        }
      ]
    }
  ]
}