{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# QLoRA: Efficient Finetuning of Quantized LLMs\n",
    "\n",
    "[QLoRA](https://arxiv.org/abs/2305.14314) is a technique primarily designed to further reduce the training memory requirement of LoRA. In QLoRA, the LoRA backbone weights are quantized to reduce the model footprint.\n",
    "\n",
    "The main features of Quark's QLoRA:\n",
    "\n",
    "1. QLoRA uses low-rank matrices to represent changes in model parameters, fine-tuning pre-trained language models by introducing additional low-rank parameters, thereby reducing the number of parameters that need to be updated. This method maintains model performance while significantly reducing the storage and computational resources required for training.\n",
    "2. Quantizing model weights typically involves converting floating-point parameters into low-bit integers to reduce memory usage and speed up computation. During the training phase, the model's weights are quantized and locked (e.g., the weights of linear layers and corresponding quantization parameters such as scale and zero point), and only the adapter matrices are trained. This reduces the number of parameters optimized during the training phase.\n",
    "3. After training, we fold the adapter weights back into the weights of the linear layer. After merging weights, the model's test metric still shows satisfactory results. As adapters are absorbed into weight, the quantized LLMs are suitable for accelerated deployment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Highlight Overview\n",
    "\n",
    "We conduct `mxfp4` quantization. We simply take the `Llama-3.2-1B-Instruct` as an example to conduct the experiments. After PTQ, we added the adapter for every Linear and performed the QAT training. After training, the adapter_A and adapter_B are merged into Linear's weight. The final PPL metric results show the effectiveness of the QLoRA tools.\n",
    "\n",
    "- **Pipeline overview :** Prepare LLM model -> PTQ -> insert adapter for liner -> QLoRA training -> merge adapter into weight.\n",
    "- **Save memory:** In a single GPU with 80GB memory, we test the `Llama-3.2-1B-Instruct` (less than 20GB) and `gpt-oss-20b` (less than 70GB).\n",
    "- **Efficient training:** during the whole training, only the adapters are updateable. Compared to training the whole model, QLoRA needs a small amount of training data and training loops.\n",
    "\n",
    "- **Merge adapter weight:** After training, adapters' weight will be merged into linear's weight. After merging, the model can still keep a relative high quality metric."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Work\n",
    "\n",
    "As mentioned above, we first conduct Quark PTQ, and then prepare the training dataset and LLM trainer. After several steps of training, we finally merge the adapters' weight into the linear weight.\n",
    "\n",
    "As we supply a simple example for you, with simple modification, you can directly run it, and test the metric.\n",
    "\n",
    "- Prepare LLM & perform the corresponding PTQ: for this section, the user can refer to the Quark PTQ example.\n",
    "  - Download the LLM model and save to the proper path. In this example;\n",
    "  - Quark PTQ example code: Quark/examples/torch/language_modeling/llm_ptq/quantize_quark.py\n",
    "  - Quark PTQ tutorials: Quark/examples/torch/language_modeling/llm_ptq/example_quark_torch_llm_ptq.rst\n",
    "- Prepare the Training Dataset & the Training & training parameters.\n",
    "\n",
    "  - Training Dataset: We use nvidia/Daring-Anteater as a fine-tune dataset. User can refer to [nvidia/Daring-Anteater](nvidia/Daring-Anteater)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## An End-to-End example of QLoRA training\n",
    "\n",
    "### Step 1: import the necessary lib:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import transformers\n",
    "from transformers import AutoProcessor, Trainer\n",
    "\n",
    "sys.path.append(\"./../../\")\n",
    "from utils import DataArguments, PTQQuantArguments, TrainingArguments, make_supervised_data_module\n",
    "\n",
    "from quark.contrib.llm_eval import eval_model\n",
    "from quark.torch import LLMTemplate, ModelQuantizer\n",
    "from quark.torch.utils.llm import get_calib_dataloader, get_model, get_tokenizer, prepare_for_moe_quant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Prepare args for PTQ and QLoRA training :\n",
    "\n",
    "For better demonstration, we predefine the majority hyperparameters. And for training, we only use a small part of the dataset and a small number of training steps. For better training results, the user can adjust the learning parameters and other hyperparameters for better performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE as runing the JuPyter Notebook will have some not desired args, so we delete.\n",
    "if len(sys.argv) >= 1 and sys.argv[1].startswith(\"--f\"):\n",
    "    del sys.argv[1]\n",
    "\n",
    "parser = transformers.HfArgumentParser((DataArguments, TrainingArguments, PTQQuantArguments))\n",
    "data_args, training_args, quant_args = parser.parse_args_into_dataclasses()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After download the training dataset and LLM, assign the dataset path and LLM model path."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_args.cache_dir = \"./hug_data\"\n",
    "quant_args.model_dir = \"./pretrained_models/meta-llama/Llama-3.2-1B-Instruct/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Perform the Quark PTQ\n",
    "\n",
    "We first conduct the PTQ, we use `mxfp4` as the default quantization config, the detail config is located in `./utils.py`. As the PTQ pipeline is mainly copied from the Quark official PTQ example, for more information users can refer to: `examples/torch/language_modeling/llm_ptq/quantize_quark.py`\n",
    "\n",
    "#### Step 3.1 Load the LLM and perform modification (if necessary).\n",
    "\n",
    "NOTE: for some MOE models, we need to replace the expert part with Linears, so as to perform calibration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = quant_args.device\n",
    "model, model_dtype = get_model(\n",
    "    quant_args.model_dir,\n",
    "    quant_args.data_type,\n",
    "    device,\n",
    "    quant_args.multi_gpu,\n",
    "    quant_args.multi_device,\n",
    "    quant_args.model_attn_implementation,\n",
    "    trust_remote_code=quant_args.trust_remote_code,\n",
    ")\n",
    "\n",
    "# prepare_for_moe_quant(model, quant_args.quant_algo)\n",
    "prepare_for_moe_quant(model)\n",
    "\n",
    "# model_type = get_model_type(model)\n",
    "model_type = model.config.model_type if hasattr(model.config, \"model_type\") else model.config.architectures[0]\n",
    "tokenizer = get_tokenizer(\n",
    "    quant_args.model_dir,\n",
    "    max_seq_len=quant_args.seq_len,\n",
    "    model_type=model_type,\n",
    "    trust_remote_code=quant_args.trust_remote_code,\n",
    ")\n",
    "\n",
    "multimodal = True if model_type in [\"mllama\", \"llama4\", \"gemma3_mllm\"] else False\n",
    "if multimodal:\n",
    "    processor = AutoProcessor.from_pretrained(quant_args.model_dir)\n",
    "    if quant_args.model_export is not None:\n",
    "        export_dir = Path(quant_args.quant_out_dir)\n",
    "        export_dir.mkdir(parents=True, exist_ok=True)\n",
    "        processor.save_pretrained(quant_args.quant_out_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3.2 Define PTQ calibration data loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main_device = model.device if quant_args.multi_gpu or quant_args.multi_device else quant_args.device\n",
    "calib_dataloader = get_calib_dataloader(\n",
    "    dataset_name=quant_args.calib_dataset,\n",
    "    processor=processor if multimodal else None,\n",
    "    tokenizer=tokenizer,\n",
    "    batch_size=quant_args.batch_size,\n",
    "    num_calib_data=quant_args.num_calib_data,\n",
    "    seqlen=quant_args.seq_len,\n",
    "    device=main_device,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3.3 Prepare PTQ configuration and Perform PTQ.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_config_type = model.config.model_type if hasattr(model.config, \"model_type\") else model.config.architectures[0]\n",
    "template = LLMTemplate.get(model_config_type)\n",
    "# add layer_quant_config\n",
    "layer_config, algo_configs = {}, {}  # NOTE for this example, skip layer_config & algo_configs\n",
    "\n",
    "quant_config = template.get_config(\n",
    "    scheme=quant_args.quant_scheme,\n",
    "    algorithm=quant_args.quant_algo,\n",
    "    kv_cache_scheme=quant_args.kv_cache_dtype,\n",
    "    min_kv_scale=quant_args.min_kv_scale,\n",
    "    layer_config=layer_config,\n",
    "    attention_scheme=quant_args.attention_dtype,\n",
    "    exclude_layers=quant_args.exclude_layers,\n",
    ")\n",
    "\n",
    "quantizer = ModelQuantizer(quant_config, quant_args.multi_device)\n",
    "model = quantizer.quantize_model(model, calib_dataloader)\n",
    "quant_args.exclude_layers = quantizer.config.exclude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3.4 Test the PPL after PTQ (Optional)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_args.use_ppl_eval_model = True\n",
    "eval_model(\n",
    "    quant_args,\n",
    "    model,\n",
    "    main_device,\n",
    "    save_metrics_to_csv=quant_args.save_metrics_to_csv,\n",
    "    output_dir=quant_args.metrics_output_dir,\n",
    "    multimodal=multimodal,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step4.1 Modify the LLM & make only adapter trainable\n",
    "\n",
    "We replace every `QuantLinear` with `QLoRaQuantLinear`. As `QLoRaQuantLinear` contains the adapter layer, which is designed for QLoRA training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace every QuantLinear to QLoRaQuantLinear\n",
    "from qlora_training import trans_quant_linear_2_qlora_quantLinear\n",
    "\n",
    "trans_quant_linear_2_qlora_quantLinear(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only make the adapter's param trainable\n",
    "from qlora_training import disable_adapters, mark_only_lora_layer_as_trainable\n",
    "\n",
    "mark_only_lora_layer_as_trainable(model)\n",
    "# take the adapter take effect in the forward path\n",
    "disable_adapters(model, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step4.2 Prepare training dataset and Trainer "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
    "    quant_args.model_dir, model_max_length=training_args.model_max_length\n",
    ")\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "\n",
    "data_module = make_supervised_data_module(\n",
    "    dataset=data_args.train_dataset,\n",
    "    tokenizer=train_tokenizer,\n",
    "    cache_dir=data_args.cache_dir,\n",
    "    train_size=data_args.train_size,\n",
    "    eval_size=data_args.eval_size,\n",
    ")\n",
    "\n",
    "# to save GPU memory during training.\n",
    "model.enable_input_require_grads()\n",
    "if model.supports_gradient_checkpointing:\n",
    "    model.config.use_cache = False\n",
    "    model.gradient_checkpointing_enable()\n",
    "    model.enable_input_require_grads()\n",
    "    print(f\"Gradient Checkpointing: {model.is_gradient_checkpointing}\")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    processing_class=tokenizer,\n",
    "    args=training_args,\n",
    "    **data_module,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step4.3 Merge the adapter weight into linear weight\n",
    "\n",
    "Merge the lora_B @ lora_A into linear's weight, this can let the model more compact and easy to deploy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qlora_training import merge_weight\n",
    "\n",
    "merge_weight(trainer.model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# freeze the model quantizer\n",
    "model = quantizer.freeze(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n[INFO]: Evaluating ...\")\n",
    "quant_args.use_ppl_eval_model = True\n",
    "eval_model(\n",
    "    quant_args,\n",
    "    model,\n",
    "    main_device,\n",
    "    save_metrics_to_csv=quant_args.save_metrics_to_csv,\n",
    "    output_dir=quant_args.metrics_output_dir,\n",
    "    multimodal=multimodal,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Appendix:\n",
    "\n",
    "We also provide a end-to-end example code, user can select the proper arguments to conduct the experiments.\n",
    "\n",
    "```shell\n",
    "export CUDA_VISIBLE_DEVICES=7\n",
    "cd /Quark/examples/torch/language_modeling/llm_qat/qlora_training/\n",
    "# using the Llama-3.2-1B-Instruct as example \n",
    "\n",
    "python qlora_training.py \n",
    "    --model_dir={PATH}/pretrained_models/meta-llama/Llama-3.2-1B-Instruct \n",
    "    --cache_dir={PATH}/hug_data \n",
    "    --kv_cache_dtype=fp8  \n",
    "    --quant_scheme=mxfp4 \n",
    "    --num_calib_data=128  \n",
    "    --train_size=100 \n",
    "    --eval_size=10 \n",
    "    --learning_rate=1e-4  \n",
    "    --lr_scheduler_type=linear\n",
    "    --train_dataset=Daring-Anteater\n",
    "```\n",
    "\n",
    "\n",
    "\n",
    "### Some results\n",
    "\n",
    "| Model              | Llama-3.2-1B-Instruct | gpt-oss-20b/ |\n",
    "| ------------------ | --------------------- | ------------ |\n",
    "| ORG Model PPL      | 13.4199               | null         |\n",
    "| PTQ‘s PPL          | 18.8498               | null         |\n",
    "| QLoRA's PPL        | 16.9987               | null         |\n",
    "| PTQ's GPU memory   | 3620MB                | 48GB         |\n",
    "| QLoRA's GPU memory | 18GB                  | 66GB         |\n",
    "\n",
    "### NOTE： \n",
    "\n",
    "   In the future, we will conduct more experiments. Next Step, we will focus on the `gpt-oss-20b` and prepare the corresponding test metrics."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "quark_nv_p310_torch27",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
