![\"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
+ }
+ ]
+ }
+ ]
+}
\ No newline at end of file