notebooks/llama_compose_trainer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Mergoo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install mergoo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Mergoo-MOE Checkpoint\n",
    "\n",
    "**Selecting Experts:**\n",
    "\n",
    "The advantage of Mergoo is that we can leverage the fine-tuned models, which excel in their respective domains. In this notebook, we merge three experts into a unified MOE architecture that we will later fine-tune.\n",
    "\n",
    "The Code Llama Python model is a fine-tuned version of Llama 7b on Python programming texts, hence it performs well for Python queries. The WikiChat model uses Wikipedia and a 7-stage pipeline to ensure its responses are factual and minimize hallucinations. Check out their [repository](https://github.com/stanford-oval/WikiChat) for a detailed explanation.\n",
    "\n",
    "We also add the base model Llama-2-7b-hf as an expert so that the final MOE model is well-balanced and performs excellently in diverse domains. The unified model will have features of all the experts and will perform well in Python programming and wiki questions.\n",
    "\n",
    "Here we are creating an MOE model using the following models:\n",
    "- [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)\n",
    "- [codellama/CodeLlama-7b-Python-hf](https://huggingface.co/codellama/CodeLlama-7b-Python-hf)\n",
    "- [stanford-oval/Llama-2-7b-WikiChat-fused](https://huggingface.co/stanford-oval/Llama-2-7b-WikiChat-fused)\n",
    "\n",
    "**Preparing Config:**\n",
    "- `model_type`: llama/mistral/bert. This is the base model family of the experts. At the moment, all the experts should come from the same base model family.\n",
    "- `num_experts_per_tok`: Total number of active experts at each step. These experts are selected sparsely.\n",
    "- `experts`: List of dictionaries of seed models that would get merged. For each expert, `model_id` is mandatory. The model_id can be either a local path or a Huggingface model id.\n",
    "- `router_layers`: These are the layer names that would be replaced with MOE layers. Weights of the rest of the layers are aggregated using averaging. In the future, we will support multiple aggregation methods from MergeKit.\n",
    "- `router_layers_index`: List of indexes. These are the indexes of transformer blocks, layers of these index would be converted to MOE. Default `router_layers_index` is empty meaning the MOE conversion gets applied on all the layers, given that `router_layers` identifier matches. `[None]` can be used when no MOE layer should be kept following the [BTM](https://arxiv.org/abs/2208.03306) architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from mergoo.compose_experts import ComposeExperts\n",
    "\n",
    "model_id = \"data/checkpoint_demo\"\n",
    "config = \\\n",
    "{\n",
    "    \"model_type\": \"llama\",\n",
    "    \"num_experts_per_tok\": 2,\n",
    "    \"experts\":[\n",
    "        {\n",
    "            \"expert_name\" : \"base_expert\",\n",
    "            \"model_id\" : \"meta-llama/Llama-2-7b-hf\"\n",
    "        },\n",
    "        {\n",
    "            \"expert_name\" : \"expert_1\",\n",
    "            \"model_id\" : \"codellama/CodeLlama-7b-Python-hf\"\n",
    "        },\n",
    "        {\n",
    "            \"expert_name\" : \"expert_2\",\n",
    "            \"model_id\" : \"stanford-oval/Llama-2-7b-WikiChat-fused\"\n",
    "        }\n",
    "    ],\n",
    "    \"router_layers\":[\n",
    "        \"gate_proj\",\n",
    "        \"up_proj\",\n",
    "        \"down_proj\"\n",
    "    ],\n",
    "}\n",
    "# create checkpoint\n",
    "expertmerger = ComposeExperts( config, torch_dtype=torch.float16 )\n",
    "expertmerger.compose()\n",
    "expertmerger.save_checkpoint(model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Now that we have created an MOE checkpoint, all the layers of this model are pretrained except for the gating/routing layers that we added. The routing layer selects the top K experts, in our case K=2. We support HuggingFace trainers: Trainer, SFTrainer. In this example, we are using the Python_code_instructions_18k_alpaca dataset for finetuning. We will train only the router layers, keeping all the other layers frozen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e1876544a5e449787d4b3defb8a62b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of LlamaForCausalLM were not initialized from the model checkpoint at data/llama_moe and are newly initialized: ['model.layers.0.mlp.down_proj.gate.weight', 'model.layers.0.mlp.gate_proj.gate.weight', 'model.layers.0.mlp.up_proj.gate.weight', 'model.layers.1.mlp.down_proj.gate.weight', 'model.layers.1.mlp.gate_proj.gate.weight', 'model.layers.1.mlp.up_proj.gate.weight', 'model.layers.10.mlp.down_proj.gate.weight', 'model.layers.10.mlp.gate_proj.gate.weight', 'model.layers.10.mlp.up_proj.gate.weight', 'model.layers.11.mlp.down_proj.gate.weight', 'model.layers.11.mlp.gate_proj.gate.weight', 'model.layers.11.mlp.up_proj.gate.weight', 'model.layers.12.mlp.down_proj.gate.weight', 'model.layers.12.mlp.gate_proj.gate.weight', 'model.layers.12.mlp.up_proj.gate.weight', 'model.layers.13.mlp.down_proj.gate.weight', 'model.layers.13.mlp.gate_proj.gate.weight', 'model.layers.13.mlp.up_proj.gate.weight', 'model.layers.14.mlp.down_proj.gate.weight', 'model.layers.14.mlp.gate_proj.gate.weight', 'model.layers.14.mlp.up_proj.gate.weight', 'model.layers.15.mlp.down_proj.gate.weight', 'model.layers.15.mlp.gate_proj.gate.weight', 'model.layers.15.mlp.up_proj.gate.weight', 'model.layers.16.mlp.down_proj.gate.weight', 'model.layers.16.mlp.gate_proj.gate.weight', 'model.layers.16.mlp.up_proj.gate.weight', 'model.layers.17.mlp.down_proj.gate.weight', 'model.layers.17.mlp.gate_proj.gate.weight', 'model.layers.17.mlp.up_proj.gate.weight', 'model.layers.18.mlp.down_proj.gate.weight', 'model.layers.18.mlp.gate_proj.gate.weight', 'model.layers.18.mlp.up_proj.gate.weight', 'model.layers.19.mlp.down_proj.gate.weight', 'model.layers.19.mlp.gate_proj.gate.weight', 'model.layers.19.mlp.up_proj.gate.weight', 'model.layers.2.mlp.down_proj.gate.weight', 'model.layers.2.mlp.gate_proj.gate.weight', 'model.layers.2.mlp.up_proj.gate.weight', 'model.layers.20.mlp.down_proj.gate.weight', 'model.layers.20.mlp.gate_proj.gate.weight', 'model.layers.20.mlp.up_proj.gate.weight', 'model.layers.21.mlp.down_proj.gate.weight', 'model.layers.21.mlp.gate_proj.gate.weight', 'model.layers.21.mlp.up_proj.gate.weight', 'model.layers.22.mlp.down_proj.gate.weight', 'model.layers.22.mlp.gate_proj.gate.weight', 'model.layers.22.mlp.up_proj.gate.weight', 'model.layers.23.mlp.down_proj.gate.weight', 'model.layers.23.mlp.gate_proj.gate.weight', 'model.layers.23.mlp.up_proj.gate.weight', 'model.layers.24.mlp.down_proj.gate.weight', 'model.layers.24.mlp.gate_proj.gate.weight', 'model.layers.24.mlp.up_proj.gate.weight', 'model.layers.25.mlp.down_proj.gate.weight', 'model.layers.25.mlp.gate_proj.gate.weight', 'model.layers.25.mlp.up_proj.gate.weight', 'model.layers.26.mlp.down_proj.gate.weight', 'model.layers.26.mlp.gate_proj.gate.weight', 'model.layers.26.mlp.up_proj.gate.weight', 'model.layers.27.mlp.down_proj.gate.weight', 'model.layers.27.mlp.gate_proj.gate.weight', 'model.layers.27.mlp.up_proj.gate.weight', 'model.layers.28.mlp.down_proj.gate.weight', 'model.layers.28.mlp.gate_proj.gate.weight', 'model.layers.28.mlp.up_proj.gate.weight', 'model.layers.29.mlp.down_proj.gate.weight', 'model.layers.29.mlp.gate_proj.gate.weight', 'model.layers.29.mlp.up_proj.gate.weight', 'model.layers.3.mlp.down_proj.gate.weight', 'model.layers.3.mlp.gate_proj.gate.weight', 'model.layers.3.mlp.up_proj.gate.weight', 'model.layers.30.mlp.down_proj.gate.weight', 'model.layers.30.mlp.gate_proj.gate.weight', 'model.layers.30.mlp.up_proj.gate.weight', 'model.layers.31.mlp.down_proj.gate.weight', 'model.layers.31.mlp.gate_proj.gate.weight', 'model.layers.31.mlp.up_proj.gate.weight', 'model.layers.4.mlp.down_proj.gate.weight', 'model.layers.4.mlp.gate_proj.gate.weight', 'model.layers.4.mlp.up_proj.gate.weight', 'model.layers.5.mlp.down_proj.gate.weight', 'model.layers.5.mlp.gate_proj.gate.weight', 'model.layers.5.mlp.up_proj.gate.weight', 'model.layers.6.mlp.down_proj.gate.weight', 'model.layers.6.mlp.gate_proj.gate.weight', 'model.layers.6.mlp.up_proj.gate.weight', 'model.layers.7.mlp.down_proj.gate.weight', 'model.layers.7.mlp.gate_proj.gate.weight', 'model.layers.7.mlp.up_proj.gate.weight', 'model.layers.8.mlp.down_proj.gate.weight', 'model.layers.8.mlp.gate_proj.gate.weight', 'model.layers.8.mlp.up_proj.gate.weight', 'model.layers.9.mlp.down_proj.gate.weight', 'model.layers.9.mlp.gate_proj.gate.weight', 'model.layers.9.mlp.up_proj.gate.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "# load the composed checkkpoint\n",
    "import torch\n",
    "from mergoo.models.modeling_llama import LlamaForCausalLM\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained(\n",
    "    model_id, \n",
    "    device_map=\"auto\", \n",
    "    torch_dtype=torch.bfloat16,\n",
    ")# 'gate' / router layers are untrained hence loaded warning would appeare for them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(579, 387)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# train only router (gating) layers\n",
    "n_weights, n_router_weights  = 0,0\n",
    "for name, weight in model.named_parameters():\n",
    "    if \"gate\" not in name:\n",
    "        weight.requires_grad_(False)\n",
    "        n_router_weights += 1\n",
    "    n_weights += 1\n",
    "n_weights, n_router_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import random\n",
    "\n",
    "dataset = datasets.load_dataset(\"iamtarun/python_code_instructions_18k_alpaca\")['train']\n",
    "dataset = dataset['prompt']\n",
    "random.shuffle(dataset)\n",
    "dataset_train =  datasets.Dataset.from_dict(dict(prompt=dataset[:-1000]))\n",
    "dataset_test = datasets.Dataset.from_dict(dict(prompt=dataset[-1000:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['prompt'],\n",
       "     num_rows: 25085\n",
       " }),\n",
       " Dataset({\n",
       "     features: ['prompt'],\n",
       "     num_rows: 1000\n",
       " }))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_train, dataset_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6edd83852a944c6aa926207849780b5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25085 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8211f65ca2641c49d5d15eb59e1e26b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:317: UserWarning: You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/accelerate/accelerator.py:436: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches', 'even_batches', 'use_seedable_sampler']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
      "dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False, even_batches=True, use_seedable_sampler=True)\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from trl import SFTTrainer\n",
    "from transformers import TrainingArguments\n",
    "\n",
    "trainer_args = TrainingArguments(\n",
    "    output_dir= \"checkpoints/llama_moe\",\n",
    "    per_device_train_batch_size = 1,\n",
    "    per_device_eval_batch_size = 1, \n",
    "    learning_rate= 1e-5,\n",
    "    save_total_limit=1,\n",
    "    num_train_epochs=1,\n",
    "    eval_steps= 5000,\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_steps= 25,\n",
    "    gradient_accumulation_steps=4,\n",
    "    bf16=True\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model,\n",
    "    args= trainer_args,\n",
    "    train_dataset= dataset_train,\n",
    "    eval_dataset= dataset_test,\n",
    "    dataset_text_field=\"prompt\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='107' max='6271' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 107/6271 03:28 < 3:24:21, 0.50 it/s, Epoch 0.02/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>25</td>\n",
       "      <td>2.713500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>1.858800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>75</td>\n",
       "      <td>1.437100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.260700</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}