models/Codegemma/lora.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3d4e2931",
   "metadata": {},
   "source": [
    "## CodeGemma Parameter Efficient Fine-Tuning with LoRA using NeMo Framework\n",
    "\n",
    "CodeGemma is a groundbreaking new open model in the Gemini family of models from Google. CodeGemma is just as powerful as previous models but compact enough to run locally on NVIDIA RTX GPUs. CodeGemma is available in 2 sizes: 2B and 7B parameters. With NVIDIA NeMo, you can customize CodeGemma to fit your usecase and deploy an optimized model on your NVIDIA GPU.\n",
    "\n",
    "In this tutorial, we'll go over a specific kind of customization -- Low-rank adapter tuning to follow a specific output format (also known as LoRA). To learn how to perform full parameter supervised fine-tuning for instruction following (also known as SFT), see the [SFT notebook on Gemma Base Model](https://github.com/NVIDIA/GenerativeAIExamples/blob/main/models/Gemma/sft.ipynb). For LoRA, we'll perform all operations within the notebook on a single GPU. The compute resources needed for training depend on which CodeGemma model you use. For the 7 billion parameter variant, you'll need a GPU with 80GB of memory. For the 2 billion parameter model, 40GB will do.\n",
    "\n",
    "We'll also learn how to export your custom model to TensorRT-LLM, an open-source library that accelerates and optimizes inference performance of the latest LLMs on the NVIDIA AI platform."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a0c3e84",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "[LoRA tuning](https://arxiv.org/abs/2106.09685) is a parameter efficient method for fine-tuning models, where we freeze the base model parameters and update an auxiliary \"adapter\" with many fewer weights. At inference time, the adapter weights are combined with the base model weights to produce a new model, customized for a particular use case or dataset. Because this adapter is so much smaller than the base model, it can be trained with far fewer resources than it would take to fine-tune the entire model. In this notebook, we'll show you how to LoRA-tune small models like the CodeGemma models on a single A100 GPU.\n",
    "\n",
    "For this example, we're going to tune our CodeGemma model on the [Alpaca Python Code Instructions Dataset](https://huggingface.co/datasets/iamtarun/python_code_instructions_18k_alpaca) and tuning our model to enhance its instruction following ability for generating Python code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8df62050",
   "metadata": {},
   "source": [
    "## Download the Pretrained CodeGemma Model\n",
    "\n",
    "For all of our customization and deployment processes, we'll need to start off with a pre-trained version of CodeGemma in the `.nemo` format. You can download the base model in `.nemo` format from the NVIDIA GPU Cloud, or convert checkpoints from another framework into a `.nemo` file. You can choose to use the 2B parameter or 7B parameter CodeGemma models for this notebook -- the 2B model will be faster to customize, but the 7B model will be more capable.\n",
    "\n",
    "You can download either model from the NVIDIA NGC Catalog, using the NGC CLI. The instructions to install and configure the NGC CLI can be found [here](https://ngc.nvidia.com/setup/installers/cli).\n",
    "\n",
    "To download the model, execute one of the following commands, based on which model you want to use:\n",
    "\n",
    "ngc registry model download-version \"nvidia/nemo/codegemma_2b_base:1.0\"\n",
    "\n",
    "or\n",
    "\n",
    "ngc registry model download-version \"nvidia/nemo/codegemma_7b_base:1.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21facc49",
   "metadata": {},
   "source": [
    "## Getting NeMo Framework\n",
    "\n",
    "NVIDIA NeMo Framework is a generative AI framework built for researchers and PyTorch developers working on large language models (LLMs), multimodal models (MM), automatic speech recognition (ASR), and text-to-speech synthesis (TTS). The primary objective of NeMo is to provide a scalable framework for researchers and developers from industry and academia to more easily implement and design new generative AI models by being able to leverage existing code and pretrained models.\n",
    "\n",
    "If you haven't already, you can pull a container that includes the version of NeMo Framework and all dependencies needed for this notebook with the following:\n",
    "\n",
    "docker pull nvcr.io/nvidia/nemo:24.03.codegemma\n",
    "\n",
    "The best way to run this notebook is from within the container. You can do that by launching the container with the following command\n",
    "\n",
    "docker run -it --rm --gpus all --ipc=host --network host -v $(pwd):/workspace nvcr.io/nvidia/nemo:24.03.codegemma\n",
    "\n",
    "Then, from within the container, start the jupyter server with\n",
    "\n",
    "jupyter lab --no-browser --port=5000 --allow-root --ip 0.0.0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f23c5e",
   "metadata": {},
   "source": [
    "## Dataset Preparation\n",
    "\n",
    "Let's download Alpaca Python Code Instructions dataset from Hugging Face:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "444aa398",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git lfs install\n",
    "!git clone https://huggingface.co/datasets/iamtarun/python_code_instructions_18k_alpaca"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99c914d4",
   "metadata": {},
   "source": [
    "Finally, the following code snippets convert the dataset into the JSONL format that NeMo defaults for PEFT. Meanwhile, we will reformat the data into list of (prompt, completion) pairs that our model can appropriately handle. Please refer to the printout for the original code instruction data format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcc7f461",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import glob\n",
    "from random import seed, shuffle\n",
    "from huggingface_hub import login\n",
    "\n",
    "login(token='your_huggingface_access_token')\n",
    "parquet_file_path = glob.glob('./python_code_instructions_18k_alpaca/data/*.parquet')\n",
    "parquet_file_list = ''.join(parquet_file_path)\n",
    "df = pd.read_parquet(parquet_file_list)\n",
    "instruct2code_list = df.to_dict('records')\n",
    "\n",
    "seed(2)\n",
    "val_percent = 5\n",
    "test_percent = 5\n",
    "instruct2code_list = instruct2code_list[:len(instruct2code_list)] \n",
    "num_train = int(len(instruct2code_list) * (100 - val_percent - test_percent) / 100)\n",
    "num_val = int(len(instruct2code_list)*(val_percent)/100)\n",
    "shuffle(instruct2code_list)\n",
    "\n",
    "instruct2code_list_train = instruct2code_list[:num_train]\n",
    "instruct2code_list_val = instruct2code_list[num_train : num_train + num_val]\n",
    "instruct2code_list_test = instruct2code_list[num_train + num_val:]\n",
    "print(f\"=== Input prompt example from the training split:\\n{instruct2code_list_train[5]['prompt']}\\n\") \n",
    "print(f\"=== Output completion example from the validation split:\\n{instruct2code_list_val[5]['output']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d30c4c0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import json\n",
    "def write_jsonl(fname, json_objs):\n",
    "    with open(fname, 'wt') as f:\n",
    "        for o in json_objs:\n",
    "            f.write(json.dumps(o)+\"\\n\")\n",
    "def form_instruction(pair):\n",
    "    outpout_loc = pair.find('### Output')\n",
    "    return(pair[:outpout_loc])\n",
    "def convert_to_jsonl(instruct2code_list, output_path):\n",
    "    json_objs = []\n",
    "    for pair in instruct2code_list:\n",
    "        prompt = form_instruction(pair['prompt'])\n",
    "        completion = pair['output']\n",
    "        json_objs.append({\"input\": prompt, \"output\": completion})\n",
    "    write_jsonl(output_path, json_objs)\n",
    "    return json_objs\n",
    "\n",
    "print(len(instruct2code_list_train))\n",
    "train_json_objs = convert_to_jsonl(instruct2code_list_train, \"alpaca_python_train.jsonl\")\n",
    "val_json_objs= convert_to_jsonl(instruct2code_list_val, \"alpaca_python_val.jsonl\")\n",
    "test_json_objs = convert_to_jsonl(instruct2code_list_test, \"alpaca_python_test.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5284142f",
   "metadata": {},
   "source": [
    "Here's an example of what the data looks like after reformatting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c892887",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_json_objs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18b28a24",
   "metadata": {},
   "source": [
    "## LoRA Configuration and Training\n",
    "\n",
    "NeMo Framework provides support for configuration and training. To proceed with the training, you'll find a script at `/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py`. The script uses config parameters to control many of its operations. An example config file allows you to quickly see what options you can change and carry out different experiments. We can start by downloading the example config file, `megatron_gpt_peft_tuning_config.yaml` from github. The file is referenced to configure the parameters for the running PEFT training jobs in NeMo with LoRA technique for language model tuning. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b40c60b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_config.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54e5428e",
   "metadata": {},
   "source": [
    "To see all of the different configuration options available, you can take a look at the file we downloaded. For this example, we're going to update a couple of settings to point to our newly-prepared datasets and to make sure the LoRA tuning runs on our A100. Feel free to experiment with these different options -- you can swap in your own datasets and change the training settings depending on what GPU you're using.\n",
    "\n",
    "For data our data configuration, we'll point to the `jsonl` files we wrote out earlier. `concat_sampling_probabilities` determines what percentage of the finetuning data you would like to come from each file -- in our example we only have 1 training file so we choose [1.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89e40bd9",
   "metadata": {},
   "source": [
    "For our model settings, we don't have much to change since we're reading in a pretrained model and can inherit the values that were already set. We need to point to our existing `.nemo` file, specify that we want to use LoRA as our scheme for finetuning, and choose our parallelism and batch size values. The values below should be appropriate for a single A100 GPU.\n",
    "\n",
    "Make sure to change the `restore_from_path` setting with the path to the `.nemo` checkpoint!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "409717ce",
   "metadata": {},
   "source": [
    "Finally, we set some options for the `Trainer`. We'll be training on 1 GPU on a single node, at bfloat16 precision. For this example we'll train for 2000 steps, with a validation check every after every 200 iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c91a7095",
   "metadata": {},
   "source": [
    "After setting the `Trainer` object configurations to handle our training loop, we set configurations for an experiment manager to handle checkpointing and logging. We can load our model from disk into memory. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58478c4f",
   "metadata": {},
   "source": [
    "Now, let's see how to add the LoRA Adapter to our model and train it. We can specify that we want to use LoRA by using the `model.peft.peft_scheme` configuration to `lora`, which stores the types of applicable adapter and the hyperparameters required to initialize the adapter module.\n",
    "\n",
    "We're now ready to start training! As the training loop runs, you'll see the validation loss drop significantly -- even with this short demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3513c286",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "PEFT_SCHEME='lora'\n",
    "MODEL_SIZE=7b\n",
    "MBS=1\n",
    "TP=1\n",
    "PP=1\n",
    "NUM_DEVICES=1\n",
    "GBS=8\n",
    "SEQ_LEN=4096\n",
    "\n",
    "EXTRA_ARGS=\"\n",
    "        +model.fp8=False \\\n",
    "        +model.fp8_e4m3=False \\\n",
    "        +model.fp8_hybrid=True \\\n",
    "        +model.fp8_margin=0 \\\n",
    "        +model.fp8_interval=1 \\\n",
    "        +model.fp8_amax_history_len=128 \\\n",
    "        +model.fp8_amax_compute_algo=max \"\n",
    "\n",
    "TRAIN_DS=[alpaca_python_train.jsonl]\n",
    "VALID_DS=[alpaca_python_val.jsonl]\n",
    "GBS=128\n",
    "PACKED=False\n",
    "MODEL=codegemma-7b_fromhf.nemo\n",
    "EXP_DIR=nemo_experiments\n",
    "    \n",
    "torchrun --nproc_per_node=1 \\\n",
    "/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py \\\n",
    "    trainer.devices=${NUM_DEVICES} \\\n",
    "        trainer.num_nodes=1 \\\n",
    "        trainer.val_check_interval=200 \\\n",
    "        trainer.max_steps=2000 \\\n",
    "        +trainer.num_sanity_val_steps=0 \\\n",
    "        +trainer.limit_val_batches=3 \\\n",
    "        model.megatron_amp_O2=True \\\n",
    "        exp_manager.resume_if_exists=False \\\n",
    "        exp_manager.exp_dir=\"${EXP_DIR}\" \\\n",
    "        exp_manager.checkpoint_callback_params.save_top_k=0 \\\n",
    "        model.tensor_model_parallel_size=${TP} \\\n",
    "        model.pipeline_model_parallel_size=${PP} \\\n",
    "        model.micro_batch_size=${MBS} \\\n",
    "        model.global_batch_size=${GBS} \\\n",
    "        model.restore_from_path=${MODEL} \\\n",
    "        model.data.train_ds.num_workers=0 \\\n",
    "        model.data.validation_ds.num_workers=0 \\\n",
    "        +model.data.train_ds.packed_sequence=${PACKED} \\\n",
    "        ++model.sequence_parallel=False \\\n",
    "        +model.log_token_counts=True \\\n",
    "        model.data.train_ds.file_names=${TRAIN_DS} \\\n",
    "        model.data.train_ds.concat_sampling_probabilities=[1.0] \\\n",
    "        model.data.validation_ds.file_names=${VALID_DS} \\\n",
    "        model.peft.peft_scheme=${PEFT_SCHEME} \\\n",
    "        model.peft.lora_tuning.target_modules=[attention_qkv] \\\n",
    "        model.data.train_ds.max_seq_length=${SEQ_LEN} \\\n",
    "        model.data.validation_ds.max_seq_length=${SEQ_LEN} \\\n",
    "        +model.apply_rope_fusion=True \\\n",
    "        ${EXTRA_ARGS} \\\n",
    "        trainer.precision=bf16 \\\n",
    "        model.answer_only_loss=True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275984c7",
   "metadata": {},
   "source": [
    "Once training is completed you should see a saved '.nemo' file in the nemo_experiments folder. This checkpoint will only contain the trained adapter weights, and not the frozen base model weights."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21dabcf2",
   "metadata": {},
   "source": [
    "Next, we'll need to merge the weights of the base model and the weights of the adapter. If you're using the `NeMo Framework` container, you'll find a script for this at `/opt/NeMo/scripts/nlp_language_modeling/merge_lora_weights/merge.py`. Otherwise, you can download the standalone script from GitHub at https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/nlp_language_modeling/merge_lora_weights/merge.py\n",
    "\n",
    "To merge weights using the merge script, you'll need the path to the base model and trained adapter, as well as a path to save the merged model to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794cf054",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "python /opt/NeMo/scripts/nlp_language_modeling/merge_lora_weights/merge.py \\\n",
    "    trainer.accelerator=gpu \\\n",
    "    tensor_model_parallel_size=1 \\\n",
    "    pipeline_model_parallel_size=1 \\\n",
    "    gpt_model_file=codegemma-7b_fromhf.nemo \\\n",
    "    lora_model_path=megatron_gpt_peft_lora_tuning.nemo \\\n",
    "    merged_model_path=gemma_lora_alpaca_python_merged.nemo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09bbc702",
   "metadata": {},
   "source": [
    "With our merged model weights, we can run evaluation on test dataset using `megatron_gpt_peft_eval.py`. We set the Set the appropriate model checkpoint path, test file path, batch sizes, number of tokens etc. and run evaluation on the test file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb943ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "python /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_eval.py \\\n",
    "    model.restore_from_path=gemma_lora_alpaca_python_merged.nemo \\\n",
    "    trainer.devices=1 \\\n",
    "    model.global_batch_size=8 \\\n",
    "    model.data.test_ds.file_names=[\"alpaca_python_test.jsonl\"] \\\n",
    "    model.data.test_ds.names=[\"alpaca_python_test_set\"] \\\n",
    "    model.data.test_ds.global_batch_size=8 \\\n",
    "    model.data.test_ds.micro_batch_size=1 \\\n",
    "    model.data.test_ds.tokens_to_generate=20 \\\n",
    "    model.tensor_model_parallel_size=1 \\\n",
    "    model.pipeline_model_parallel_size=1 \\\n",
    "    inference.greedy=True \\\n",
    "    model.data.test_ds.output_file_path_prefix=/results \\\n",
    "    model.data.test_ds.write_predictions_to_file=True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "579048ac",
   "metadata": {},
   "source": [
    "Check the output from the result file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55756a1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tail -n 4 /results_test_alpaca_python_test_set_inputs_preds_labels.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "389f2115",
   "metadata": {},
   "source": [
    "Note, This is only a sample output (based of a toy LoRA example) and your output may vary. The performance can be further improved by fine tuning the model for more steps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fa0eb3a",
   "metadata": {},
   "source": [
    "Finally, let's continue on to the \"Exporting to TensorRT-LLM\" section, to learn how to export our new model for optimized inference using TensorRT-LLM! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86e36e7b",
   "metadata": {},
   "source": [
    "## Exporting to TensorRT-LLM\n",
    "\n",
    "TensorRT-LLM is an open-source library for optimizing inference performance to acheive state-of-the-art speed on NVDIA GPUs. The NeMo framework offers an easy way to compile .nemo models into optimized TensorRT-LLM engines which you can run locally embedded in another application, or serve to other applications using a server like Triton Inference Server.\n",
    "\n",
    "To start with, lets create a folder where our exported model will land"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b32a604",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir codegemma_trt_llm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1f8ea73",
   "metadata": {},
   "source": [
    "With our merged model weights, we just need to create an instance of the TensorRTLLM class and call the TensorRTLLM.export() function -- pointing the nemo_checkpoint_path argument to the newly merged model from above.\n",
    "\n",
    "This creates a couple of files in the folder we created -- an engine file that holds the weights and the compiled execution graph of the model, a tokenizer.model file which holds the tokenizer information, and config.json which holds some metadata about the model (along with model.cache, which caches some operations and makes it faster to re-compile the model in the future.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2eca7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.export import TensorRTLLM\n",
    "trt_llm_exporter = TensorRTLLM(model_dir=\"gemma_alpaca_python_merged_trt_llm\")\n",
    "trt_llm_exporter.export(nemo_checkpoint_path=\"gemma_lora_alpaca_python_merged.nemo\", model_type=\"gemma\", n_gpus=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2de9dde1",
   "metadata": {},
   "source": [
    "With the model exported into TensorRTLLM, we can perform very fast inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072e9474",
   "metadata": {},
   "outputs": [],
   "source": [
    "trt_llm_exporter.forward([\"Implement Fibonacci sequence in Python\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b2f968a",
   "metadata": {},
   "source": [
    "There's also a convenient function to deploy a the model as a service, backed by Triton Inference Server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70fb870a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.deploy import DeployPyTriton\n",
    "\n",
    "nm = DeployPyTriton(model=trt_llm_exporter, triton_model_name=\"gemma\")\n",
    "nm.deploy()\n",
    "nm.serve()"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}