mlx-usft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "985e53bd-c394-4786-9bc2-4ac82d11aec8",
   "metadata": {},
   "source": [
    "# Text Completion / Raw Text Training\n",
    "### This is a notebook by [Mark Lord](https://twitter.com/priontific); Communications Lead of [ValleyDAO](https://valleydao.bio) and leader of [The Curriculum Protocol](https://discord.gg/U8TZ2p4z).\n",
    "#### It borrows heavily from the Unsloth team's [fantastic Colab notebooks](https://colab.research.google.com/drive/1ef-tab5bhkvWmBOObepl1WgJvfvSzn5Q?usp=sharing#scrollTo=IqM-T1RTzY6C)!\n",
    "\n",
    "#### This Notebook is designed for training an LLM with the default loss function of the MLX Lora.py script. This means you can finetune on any dataset and let your model act as a text completion model, like for novel writing.\n",
    "----\n",
    "#### We'll be training on [Tiny Stories](https://huggingface.co/datasets/roneneldan/TinyStories) which is a collection of small stories. For example:\n",
    "\n",
    "`Once upon a time, there was a little car named Beep. Beep loved to go fast and play in the sun.\n",
    "Beep was a healthy car because he always had good fuel....`\n",
    "\n",
    "#### To run this Notebook, first make sure you've installed everything in the requirements.txt file (recommended that you use a virtualenv), then press \"Run\" in the top left and select \"Run All Cells\".\n",
    "#####"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "12d2f159-4a78-4ba9-a0bd-796c8db19934",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remember to install all requirements. Strongly recommended that you set up a virtual environment!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ce9b7bb-4121-4f66-a847-72ab06655f65",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Set your model here 👇\n",
    "\n",
    "#### (You can also download pre-quantised models from https://huggingface.co/mlx-community to save on bandwidth!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "48e4f2d7-e0e2-42a0-9c1b-1b483adfaa76",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Model set to TinyLlama/TinyLlama-1.1B-Chat-v1.0.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Set the github path to the model. If you've already got a model on disk, you'll have to edit this notebook to point to that instead.\n",
    "\n",
    "github_path = \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\"\n",
    "print(f\"\\nModel set to {github_path}.\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb60e26-7565-4834-ada6-5324dc87e672",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Now download it and convert it to 4-bit.\n",
    "#### (Unless you successfully downloaded a pre-quantised model, in which case **you do not need to run this step.**)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8d509e7e-1d9d-4a14-b42a-41f793525135",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Loading\n",
      "Fetching 8 files: 100%|████████████████████████| 8/8 [00:00<00:00, 18236.10it/s]\n",
      "[INFO] Quantizing\n",
      "\n",
      "Model converted.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Convert a model of choice. The -q flag quantizes the model to 4-bit by default.\n",
    "\n",
    "!python3 -m mlx_lm.convert \\\n",
    "--hf-path {github_path} \\\n",
    "--mlx-path {github_path} \\\n",
    "-q\n",
    "\n",
    "print(f\"\\nModel converted.\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd9a98e-94c0-4689-be35-c7fc1c1d708c",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "### By default, MLX determines how many entries are seen using iters and batches rather than epochs. \n",
    "\n",
    "#### The code for this is extremely simple, and you can run it just using the configuration in the code block below.\n",
    "\n",
    "#### Training in epochs however is more desirable for me, as I'd prefer to make sure that every example is seen as many times as every other, so the block below is disabled. If you would rather use iters and batch size, un-comment out the block below and put it into the later epoch-based training code instead.\n",
    "\n",
    "----\n",
    "# Set training parameters here 👇\n",
    "#### The TinyStories' dataset is a 1GB install; depending on your internet speed, you might be here a while."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "9e447afd-7e7b-4e30-93d9-1a715f0531ac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Repo card metadata block was not found. Setting CardData to empty.\n",
      "Repo card metadata block was not found. Setting CardData to empty.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Datasets loaded.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"roneneldan/TinyStories\", split = \"train[:500]\")\n",
    "valid_dataset = load_dataset(\"roneneldan/TinyStories\", split=\"train[501:601]\")\n",
    "\n",
    "print(f\"\\nDatasets loaded.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "85b1f80f-5b2c-4796-8b4a-09e92518b154",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "EOS token added to each entry.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(github_path)\n",
    "EOS_TOKEN = tokenizer.eos_token\n",
    "\n",
    "# Set the path for your modified training file\n",
    "data_dir = './data'  # Directory for the modified training file\n",
    "if not os.path.exists(data_dir):\n",
    "    os.makedirs(data_dir)  # Create the data directory if it doesn't exist\n",
    "\n",
    "modified_train_file = f'{data_dir}/train.jsonl'\n",
    "valid_file = f'{data_dir}/valid.jsonl'\n",
    "\n",
    "# Open a new file and write each story with the EOS token appended\n",
    "with open(modified_train_file, 'w') as file:\n",
    "    for item in dataset:\n",
    "        story_with_eos = item['text'] + EOS_TOKEN  # Append the EOS token to each story\n",
    "        # Write this story as a JSON object to the file\n",
    "        file.write(json.dumps({\"text\": story_with_eos}) + '\\n')\n",
    "\n",
    "# Same but for the validation file, though not strictly necessary\n",
    "with open(valid_file, 'w') as file:\n",
    "    for item in dataset:\n",
    "        story_with_eos = item['text'] + EOS_TOKEN  # Append the EOS token to each story\n",
    "        # Write this story as a JSON object to the file\n",
    "        file.write(json.dumps({\"text\": story_with_eos}) + '\\n')\n",
    "\n",
    "print(f\"\\nEOS token added to each entry.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "828f3d91-df98-44cc-9304-0a8cec1f189e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Below are your entries with the EOS token appended.\n",
      "\n",
      "=========================\n",
      "One day, a little girl named Lily found a needle in her room. She knew it was difficult to play with it because it was sharp. Lily wanted to share the needle with her mom, so she could sew a button on her shirt.\n",
      "\n",
      "Lily went to her mom and said, \"Mom, I found this needle. Can you share it with me and sew my shirt?\" Her mom smiled and said, \"Yes, Lily, we can share the needle and fix your shirt.\"\n",
      "\n",
      "Together, they shared the needle and sewed the button on Lily's shirt. It was not difficult for them because they were sharing and helping each other. After they finished, Lily thanked her mom for sharing the needle and fixing her shirt. They both felt happy because they had shared and worked together.</s>\n",
      "=========================\n",
      "Once upon a time, there was a little car named Beep. Beep loved to go fast and play in the sun. Beep was a healthy car because he always had good fuel. Good fuel made Beep happy and strong.\n",
      "\n",
      "One day, Beep was driving in the park when he saw a big tree. The tree had many leaves that were falling. Beep liked how the leaves fall and wanted to play with them. Beep drove under the tree and watched the leaves fall on him. He laughed and beeped his horn.\n",
      "\n",
      "Beep played with the falling leaves all day. When it was time to go home, Beep knew he needed more fuel. He went to the fuel place and got more healthy fuel. Now, Beep was ready to go fast and play again the next day. And Beep lived happily ever after.</s>\n",
      "=========================\n",
      "One day, a little fish named Fin was swimming near the shore. He saw a big crab and wanted to be friends. \"Hi, I am Fin. Do you want to play?\" asked the little fish. The crab looked at Fin and said, \"No, I don't want to play. I am cold and I don't feel fine.\"\n",
      "\n",
      "Fin felt sad but wanted to help the crab feel better. He swam away and thought of a plan. He remembered that the sun could make things warm. So, Fin swam to the top of the water and called to the sun, \"Please, sun, help my new friend feel fine and not freeze!\"\n",
      "\n",
      "The sun heard Fin's call and shone its warm light on the shore. The crab started to feel better and not so cold. He saw Fin and said, \"Thank you, little fish, for making me feel fine. I don't feel like I will freeze now. Let's play together!\" And so, Fin and the crab played and became good friends.</s>\n",
      "=========================\n",
      "Once upon a time, in a land full of trees, there was a little cherry tree. The cherry tree was very sad because it did not have any friends. All the other trees were big and strong, but the cherry tree was small and weak. The cherry tree was envious of the big trees.\n",
      "\n",
      "One day, the cherry tree felt a tickle in its branches. It was a little spring wind. The wind told the cherry tree not to be sad. The wind said, \"You are special because you have sweet cherries that everyone loves.\" The cherry tree started to feel a little better.\n",
      "\n",
      "As time went on, the cherry tree grew more and more cherries. All the animals in the land came to eat the cherries and play under the cherry tree. The cherry tree was happy because it had many friends now. The cherry tree learned that being different can be a good thing. And they all lived happily ever after.</s>\n",
      "=========================\n",
      "Once upon a time, there was a little girl named Lily. Lily liked to pretend she was a popular princess. She lived in a big castle with her best friends, a cat and a dog.\n",
      "\n",
      "One day, while playing in the castle, Lily found a big cobweb. The cobweb was in the way of her fun game. She wanted to get rid of it, but she was scared of the spider that lived there.\n",
      "\n",
      "Lily asked her friends, the cat and the dog, to help her. They all worked together to clean the cobweb. The spider was sad, but it found a new home outside. Lily, the cat, and the dog were happy they could play without the cobweb in the way. And they all lived happily ever after.</s>\n"
     ]
    }
   ],
   "source": [
    "print(f\"\\nBelow are your entries with the EOS token appended.\\n\")\n",
    "\n",
    "with open(modified_train_file, 'r') as file:\n",
    "    for i, line in enumerate(file):\n",
    "        if i >= 5:  # Only read the first 5 entries\n",
    "            break\n",
    "        entry = json.loads(line)  # Parse the JSON string into a Python dictionary\n",
    "        print(\"=========================\")\n",
    "        print(entry[\"text\"])  # Print the \"text\" field of the entry\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "88333717-302d-491f-b389-e8ea5761283f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Automatically detected 5000 data entries.\n",
      "For 1 epoch(s) with a batch size of 1, we will set iters to: 5000\n",
      "Total number of tokens in the JSONL file: 1126679\n",
      "Estimated training rate in tokens/second if fits in GPU: 734.0\n",
      "\n",
      "If model fits in GPU: Estimated time for 1 epoch(s) with 22 LoRA layer(s) with a token amount of 1126679: \n",
      "25 minutes and 34 seconds\n",
      "\n",
      "Else if model doesn't fit in GPU, could be up to:\n",
      "178 minutes and 58 seconds\n",
      "\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1                    # Change this as needed; but more batches takes more memory, and my tests show no increases in speed.\n",
    "epochs = 1                        # For newbies, this means how many times the model gets to see the dataset. 1 means it goes through once and that's it.\n",
    "train_set = \"./data\"              # Change this to point to the folder containing your train.JSONL file.\n",
    "lora_layers = 22                  # Set the number of layers e.g. 4, 16, 32. Less layers = less effect on style, but saves on RAM.\n",
    "context_length = 1024              # Set the context length during training. More context takes more memory.\n",
    "learning_rate = 2e-5              # Set the learning rate. Higher e.g. 2e-5 can make your model go a bit nuts. Lower e.g. 2e-6 and it'll learn real slow.\n",
    "\n",
    "\n",
    "'''MLX doesn't have a native flag for training epochs, so I've implemented some code below to convert iters to epochs.'''\n",
    "\n",
    "train_file = \"./data/train.jsonl\"\n",
    "iters_per_epoch = -(-sum(1 for _ in open(train_file)) // batch_size)\n",
    "iters = iters_per_epoch * epochs # Swap this out to iters = 100 or whatever you like if you'd prefer to train in iters.\n",
    "\n",
    "'''The code below counts the number of characters in your training dataset to figure out how long training will take. It's not high accuracy.'''\n",
    "import sentencepiece as spm\n",
    "import json\n",
    "\n",
    "# Use the tokenizer to split your dataset into tokens:\n",
    "model_path = f\"{github_path}/tokenizer.model\"\n",
    "sp_processor = spm.SentencePieceProcessor()\n",
    "sp_processor.load(model_path)\n",
    "\n",
    "def count_tokens_in_file(sp_processor, file_path):\n",
    "    total_tokens = 0\n",
    "    with open(file_path, 'r') as f:\n",
    "        for line in f:\n",
    "            entry = json.loads(line)\n",
    "            text = entry.get('text', '')  # Make sure 'text' is the correct key for your JSONL entries\n",
    "            tokens = sp_processor.encode_as_pieces(text)\n",
    "            total_tokens += len(tokens)\n",
    "    return total_tokens\n",
    "\n",
    "# Use the correct variable for the SentencePiece processor here\n",
    "total_tokens = count_tokens_in_file(sp_processor, train_file) * epochs\n",
    "multiplier_for_layers = 1/(1+(((32-lora_layers)/32)*1.5)) # Really approximate maths to get a multiplier based on number of layers.\n",
    "training_rate = 500 // multiplier_for_layers\n",
    "estimated_total_time = int(total_tokens // training_rate)\n",
    "estimated_minutes = int(estimated_total_time // 60)\n",
    "estimated_seconds = int(estimated_total_time % 60)\n",
    "slow_time = estimated_total_time * 7\n",
    "slow_minutes = int(slow_time // 60)\n",
    "slow_seconds = int(slow_time % 60)\n",
    "\n",
    "print(f\"\\nAutomatically detected {iters_per_epoch} data entries.\")\n",
    "print(f\"For {epochs} epoch(s) with a batch size of {batch_size}, we will set iters to: {iters}\")\n",
    "print(f\"Total number of tokens in the JSONL file: {total_tokens}\")\n",
    "print(f\"Estimated training rate in tokens/second if fits in GPU: {training_rate}\")\n",
    "print(f\"\\nIf model fits in GPU: Estimated time for {epochs} epoch(s) with {lora_layers} LoRA layer(s) with a token amount of {total_tokens}: \\n{estimated_minutes} minutes and {estimated_seconds} seconds\")\n",
    "print(f\"\\nElse if model doesn't fit in GPU, could be up to:\\n{slow_minutes} minutes and {slow_seconds} seconds\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd381122-f7cf-40fd-9cd8-08e9392e0d63",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "#### Remember, the (very roughly) estimated time 👆 is for an M1 Max - you'll have to scale this based on the processor you have.\n",
    "-----\n",
    "#### To work out how long it'll take your machine, you can apply a rough calculation:\n",
    "\n",
    "##### M(1/2/3) Pro: 2x as long (double the estimated time)\n",
    "##### M(1/2/3) Base: 4x as long (quadruple the estimated time)\n",
    "##### M(1/2/3) Ultra: 1/2 as long (halve the estimated time)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd16a079-6d24-41b9-8d46-cd630d33248a",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "### The cell below 👇 will train and produce your LoRA adapters as an adapters.npz file.\n",
    "\n",
    "##### Make sure that your data is present as a train.JSONL file in the data folder of this repo. The model will automatically be in the mlx_model folder if you downloaded it using the code at the beginning.\n",
    "\n",
    "###### (By default, a checkpoint will be saved every 100 iterations, though you can change this by adding a '--save-every 100' flag and changing the number.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "7ae03fcd-6883-4507-9d24-73ce204b897a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Adapter file set to trial1.npz.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "adapters = \"trial1.npz\"\n",
    "print(f\"\\nAdapter file set to {adapters}.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "b6e0ad10-57f3-4202-950d-cb28d1f2a474",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Any previous adapter file trial1.npz cleared (and all checkpoints deleted).\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Uncomment out this cell to delete your adapters if your model generation has gone hay-wire.\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"  # or \"true\" if you want to enforce parallelism\n",
    "\n",
    "!rm -f {adapters}\n",
    "!rm -rf checkpoints  # Be careful: This will delete the directory and all its contents!\n",
    "!mkdir checkpoints\n",
    "print(f\"\\nAny previous adapter file {adapters} cleared (and all checkpoints deleted).\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d1b55bb4-ef7d-4648-bcf9-40c55e4991de",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained model\n",
      "Total parameters 228.383M\n",
      "Trainable parameters 1.126M\n",
      "Loading datasets\n",
      "Training\n",
      "Starting training..., iters: 5000\n",
      "Iter 1: Val loss 1.677, Val took 3.820s\n",
      "Iter 10: Train loss 1.566, Learning Rate 2.000e-05, It/sec 4.018, Tokens/sec 817.344, Trained Tokens 2034\n",
      "Iter 20: Train loss 1.569, Learning Rate 2.000e-05, It/sec 3.488, Tokens/sec 732.439, Trained Tokens 4134\n",
      "Iter 30: Train loss 1.611, Learning Rate 2.000e-05, It/sec 3.616, Tokens/sec 726.731, Trained Tokens 6144\n",
      "Iter 40: Train loss 1.494, Learning Rate 2.000e-05, It/sec 3.445, Tokens/sec 732.391, Trained Tokens 8270\n",
      "Iter 50: Train loss 1.457, Learning Rate 2.000e-05, It/sec 3.781, Tokens/sec 717.663, Trained Tokens 10168\n",
      "Iter 60: Train loss 1.504, Learning Rate 2.000e-05, It/sec 3.139, Tokens/sec 753.024, Trained Tokens 12567\n",
      "Iter 70: Train loss 1.442, Learning Rate 2.000e-05, It/sec 2.806, Tokens/sec 755.340, Trained Tokens 15259\n",
      "Iter 80: Train loss 1.474, Learning Rate 2.000e-05, It/sec 3.474, Tokens/sec 731.648, Trained Tokens 17365\n",
      "Iter 90: Train loss 1.649, Learning Rate 2.000e-05, It/sec 3.311, Tokens/sec 625.861, Trained Tokens 19255\n",
      "Iter 100: Train loss 1.388, Learning Rate 2.000e-05, It/sec 3.475, Tokens/sec 739.504, Trained Tokens 21383\n",
      "Iter 100: Saved adapter weights to checkpoints/100_trial1.npz.\n",
      "Iter 110: Train loss 1.363, Learning Rate 2.000e-05, It/sec 2.456, Tokens/sec 733.839, Trained Tokens 24371\n",
      "Iter 120: Train loss 1.442, Learning Rate 2.000e-05, It/sec 3.270, Tokens/sec 707.974, Trained Tokens 26536\n",
      "Iter 130: Train loss 1.493, Learning Rate 2.000e-05, It/sec 3.163, Tokens/sec 691.506, Trained Tokens 28722\n",
      "Iter 140: Train loss 1.478, Learning Rate 2.000e-05, It/sec 2.991, Tokens/sec 731.268, Trained Tokens 31167\n",
      "Iter 150: Train loss 1.419, Learning Rate 2.000e-05, It/sec 3.305, Tokens/sec 692.101, Trained Tokens 33261\n",
      "Iter 160: Train loss 1.342, Learning Rate 2.000e-05, It/sec 2.910, Tokens/sec 731.043, Trained Tokens 35773\n",
      "Iter 170: Train loss 1.403, Learning Rate 2.000e-05, It/sec 3.347, Tokens/sec 699.844, Trained Tokens 37864\n",
      "Iter 180: Train loss 1.383, Learning Rate 2.000e-05, It/sec 3.216, Tokens/sec 689.765, Trained Tokens 40009\n",
      "Iter 190: Train loss 1.522, Learning Rate 2.000e-05, It/sec 3.624, Tokens/sec 690.703, Trained Tokens 41915\n",
      "Iter 200: Train loss 1.436, Learning Rate 2.000e-05, It/sec 3.012, Tokens/sec 701.417, Trained Tokens 44244\n",
      "Iter 200: Val loss 1.443, Val took 3.860s\n",
      "Iter 200: Saved adapter weights to checkpoints/200_trial1.npz.\n",
      "Iter 210: Train loss 1.437, Learning Rate 2.000e-05, It/sec 3.475, Tokens/sec 660.973, Trained Tokens 46146\n",
      "Iter 220: Train loss 1.437, Learning Rate 2.000e-05, It/sec 3.288, Tokens/sec 688.910, Trained Tokens 48241\n",
      "Iter 230: Train loss 1.457, Learning Rate 2.000e-05, It/sec 3.416, Tokens/sec 696.257, Trained Tokens 50279\n",
      "Iter 240: Train loss 1.376, Learning Rate 2.000e-05, It/sec 3.292, Tokens/sec 684.751, Trained Tokens 52359\n",
      "Iter 250: Train loss 1.561, Learning Rate 2.000e-05, It/sec 3.415, Tokens/sec 724.766, Trained Tokens 54481\n",
      "Iter 260: Train loss 1.594, Learning Rate 2.000e-05, It/sec 3.389, Tokens/sec 711.272, Trained Tokens 56580\n",
      "Iter 270: Train loss 1.420, Learning Rate 2.000e-05, It/sec 3.723, Tokens/sec 679.744, Trained Tokens 58406\n",
      "Iter 280: Train loss 1.525, Learning Rate 2.000e-05, It/sec 3.112, Tokens/sec 701.063, Trained Tokens 60659\n",
      "Iter 290: Train loss 1.425, Learning Rate 2.000e-05, It/sec 3.258, Tokens/sec 727.116, Trained Tokens 62891\n",
      "Iter 300: Train loss 1.520, Learning Rate 2.000e-05, It/sec 3.572, Tokens/sec 720.066, Trained Tokens 64907\n",
      "Iter 300: Saved adapter weights to checkpoints/300_trial1.npz.\n",
      "Iter 310: Train loss 1.377, Learning Rate 2.000e-05, It/sec 3.346, Tokens/sec 732.195, Trained Tokens 67095\n",
      "Iter 320: Train loss 1.459, Learning Rate 2.000e-05, It/sec 3.040, Tokens/sec 692.415, Trained Tokens 69373\n",
      "Iter 330: Train loss 1.523, Learning Rate 2.000e-05, It/sec 3.363, Tokens/sec 722.733, Trained Tokens 71522\n",
      "Iter 340: Train loss 1.481, Learning Rate 2.000e-05, It/sec 3.340, Tokens/sec 741.208, Trained Tokens 73741\n",
      "Iter 350: Train loss 1.487, Learning Rate 2.000e-05, It/sec 3.470, Tokens/sec 701.897, Trained Tokens 75764\n",
      "Iter 360: Train loss 1.491, Learning Rate 2.000e-05, It/sec 2.825, Tokens/sec 635.167, Trained Tokens 78012\n",
      "Iter 370: Train loss 1.503, Learning Rate 2.000e-05, It/sec 3.600, Tokens/sec 698.360, Trained Tokens 79952\n",
      "Iter 380: Train loss 1.445, Learning Rate 2.000e-05, It/sec 2.902, Tokens/sec 727.599, Trained Tokens 82459\n",
      "Iter 390: Train loss 1.480, Learning Rate 2.000e-05, It/sec 3.423, Tokens/sec 706.547, Trained Tokens 84523\n",
      "Iter 400: Train loss 1.361, Learning Rate 2.000e-05, It/sec 2.829, Tokens/sec 723.171, Trained Tokens 87079\n",
      "Iter 400: Val loss 1.400, Val took 4.169s\n",
      "Iter 400: Saved adapter weights to checkpoints/400_trial1.npz.\n",
      "Iter 410: Train loss 1.426, Learning Rate 2.000e-05, It/sec 3.189, Tokens/sec 664.859, Trained Tokens 89164\n",
      "Iter 420: Train loss 1.466, Learning Rate 2.000e-05, It/sec 3.722, Tokens/sec 665.106, Trained Tokens 90951\n",
      "Iter 430: Train loss 1.455, Learning Rate 2.000e-05, It/sec 3.378, Tokens/sec 665.080, Trained Tokens 92920\n",
      "Iter 440: Train loss 1.582, Learning Rate 2.000e-05, It/sec 2.951, Tokens/sec 745.747, Trained Tokens 95447\n",
      "Iter 450: Train loss 1.380, Learning Rate 2.000e-05, It/sec 2.612, Tokens/sec 707.778, Trained Tokens 98157\n",
      "Iter 460: Train loss 1.404, Learning Rate 2.000e-05, It/sec 2.719, Tokens/sec 669.474, Trained Tokens 100619\n",
      "Iter 470: Train loss 1.391, Learning Rate 2.000e-05, It/sec 2.466, Tokens/sec 712.111, Trained Tokens 103507\n",
      "Iter 480: Train loss 1.450, Learning Rate 2.000e-05, It/sec 3.548, Tokens/sec 706.836, Trained Tokens 105499\n",
      "Iter 490: Train loss 1.290, Learning Rate 2.000e-05, It/sec 3.499, Tokens/sec 677.421, Trained Tokens 107435\n",
      "Iter 500: Train loss 1.359, Learning Rate 2.000e-05, It/sec 2.840, Tokens/sec 731.354, Trained Tokens 110010\n",
      "Iter 500: Saved adapter weights to checkpoints/500_trial1.npz.\n",
      "Model training complete.\n",
      "\n",
      "CPU times: user 1.19 s, sys: 367 ms, total: 1.56 s\n",
      "Wall time: 3min 33s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "!python3 -m mlx_lm.lora \\\n",
    "--train \\\n",
    "--model {github_path} \\\n",
    "--data {train_set} \\\n",
    "--batch-size {batch_size} \\\n",
    "--lora-layers {lora_layers} \\\n",
    "--iters {iters} \\\n",
    "--max-seq-length {context_length} \\\n",
    "--learning-rate {learning_rate} \\\n",
    "--adapter-file {adapters} \\\n",
    "--save-every 100\n",
    "\n",
    "print(f\"\\nModel training complete.\\n\")\n",
    "print(f\"\\nIf you're getting a strange error and the training isn't happening (will be obvious as it'll end instantly).\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b17d61f-07d4-4ea4-9e13-a1e2de2ca35d",
   "metadata": {},
   "source": [
    "### Congratulations! You've trained a model! \n",
    "#### Well done! 🎉 \n",
    "#### You'll now have an adapters.npz file in the directory of this notebook. We can now apply this to the model in order to generate some text in the style of our training data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "a3ffcb4a-1643-4b9c-9795-3c9e237bf530",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "**Without your trained adapters:**\n",
      "==========\n",
      "Prompt: Below is a story about star wars\n",
      "\n",
      "Star Wars is a series of films that have been released by Lucasfilm Ltd. The first film was released in 1977 and the most recent one was released in 2015. The series has been a huge success and has been a major part of pop culture for over 40 years. The series has been adapted into various forms of media, including books, comics, video games, and even a theme park. The series has been a huge influence on pop culture\n",
      "==========\n",
      "Prompt: 142.209 tokens-per-sec\n",
      "Generation: 90.163 tokens-per-sec\n",
      "\n",
      "**And now with your trained adapters:**\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "The adapter file does not exist: trial1.npz",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[36], line 12\u001b[0m\n\u001b[1;32m      9\u001b[0m response \u001b[38;5;241m=\u001b[39m generate(model, tokenizer, prompt\u001b[38;5;241m=\u001b[39m prompt, verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m**And now with your trained adapters:**\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m---> 12\u001b[0m model, tokenizer \u001b[38;5;241m=\u001b[39m \u001b[43mload\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgithub_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43madapter_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madapters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     13\u001b[0m response \u001b[38;5;241m=\u001b[39m generate(model, tokenizer, prompt\u001b[38;5;241m=\u001b[39m prompt, verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.11/site-packages/mlx_lm/utils.py:386\u001b[0m, in \u001b[0;36mload\u001b[0;34m(path_or_hf_repo, tokenizer_config, adapter_file, lazy)\u001b[0m\n\u001b[1;32m    384\u001b[0m model \u001b[38;5;241m=\u001b[39m load_model(model_path, lazy)\n\u001b[1;32m    385\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m adapter_file \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 386\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mapply_lora_layers\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43madapter_file\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    387\u001b[0m     model\u001b[38;5;241m.\u001b[39meval()\n\u001b[1;32m    389\u001b[0m tokenizer \u001b[38;5;241m=\u001b[39m AutoTokenizer\u001b[38;5;241m.\u001b[39mfrom_pretrained(model_path, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mtokenizer_config)\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.11/site-packages/mlx_lm/tuner/utils.py:76\u001b[0m, in \u001b[0;36mapply_lora_layers\u001b[0;34m(model, adapter_file)\u001b[0m\n\u001b[1;32m     65\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;124;03mApply LoRA layers to the model.\u001b[39;00m\n\u001b[1;32m     67\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     73\u001b[0m \u001b[38;5;124;03m    nn.Module: The updated model with LoRA layers applied.\u001b[39;00m\n\u001b[1;32m     74\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mexists(adapter_file):\n\u001b[0;32m---> 76\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mFileNotFoundError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe adapter file does not exist: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00madapter_file\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     78\u001b[0m adapters \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m(mx\u001b[38;5;241m.\u001b[39mload(adapter_file)\u001b[38;5;241m.\u001b[39mitems())\n\u001b[1;32m     80\u001b[0m linear_replacements \u001b[38;5;241m=\u001b[39m []\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: The adapter file does not exist: trial1.npz"
     ]
    }
   ],
   "source": [
    "# Test the model's generation.\n",
    "\n",
    "from mlx_lm import load, generate\n",
    "\n",
    "prompt = \"Below is a story about star wars\\n\"\n",
    "\n",
    "print(f\"\\n**Without your trained adapters:**\")\n",
    "model, tokenizer = load(github_path)\n",
    "response = generate(model, tokenizer, prompt= prompt, verbose=True)\n",
    "\n",
    "print(f\"\\n**And now with your trained adapters:**\")\n",
    "model, tokenizer = load(github_path, adapter_file=adapters)\n",
    "response = generate(model, tokenizer, prompt= prompt, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2620ef93-5891-4c66-b130-35f7fa409441",
   "metadata": {},
   "source": [
    "### If you found this Notebook useful, I'd love to know - feel free to reach out to me on my [Twitter](https://twitter.com/priontific) or come chat with me on the [Curriculum Protocol](https://discord.gg/U8TZ2p4z) Discord.\n",
    "#### As MLX functionality continues to expand (different loss functions to enable RLHF, instruct fine-tuning etc.) I'll make more of these Notebooks.\n",
    "\n",
    "Well done on fine-tuning TinyLlama! 🎉  "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}