examples/colab-notebooks/colab-axolotl-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AKjdG7tbTb-n"
   },
   "source": [
    "# Example notebook for running Axolotl on google colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RcbNpOgWRcii"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "# Check so there is a gpu available, a T4(free tier) is enough to run this notebook\n",
    "assert (torch.cuda.is_available()==True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h3nLav8oTRA5"
   },
   "source": [
    "## Install Axolotl and dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3c3yGAwnOIdi",
    "outputId": "e3777b5a-40ef-424f-e181-62dfecd1dd01"
   },
   "outputs": [],
   "source": [
    "!pip install -e git+https://github.com/axolotl-ai-cloud/axolotl#egg=axolotl\n",
    "!pip install flash-attn==\"2.5.0\"\n",
    "!pip install deepspeed==\"0.13.1\"!pip install mlflow==\"2.13.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BW2MFr7HTjub"
   },
   "source": [
    "## Create an yaml config file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9pkF2dSoQEUN"
   },
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "# Your YAML string\n",
    "yaml_string = \"\"\"\n",
    "base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\n",
    "model_type: LlamaForCausalLM\n",
    "tokenizer_type: LlamaTokenizer\n",
    "\n",
    "load_in_8bit: false\n",
    "load_in_4bit: true\n",
    "strict: false\n",
    "\n",
    "datasets:\n",
    "  - path: mhenrichsen/alpaca_2k_test\n",
    "    type: alpaca\n",
    "dataset_prepared_path:\n",
    "val_set_size: 0.05\n",
    "output_dir: ./outputs/qlora-out\n",
    "\n",
    "adapter: qlora\n",
    "lora_model_dir:\n",
    "\n",
    "sequence_len: 4096\n",
    "sample_packing: true\n",
    "eval_sample_packing: false\n",
    "pad_to_sequence_len: true\n",
    "\n",
    "lora_r: 32\n",
    "lora_alpha: 16\n",
    "lora_dropout: 0.05\n",
    "lora_target_modules:\n",
    "lora_target_linear: true\n",
    "lora_fan_in_fan_out:\n",
    "\n",
    "wandb_project:\n",
    "wandb_entity:\n",
    "wandb_watch:\n",
    "wandb_name:\n",
    "wandb_log_model:\n",
    "\n",
    "gradient_accumulation_steps: 4\n",
    "micro_batch_size: 2\n",
    "num_epochs: 4\n",
    "optimizer: paged_adamw_32bit\n",
    "lr_scheduler: cosine\n",
    "learning_rate: 0.0002\n",
    "\n",
    "train_on_inputs: false\n",
    "group_by_length: false\n",
    "bf16: auto\n",
    "fp16:\n",
    "tf32: false\n",
    "\n",
    "gradient_checkpointing: true\n",
    "early_stopping_patience:\n",
    "resume_from_checkpoint:\n",
    "local_rank:\n",
    "logging_steps: 1\n",
    "xformers_attention:\n",
    "flash_attention: true\n",
    "\n",
    "warmup_steps: 10\n",
    "evals_per_epoch: 4\n",
    "saves_per_epoch: 1\n",
    "debug:\n",
    "deepspeed:\n",
    "weight_decay: 0.0\n",
    "fsdp:\n",
    "fsdp_config:\n",
    "special_tokens:\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "# Convert the YAML string to a Python dictionary\n",
    "yaml_dict = yaml.safe_load(yaml_string)\n",
    "\n",
    "# Specify your file path\n",
    "file_path = 'test_axolotl.yaml'\n",
    "\n",
    "# Write the YAML file\n",
    "with open(file_path, 'w') as file:\n",
    "    yaml.dump(yaml_dict, file)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bidoj8YLTusD"
   },
   "source": [
    "## Launch the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ydTI2Jk2RStU",
    "outputId": "d6d0df17-4b53-439c-c802-22c0456d301b"
   },
   "outputs": [],
   "source": [
    "# By using the ! the comand will be executed as a bash command\n",
    "!accelerate launch -m axolotl.cli.train /content/test_axolotl.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play with inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# By using the ! the comand will be executed as a bash command\n",
    "!accelerate launch -m axolotl.cli.inference /content/test_axolotl.yaml \\\n",
    "    --qlora_model_dir=\"./qlora-out\" --gradio"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}