lama2-7b_galore.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4e89d054-2236-4e32-b068-9be317814f6b",
   "metadata": {},
   "source": [
    "# GaLore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b5273f-0d20-4df5-8714-861ea063ca53",
   "metadata": {},
   "source": [
    "GaLore is a novel approach detailed in the publication GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection. \n",
    "GaLore reduces VRAM requirements not by decreasing the number of parameters directly but by optimizing how these parameters are trained.\n",
    "\n",
    "GaLore focuses on two primary strategies:\n",
    "* Gradient Low-Rank Projection: GaLore shifts away from handling the full, high-dimensional gradients of weight matrices. Instead, it projects these gradients onto a low-rank space, significantly reducing the computational load while retaining essential informat\n",
    "* Per-Layer Weight Updates: Unlike the conventional method where an optimizer updates all layers simultaneously after backpropagation, GaLore implements updates on a per-layer basis during backpropagation. This approach further reduces the memory footprint throughout the training process.\n",
    "\n",
    "Paper: https://huggingface.co/papers/2403.03507\n",
    "\n",
    "Official HF blog post: https://huggingface.co/blog/galore\n",
    "\n",
    "My experiments: https://medium.com/@geronimo7/llm-training-on-consumer-gpus-with-galore-d25075143cfb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "359d230d-8d6f-43ae-b2ad-40df8880b156",
   "metadata": {},
   "source": [
    "**Here's how to train the llama2-7b model using the GaLore layerwise method, operational on a single NVIDIA GeForce RTX 3090 GPU.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3e306f-c2d2-45bc-9539-792f65575c6e",
   "metadata": {},
   "source": [
    "# Prerequisistes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16be4000-bfd2-47f0-9bc3-d78b899e7da9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U bitsandbytes wandb huggingface_hub flash_attn sentencepiece accelerate datasets trl transformers tokenizers galore_torch flash_attn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f841c61-da3a-4f3b-9b23-3991cce7f9f4",
   "metadata": {},
   "source": [
    "Versions known to work:\n",
    "```\n",
    "accelerate==0.28.0\n",
    "bitsandbytes==0.43.0\n",
    "datasets==2.18.0\n",
    "tokenizers==0.15.2\n",
    "transformers==4.39.1\n",
    "trl==0.8.1\n",
    "wandb==0.16.4\n",
    "torch==2.2.1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f96903fc-b33c-4141-9922-ef440546572b",
   "metadata": {},
   "source": [
    "# Log into wandb and huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ced95ce1-bb29-46cb-953f-fd15626d68fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from wandb import login\n",
    "login(key=\"XXX\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c52bee-0de5-44a1-8273-7b3284452bfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import huggingface_hub\n",
    "huggingface_hub.login(token=\"XXX\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43d0d9aa-3f71-45c4-b6e7-6cc08714ff10",
   "metadata": {},
   "source": [
    "# Setup GaLore hyperparameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b424738d-6e26-4fca-aaae-cc9c3378a4a3",
   "metadata": {},
   "source": [
    "The GaLore optimizer comes with a few hyperparameters to set:\n",
    "\n",
    "* `target_modules_list`: Specifies the layers targeted by GaLore\n",
    "* `rank`: The rank of the projection matrices. Similar to LoRA, the higher the rank the more closely the finetuning will resemble a full parameter finetune. The GaLore authors recomment 1024 for a 7B model.\n",
    "* `update_proj_gap`: The number of steps after which the projections are updated. The update is an expensive step and takes around 15 minutes for a 7B model. Defines the interval for updating projections, with a suggested range between 50 and 1000 steps.\n",
    "* `scale`: A scale factor akin to LoRA’s alpha, adjusting the update strength. After trying a few values I found scale=2 to most closely resemble a classic full-parameter finetune."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90fcf0c-023e-40c1-a601-4447f5bc6eac",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 1e-5\n",
    "\n",
    "rank = 1024\n",
    "update_proj_gap = 200\n",
    "scale = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb312db-0d10-4079-a220-99b9a47817db",
   "metadata": {},
   "source": [
    "# Load model, tokenizer, dataset and setup trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5f479e-7606-495f-ab3b-1ee618875639",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change me!\n",
    "modelpath = \"meta-llama/Llama-2-7b-hf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14e14c8d-72c0-4bf8-82be-d0e5ec5e7ef8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, set_seed, get_constant_schedule\n",
    "from trl import SFTTrainer, setup_chat_format, DataCollatorForCompletionOnlyLM\n",
    "from datasets import load_dataset\n",
    "import torch, uuid, wandb\n",
    "\n",
    "set_seed(42)\n",
    "run_id = f\"galore-{str(uuid.uuid4())}\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    modelpath,    \n",
    "    torch_dtype = torch.bfloat16,\n",
    "    attn_implementation = \"flash_attention_2\",  \n",
    "    device_map = \"auto\",\n",
    "    use_cache = False,\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(modelpath, use_fast = False)\n",
    "\n",
    "model, tokenizer = setup_chat_format(model, tokenizer)\n",
    "if tokenizer.pad_token in [None, tokenizer.eos_token]: \n",
    "    tokenizer.pad_token = tokenizer.unk_token\n",
    "\n",
    "dataset = load_dataset(\"g-ronimo/oasst2_top4k_en\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74ace482-0051-4916-8e24-a935c855587d",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_arguments = TrainingArguments(\n",
    "    output_dir = f\"out_{run_id}\",\n",
    "    evaluation_strategy = \"steps\",\n",
    "    label_names = [\"labels\"],\n",
    "    per_device_train_batch_size = 16,\n",
    "    save_steps = 250,\n",
    "    eval_steps = 250,\n",
    "    logging_steps = 1, \n",
    "    learning_rate = lr,\n",
    "    num_train_epochs = 3,\n",
    "    lr_scheduler_type = \"constant\",\n",
    "    gradient_checkpointing = True,\n",
    "    group_by_length = False,\n",
    "    optim=\"galore_adamw_8bit_layerwise\",\n",
    "    optim_target_modules=[\"attn\", \"mlp\"],\n",
    "    optim_args=f\"rank={rank}, update_proj_gap={update_proj_gap}, scale={scale}\",\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model = model,\n",
    "    tokenizer = tokenizer,\n",
    "    train_dataset = dataset[\"train\"],\n",
    "    eval_dataset = dataset['test'],\n",
    "    data_collator = DataCollatorForCompletionOnlyLM(\n",
    "        instruction_template = \"<|im_start|>user\", \n",
    "        response_template = \"<|im_start|>assistant\", \n",
    "        tokenizer = tokenizer, \n",
    "        mlm = False),\n",
    "    max_seq_length = 256,\n",
    "    dataset_kwargs = dict(add_special_tokens = False),\n",
    "    args = training_arguments,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef43edd-639d-476c-b26b-6e6110d14d8d",
   "metadata": {},
   "source": [
    "# Init wandb and start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7796bce9-820d-4b98-a6ec-2b2e4c2c236f",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(\n",
    "    project = \"galore-7B\", \n",
    "    name = run_id,\n",
    ").log_code(include_fn=lambda path: path.endswith(\".py\") or path.endswith(\".ipynb\"))\n",
    "\n",
    "trainer.train()"
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}