wehub-resource-sync
caf324b09d
Build documentation / build (push) Failing after 0s
Deploy "method_comparison" Gradio to Spaces / deploy (push) Has been cancelled
Deploy "PEFT shop" Gradio app to Spaces / deploy (push) Has been cancelled
tests on transformers main / tests (push) Has been cancelled
tests / check_code_quality (push) Has been cancelled
tests / tests (ubuntu-latest, 3.10) (push) Has been cancelled
tests / tests (ubuntu-latest, 3.11) (push) Has been cancelled
tests / tests (ubuntu-latest, 3.12) (push) Has been cancelled
tests / tests (ubuntu-latest, 3.13) (push) Has been cancelled
tests / tests (windows-latest, 3.10) (push) Has been cancelled
tests / tests (windows-latest, 3.11) (push) Has been cancelled
tests / tests (windows-latest, 3.12) (push) Has been cancelled
tests / tests (windows-latest, 3.13) (push) Has been cancelled
Secret Leaks / trufflehog (push) Has been cancelled
CI security linting / zizmor latest via Cargo (push) Has been cancelled
440 lines
17 KiB
Plaintext
440 lines
17 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "e2b69494",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Block-Diagonal LoRA for Eliminating Communication Overhead in Tensor Parallel LoRA Serving"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "2fb2b221",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Introduction\n",
|
|
"Block-Diagonal LoRA (BD-LoRA) is a LoRA variant in which some LoRA factors are constrained to be block-diagonal. This allows faster serving by eliminating communication overheads \n",
|
|
"when running inference on multiple GPUs. Despite the block-diagonal constraint, BD-LoRA is similarly performant to vanilla LoRA at similar parameter counts.\n",
|
|
"\n",
|
|
"BD-LoRA is designed to be used with tensor parallelism, which means sharding the weights of a model among multiple GPUs. A popular sharding strategy is the [Megatron Sharding Strategy](https://arxiv.org/abs/1909.08053). For two linear layers $W_1$, $W_2$ that follow each other (for example the up and down projections in a transformer MLP module), we will shard the first layer in a column-parallel way (which requires LoRA B to be block-diagonal) and the second layer in a row-parallel way (which requires LoRA A to be block-diagonal). For the attention module, this can be similarly achieved by taking the Q, K and V projections together as $W_1$ and the out projection as $W_2$, sharding accordingly. This sharding allows a compatible inference engine to distribute each block-diagonal shard over a a different GPU, cutting the need to communicate partial results among GPUs. In the image below, you can see the exact sharding strategy and how this saves computational efforts.\n",
|
|
"\n",
|
|
"Paper: https://arxiv.org/html/2510.23346v1\n",
|
|
"\n",
|
|
"<div>\n",
|
|
"<img src=\"bdlora-sharding.png\" width=\"800\"/>\n",
|
|
"</div>\n",
|
|
"\n",
|
|
"### Performance, rank and parameter count\n",
|
|
"BD-LoRA achieves similar performance to LoRA (see image below, or the `method_comparison` folder in the peft repository root) at the same parameter count. However, as every other factor in BD-LoRA is block-diagonal, a BD-LoRA adapter will have less parameters than a LoRA adapter at the same rank. The performance of BD-LoRA is only competitive when the rank is then increased accordingly. We provide example code for rank-matching at the end of this example notebook.\n",
|
|
"\n",
|
|
"<div>\n",
|
|
"<img src=\"bdlora-performance.png\" width=\"600\"/>\n",
|
|
"</div>\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "40eea544",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/efs/aconzel/workspace/peft/bdlora-peft/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
|
" from .autonotebook import tqdm as notebook_tqdm\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"from peft.tuners import BdLoraConfig, LoraConfig\n",
|
|
"from peft import get_peft_model\n",
|
|
"from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoModelForCausalLM, AutoTokenizer\n",
|
|
"from datasets import load_dataset\n",
|
|
"import torch"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "5959e4ab",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Quick Start\n",
|
|
"To use BD-LoRA, we can follow standard LoRA-training procedures. We only need to change the `LoraConfig` to a `BdLoraConfig` and specify which LoRA should be block-diagonal. \n",
|
|
"As an example, we will train a LLama-Model in such a way that it can later benefit from inference speed-up as specified in the BD-LoRA paper. However, BD-LoRA can be used with all other models that follow a transformer architecture. \n",
|
|
"\n",
|
|
"As explained in the introduction, we want to shard each module (MLP and attention) in an alternating fashion, first column-parallel with LoRA-B block-diagonal, then row-parallel with LoRA-A block-diagonal. Different from standard MLP modules, Llama also uses a gate projection, which we can fuse together with the up-projection.\n",
|
|
"\n",
|
|
"Therefore, we want the following block-diagonal factors (following the naming convention from the Llama architecture):\n",
|
|
"\n",
|
|
"- LoRA-A Block-Diagonal (Row-parallel sharding): Out (`out_proj`), Down (`down_proj`)\n",
|
|
"- LoRA-B Block-Diagonal (Column-parallel sharding): QKV (`q_proj, k_proj, v_proj`), Up+Gate (`up_proj, gate_proj`)\n",
|
|
"\n",
|
|
"Additionally, we need to know on how many GPUs we want to serve before we start training, as this corresponds to the number of block we will use for each block-diagonal factor. For this experiment, we will use 2 blocks (equivalent to a tensor-parallelism degree of 2). Caveat: For a small model such as Llama 3.2-1B which we are using, one would use a single GPU for serving, and use TP=2 or TP=8 only for larger models, like Llama 3.1-8B or Llama 3.3-70B respectively. "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "6180c1f9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"model_name = \"meta-llama/Llama-3.2-1B\"\n",
|
|
"model = AutoModelForCausalLM.from_pretrained(model_name)\n",
|
|
"tokenizer = AutoTokenizer.from_pretrained(model_name)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "6a50e350",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"trainable params: 44,826,624 || all params: 1,280,641,024 || trainable%: 3.5003\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"target_modules=[\"q_proj\", \"v_proj\", \"k_proj\", \"up_proj\", \"gate_proj\", \"o_proj\", \"down_proj\"]\n",
|
|
"# Set this equal to the number of GPUs you want to serve the model with later\n",
|
|
"nblocks = 2\n",
|
|
"\n",
|
|
"bdlora_config = BdLoraConfig(\n",
|
|
" target_modules_bd_a=[\"o_proj\", \"down_proj\"],\n",
|
|
" target_modules_bd_b=[\"q_proj\", \"v_proj\", \"k_proj\", \"up_proj\", \"gate_proj\"],\n",
|
|
" nblocks=nblocks\n",
|
|
")\n",
|
|
"\n",
|
|
"config = LoraConfig(\n",
|
|
" r=96,\n",
|
|
" # adjust target modules and the ...target_modules_bd attributes according to model architecture (for example renaming)\n",
|
|
" target_modules=target_modules,\n",
|
|
" use_bdlora=bdlora_config,\n",
|
|
" lora_bias=False\n",
|
|
")\n",
|
|
"\n",
|
|
"peft_model = get_peft_model(model, config)\n",
|
|
"peft_model.print_trainable_parameters()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "c1facfea",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Training\n",
|
|
"We train the model for 10 steps, this training block is just intended to showcase how BD-LoRA integrates into other huggingface tools."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"id": "64e4681b",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"\n",
|
|
" <div>\n",
|
|
" \n",
|
|
" <progress value='10' max='10' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
|
|
" [10/10 00:31, Epoch 10/10]\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>1</td>\n",
|
|
" <td>3.137500</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>2</td>\n",
|
|
" <td>3.137500</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>3</td>\n",
|
|
" <td>3.126700</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>4</td>\n",
|
|
" <td>3.104900</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>5</td>\n",
|
|
" <td>3.085900</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>6</td>\n",
|
|
" <td>3.069800</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>7</td>\n",
|
|
" <td>3.056500</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>8</td>\n",
|
|
" <td>3.046100</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>9</td>\n",
|
|
" <td>3.038300</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <td>10</td>\n",
|
|
" <td>3.033100</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table><p>"
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.HTML object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"TrainOutput(global_step=10, training_loss=3.0836387634277345, metrics={'train_runtime': 38.346, 'train_samples_per_second': 66.761, 'train_steps_per_second': 0.261, 'total_flos': 1954507653120000.0, 'train_loss': 3.0836387634277345, 'epoch': 10.0})"
|
|
]
|
|
},
|
|
"execution_count": 4,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"dataset = load_dataset(\"imdb\", split=\"train[:1%]\")\n",
|
|
"\n",
|
|
"tokenizer.pad_token = tokenizer.eos_token\n",
|
|
"def tokenize(batch):\n",
|
|
" return tokenizer(batch[\"text\"], truncation=True, padding=\"max_length\", max_length=128)\n",
|
|
"\n",
|
|
"dataset = dataset.map(tokenize, batched=True, remove_columns=[\"text\"])\n",
|
|
"training_args = TrainingArguments(\n",
|
|
" output_dir=\"./results\",\n",
|
|
" per_device_train_batch_size=8,\n",
|
|
" gradient_accumulation_steps=4,\n",
|
|
" warmup_steps=2,\n",
|
|
" max_steps=10,\n",
|
|
" learning_rate=2e-4,\n",
|
|
" logging_steps=1,\n",
|
|
")\n",
|
|
"\n",
|
|
"trainer = Trainer(\n",
|
|
" model=peft_model,\n",
|
|
" args=training_args,\n",
|
|
" train_dataset=dataset,\n",
|
|
" data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
|
|
")\n",
|
|
"\n",
|
|
"peft_model.config.use_cache = False\n",
|
|
"trainer.train()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "3ac31a38",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Saving Model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "28f49ee4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"peft_model.save_pretrained(\"example_bd_lora_adapter\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "1a3d6fab",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Example Output"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"id": "20c36df8",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"The Batman Trilogy by Christopher Nolan\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"text = \"The Batman Trilogy by Christopher Nolan\"\n",
|
|
"inputs = tokenizer(text, return_tensors=\"pt\").to(model.device) \n",
|
|
"\n",
|
|
"outputs = peft_model.generate(**inputs, max_length=50)\n",
|
|
"decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
|
|
"print(decoded)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "711eb30d",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Investigating the shapes of LoRA Adapters\n",
|
|
"We can check out the adapter shapes to see if they follow the sharding patterns that we have discussed. To make the implementation more memory efficient, \n",
|
|
"the block-diagonal matrices are not saved in a block-diagonal manner, but the blocks are stacked along the non-rank dimensions. \n",
|
|
"\n",
|
|
"For example, if a layer is column sharded, such as the q-proj in Llama, then the LoRA-B factor is block-diagonal. Assume that the q-proj has layer weights (out_features, in_features), \n",
|
|
"then LoRA-A will have shape (rank, in_features), and LoRA-B will have shape (out_features, rank / TP), which corresponds to TP blocks of shape (out_features/TP, rank/TP) each. This can be checked by investigating the weight shapes:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"id": "3660c4de",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Base layer has shape: [512, 2048]\n",
|
|
"LoRA-A (vanilla): [96, 2048]\n",
|
|
"LoRA-B (block-diagonal): [512, 48 ]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"shape_base = list(peft_model.state_dict()['base_model.model.model.layers.0.self_attn.v_proj.base_layer.weight'].shape)\n",
|
|
"shape_a = list(peft_model.state_dict()['base_model.model.model.layers.0.self_attn.v_proj.lora_A.default.weight'].shape)\n",
|
|
"shape_b = list(peft_model.state_dict()['base_model.model.model.layers.0.self_attn.v_proj.lora_B.default.weight'].shape)\n",
|
|
"print(f\"Base layer has shape: [{shape_base[0]}, {shape_base[1]}]\\nLoRA-A (vanilla): [{shape_a[0]}, {shape_a[1]}]\\nLoRA-B (block-diagonal): [{shape_b[0]}, {shape_b[1]} ]\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d95118af",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Matching the rank\n",
|
|
"Assuming we want to achieve the same performance of a LoRA adapter of a given rank, at which rank would we have to train BD-LoRA? We can find this out by matching the number of trainable parameters. A simple iteration over the ranks of the BD-LoRA adapter is sufficient to do that:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"id": "cd4fc45a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"BD-LoRA rank to match vanilla LoRA performance at rank 64: 96 at 45088768 vanilla LoRA params and 44826624 BD-LoRA params.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"def rank_to_params(r: int, bd_lora: bool, nblocks: int):\n",
|
|
" model = AutoModelForCausalLM.from_pretrained(model_name)\n",
|
|
" if bd_lora:\n",
|
|
" config = LoraConfig(\n",
|
|
" r=r,\n",
|
|
" target_modules=target_modules,\n",
|
|
" use_bdlora=bdlora_config,\n",
|
|
" lora_bias=False\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" config = LoraConfig(\n",
|
|
" r=r,\n",
|
|
" # If you use a model different from Llama, change the settings below\n",
|
|
" target_modules=target_modules,\n",
|
|
" lora_bias=False\n",
|
|
" )\n",
|
|
"\n",
|
|
"\n",
|
|
" peft_model = get_peft_model(model, config)\n",
|
|
" return peft_model.get_nb_trainable_parameters()[0]\n",
|
|
"\n",
|
|
"r_orig = 64\n",
|
|
"r = r_orig\n",
|
|
"lora_nparams = rank_to_params(r, False, nblocks)\n",
|
|
"bdlora_nparams = 0\n",
|
|
"while bdlora_nparams < lora_nparams:\n",
|
|
" r += nblocks\n",
|
|
" bdlora_nparams = rank_to_params(r, True, nblocks)\n",
|
|
"# subtract nblocks again to be just under the parameter count of vanilla LoRA, following the original papers methodology\n",
|
|
"print(f\"BD-LoRA rank to match vanilla LoRA performance at rank {r_orig}: {r-nblocks} at {lora_nparams} vanilla LoRA params and {rank_to_params(r-nblocks, True, nblocks)} BD-LoRA params.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "d27b9545",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Integration with vLLM\n",
|
|
"Currently, vLLM has an experimental PR that allows you to use it with BD-LoRA. Clone the github repository and check out the commit of \n",
|
|
"the pull request at https://github.com/vllm-project/vllm/pull/28136#. Then, install vLLM following the usual instructions: https://docs.vllm.ai/en/stable/getting_started/installation/. We assume that you have a hardware setup with at least 2 available GPUs. \n",
|
|
"\n",
|
|
"We have included a script that starts a vLLM server with two BD-LoRA modules at `vllm_server.bash` (you might have to kill this jupyter server beforehand, as it is likely already using your GPU resources).\n",
|
|
"Once the server has started, you can query it via `python3 chat.py \"Please write your message here.\" --target lora1`."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "bdlora-peft",
|
|
"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.11.13"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|