![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "markdown",
"source": [
"# nanoLoRA\n",
"> A Minimalistic Implementation of Low-Rank Adaptation\n",
"\n",
"❤️ Created by [@maximelabonne](https://twitter.com/maximelabonne) as part of the 🗣️ [Large Language Model Course](https://github.com/mlabonne/llm-course).\n",
"\n",
"1. **Low-Rank Decomposition**: LoRA represents the updates to the weight matrices with a low-rank decomposition. For a pre-trained weight matrix $W_0$, its update is represented by $W_0 + BA$, where $B$ and $A$ are matrices with lower rank $r$, and $r$ is less than the minimum of the dimensions of $W_0$.\n",
"\n",
"2. **Freezing Pre-trained Weights**: During training, $W_0$ remains constant and does not receive gradient updates. The trainable parameters are contained in matrices $A$ and $B$, thereby constraining the updates to a lower intrinsic rank.\n",
"\n",
"3. **Forward Pass Modification**: The modified forward pass includes both the original weight matrix and the low-rank update, yielding $h=W_0 x+BAx$.\n",
"\n",
"4. **Deployment Efficiency**: When deployed, LoRA can explicitly compute and store $W=W_0 + BA$ and perform inference as usual. This approach introduces no additional inference latency compared to a fully fine-tuned model.\n"
],
"metadata": {
"id": "R-PRf6eGviwH"
}
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"id": "UI7HPDt_51_C"
},
"outputs": [],
"source": [
"!pip install -q transformers datasets"
]
},
{
"cell_type": "code",
"source": [
"from datasets import load_dataset\n",
"from transformers import DataCollatorForLanguageModeling\n",
"from transformers import GPT2Model, GPT2Config, GPT2Tokenizer\n",
"\n",
"def tokenize(element):\n",
" outputs = tokenizer(\n",
" element[\"text\"],\n",
" truncation=True,\n",
" max_length=context_length,\n",
" return_overflowing_tokens=True,\n",
" return_length=True,\n",
" )\n",
" input_batch = []\n",
" for length, input_ids in zip(outputs[\"length\"], outputs[\"input_ids\"]):\n",
" if length == context_length:\n",
" input_batch.append(input_ids)\n",
" return {\"input_ids\": input_batch}\n",
"\n",
"dataset = load_dataset('wikitext', 'wikitext-2-raw-v1')\n",
"\n",
"tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
"tokenizer.pad_token = tokenizer.eos_token\n",
"\n",
"context_length = 128\n",
"tokenized_datasets = dataset.map(\n",
" tokenize, batched=True, remove_columns=dataset[\"train\"].column_names\n",
")\n",
"\n",
"data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 337,
"referenced_widgets": [
"88e51d370c39487a8d21212a72a40701",
"bc968b86e959408585cb7087b1df7004",
"5cb49b91d78a4b92825f0983383e95d4",
"2501f7a4382746f09c4c810400c6ba59",
"9dd182e7908c4fd7a9cc9fceaa4bfcc0",
"ff758a2fbc67439bbb58f50cfc05bd3c",
"f4aad55b509344eda8526755eb62f5e1",
"e23e7eec5da1427491e77d69f602723a",
"163fdb5c3cbe4e4a969b577085a759d5",
"a8c77a0dbbc1490aa169c7a628f30e66",
"9de9e4c3f7524f7d8a6722116b60a655",
"0bd01c4341804a448d6010ed308a5282",
"ee1f5a205a124d7aaf750925609c25c9",
"a5f2eb5027a441bb82d56bde705e1990",
"36ad125639c8470b9219fc9db2544dd0",
"380b37ab8bd54cdfa4a46f761fa8c39e",
"2ee7eba2c86e4e4fae4ebf1ecb35e53f",
"175db7688b574ba289f8efc451dcc46b",
"65c7c6f768194fff8990962acd91824a",
"998d80875e9a4c77bc3b7191de16babd",
"217e998925c4404bb2d03964ccd91ee1",
"abf166e0da6747b18650b332033c4be9",
"0b3167e5e91a4366bd1366b5a3ee7108",
"2a9eed8424a349508d82c8db1ec13bdd",
"31d44a7bb0894242a2122b8ce92c59b2",
"441df0d778e94eb1a37ab5459e539be7",
"142099da23a44b2db575bfc559f9ac51",
"27ad27b592d34da4b54e6353adbdbd72",
"15712126b47642a9b8c6963ec8c0f639",
"da841539ce824c9494d739b41e3a9ff5",
"db3d66016c41483cbd407984b00a6cb4",
"a8bd2f0c901045f2bce630cd9a3f4a14",
"0265ddffc63044dbaabb9ea3b4debc7b",
"9b9159a288244d2db0542effbe2b40a0",
"28fd0acd6abc4d48bf618e3e17cf3b03",
"87dc0a735e1c43a3a1a744cdcfce9e32",
"8b51e977b82d4e05ac54f08b06196f27",
"49f529030ceb473e9aa36bac22b77f17",
"e3a3d6dc69204f2b8f742241257ead99",
"a3a6aaf5ccc54c14851425b4d761363d",
"2c068da47f064c3e8df9f9ed0ffaaa8c",
"0502279be3124defbdef35a05d0dc4af",
"ca7c1452f3f4413fbd720b9b043ed540",
"2bfb99bbbeaa46f3b7b959dc146f7b42",
"cfd15702aa944969b1102d8f3422748e",
"96c282c5878a4d18983aea8b5825e7c3",
"932a26af25cf4ce2abd338c16aaa9aa2",
"b10bb79cdcc4474fb7896a5d10580253",
"56d028d8b95b4f8282841fa62b12c4c9",
"eeb1186a3f27435bb99c280d22867fb8",
"4b09af85e0f246f38ab7203c3eebcc11",
"10ce9c2b121b4b5fbab6e4d36dd115fb",
"ad478d1416f144649d295158a30bbf8d",
"1856e01e7ea648b78a8a8c3c2bcbb9ac",
"162480d103254d2a8fe1c6b6da087cde",
"b26b51c7a2b645698e0804dbb2ea9b02",
"e2e7ce1ac12e4eaebb422e6ebdaaec4e",
"81aeb9aede8545d88eeae05513fcdec3",
"72ad35071b684d3c880f9fbc6df1713b",
"e31735c84a7d4dfb836d16a0699c095e",
"b77bbda3a7c84e5b9f838783629a8fd9",
"5c64a41cbccd4a688b2705409a1c7b71",
"a9196f018e7545df884644cb0d49dd0b",
"c8c6c4e61e51417e9888c8bade1b9d34",
"680c886c17bc4a0db03fff5e394b3f37",
"c0fde925024c4a6b9ca7da408cbb7c45",
"9593bbb01b8240b19c880b4a7c2328cf",
"8fd936bc6b55402c93419eb0a4424196",
"6d40dc917b464673af71517733c4934d",
"f5a543c831514eeb9ec68c92faedaeee",
"ae74eb0a595a4cbb830969a4ae68f709",
"ad676eedaab34389807e3d55acde46f6",
"c8bb4ed0327949e0a89ea922d1ed79e3",
"2e29360c4c4b479f9b7592ca71771e1c",
"4ae52492d017473ea4e85ebd2c491476",
"26fca6a701e94a9bb23847c2175b3c8d",
"c5a90735c23a4732b57541c7dbb3d7d8",
"371b1138cbdf4ebda9d89219fb08c7bb",
"fe1c922a53334a41a7e91e5f648a4bca",
"df5c97ba4822468d9155df2141e13f3e",
"8a22da4f94274b97aef0aac8098a0dd3",
"efda83fd562b4a2196926a307d76fe7a",
"2bcf9da3f73a4be08f6689bf078169c0",
"806dd64f7f7d433da1dd6787ba0aef85",
"dbe0228febda48749773b879aadee6ec",
"e16fd1d8b9964497aa1361cc421512ea",
"a4565932d1e243cfad243a5afd57d25e",
"111ff211564d498cbf7c532c43a0d49e",
"5726ac78e4b34ab49a1f9861d5a4d36b",
"76761d97e10b4703aed72d21c4e111ff",
"115c52d4f6244fd4919df84319dba888",
"19b15cb308414aad9069f3cbe9eb5f51",
"63123697618f40778c57690cf44ed7bd",
"57b87ce9cef64cb08b57612bcc264976",
"054efa23f043447db19e5d84676fb43b",
"5f92a0323cb246f5bd022bc47626a3ae",
"df1b1f711d7f4020b2547fe052690d54",
"acf1e8dbc8e64f91ad3233e6a64fd805",
"e6aeae0d76f54544bcc294eb5b53249f",
"424555ef71d44e5187e1bfcd420d8094",
"ed48fea4eaef4186ac412fb214636406",
"1e7580ef696649d49500c30812d8ac98",
"8ea4af3fb46e45329e407779849bceb8",
"52e81d828e6a49f2ae1fd52eb941d55a",
"ed259669f9a6495fa05b41a9541a0ce5",
"159ee8dc836b49d69cc5e57472523bc8",
"6205b074a9f1445eabe40bb145cc2401",
"fc16b104f06a440a8c1db7d02cee4da2",
"0e4e5f1d4c1d48bb83acc442d6e7cc1e",
"8ba4fc9575824920aa7245905ef1bc7c"
]
},
"id": "diEMvtXR7f7Q",
"outputId": "88b55d05-980c-4eb7-dfa6-17092426ec36"
},
"execution_count": 2,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"Downloading data: 0%| | 0.00/4.72M [00:00 0, \"r must be > 0\"\n",
"\n",
" self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))\n",
" self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))\n",
" self.scaling = lora_alpha / r\n",
" self.weight.requires_grad = False\n",
" self.merged = False\n",
" self.reset_lora_parameters()\n",
"\n",
" def reset_lora_parameters(self):\n",
" super().reset_parameters()\n",
" nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))\n",
" nn.init.zeros_(self.lora_B)\n",
"\n",
" def train(self, mode: bool = True):\n",
" super().train(mode)\n",
" if mode:\n",
" if self.merged:\n",
" # Make sure that the weights are not merged for training\n",
" self.weight.data -= (self.lora_B @ self.lora_A) * self.scaling\n",
" self.merged = False\n",
" else:\n",
" if not self.merged:\n",
" # Merge the weights for inference\n",
" self.weight.data += (self.lora_B @ self.lora_A) * self.scaling\n",
" self.merged = True\n",
"\n",
" def forward(self, x: torch.Tensor):\n",
" if not self.merged:\n",
" out = F.linear(x, self.weight, bias=self.bias)\n",
" out += (x @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling\n",
" return out\n",
" else:\n",
" return F.linear(x, self.weight, bias=self.bias)"
],
"metadata": {
"id": "Bjwhp1B2bY5z"
},
"execution_count": 3,
"outputs": []
},
{
"cell_type": "code",
"source": [
"from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig\n",
"\n",
"# Initialize the model\n",
"config = AutoConfig.from_pretrained(\n",
" \"gpt2\",\n",
" vocab_size=len(tokenizer),\n",
" n_ctx=1024,\n",
" bos_token_id=tokenizer.bos_token_id,\n",
" eos_token_id=tokenizer.eos_token_id,\n",
")\n",
"model = GPT2LMHeadModel(config)\n",
"\n",
"# Freeze all parameters in the model\n",
"for param in model.parameters():\n",
" param.requires_grad = False\n",
"\n",
"# Modify the layers and unfreeze these parameters\n",
"for i in range(config.n_layer):\n",
" model.transformer.h[i].mlp.c_fc = nanoLoRA(config.n_embd, config.n_embd, r=8)\n",
" model.transformer.h[i].mlp.c_proj = nanoLoRA(config.n_embd, config.n_embd, r=8)\n",
"\n",
"num_train_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
"print(f'The model has {num_train_params:,} trainable parameters.')"
],
"metadata": {
"id": "9oWWdIkJ56Wd",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "72c38f6a-98e8-4278-e998-a3e387f29277"
},
"execution_count": 4,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"The model has 313,344 trainable parameters.\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"from transformers import Trainer, TrainingArguments\n",
"\n",
"args = TrainingArguments(\n",
" output_dir=\"results\",\n",
" per_device_train_batch_size=32,\n",
" per_device_eval_batch_size=32,\n",
" evaluation_strategy=\"steps\",\n",
" eval_steps=50,\n",
" logging_steps=50,\n",
" gradient_accumulation_steps=1,\n",
" num_train_epochs=1,\n",
" weight_decay=0.1,\n",
" warmup_steps=10,\n",
" lr_scheduler_type=\"cosine\",\n",
" learning_rate=5e-4,\n",
" fp16=True,\n",
" report_to=\"tensorboard\",\n",
")\n",
"\n",
"trainer = Trainer(\n",
" model=model,\n",
" tokenizer=tokenizer,\n",
" args=args,\n",
" data_collator=data_collator,\n",
" train_dataset=tokenized_datasets[\"train\"],\n",
" eval_dataset=tokenized_datasets[\"validation\"],\n",
")\n",
"\n",
"trainer.train()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 321
},
"id": "-eHDB7xB2HMk",
"outputId": "8007f1f4-3c57-4f1f-abd3-2b76a96713d2"
},
"execution_count": 5,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"/usr/local/lib/python3.10/dist-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
" warnings.warn(\n"
]
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"| Step | \n", "Training Loss | \n", "Validation Loss | \n", "
|---|---|---|
| 50 | \n", "10.368600 | \n", "9.695271 | \n", "
| 100 | \n", "9.660800 | \n", "9.591634 | \n", "
| 150 | \n", "9.634500 | \n", "9.585637 | \n", "
| 200 | \n", "9.623000 | \n", "9.583360 | \n", "
| 250 | \n", "9.613100 | \n", "9.582972 | \n", "
"
]
},
"metadata": {}
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"TrainOutput(global_step=266, training_loss=9.771246601764421, metrics={'train_runtime': 96.2756, 'train_samples_per_second': 88.164, 'train_steps_per_second': 2.763, 'total_flos': 279368589901824.0, 'train_loss': 9.771246601764421, 'epoch': 1.0})"
]
},
"metadata": {},
"execution_count": 5
}
]
},
{
"cell_type": "code",
"source": [
"trainer.evaluate()"
],
"metadata": {
"id": "grVjJRyP8Cvo",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 124
},
"outputId": "14d33957-ccc6-4fd6-db88-6aa977c83cf9"
},
"execution_count": 6,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"\n",
" \n",
" \n",
" [28/28 00:03]\n",
" \n",
" "
]
},
"metadata": {}
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"{'eval_loss': 9.582963943481445,\n",
" 'eval_runtime': 3.3787,\n",
" 'eval_samples_per_second': 262.824,\n",
" 'eval_steps_per_second': 8.287,\n",
" 'epoch': 1.0}"
]
},
"metadata": {},
"execution_count": 6
}
]
}
]
}