Lora scheduler

.gitignore

@@ -203,3 +203,5 @@ benchmarks/*.json
 # Linting
 actionlint
 shellcheck*/
+
+out/

examples/generate_rand_loras.ipynb

@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 148,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1, Loss: 7.922913551330566\n",
+      "Epoch 2, Loss: 7.8935418128967285\n",
+      "Epoch 3, Loss: 7.1147661209106445\n",
+      "Epoch 4, Loss: 6.443078517913818\n",
+      "Epoch 5, Loss: 4.377083778381348\n",
+      "Epoch 6, Loss: 2.9477269649505615\n",
+      "Epoch 7, Loss: 1.8892734050750732\n",
+      "Epoch 8, Loss: 0.8281649351119995\n",
+      "Epoch 9, Loss: 0.8178035616874695\n",
+      "Epoch 10, Loss: 0.582257866859436\n",
+      "Epoch 11, Loss: 0.9278958439826965\n",
+      "Epoch 12, Loss: 0.75615394115448\n",
+      "Epoch 13, Loss: 1.5576722621917725\n",
+      "Epoch 14, Loss: 0.056732214987277985\n",
+      "Epoch 15, Loss: 0.17235752940177917\n",
+      "Epoch 16, Loss: 0.09152041375637054\n",
+      "Epoch 17, Loss: 0.13022735714912415\n",
+      "Epoch 18, Loss: 0.23271627724170685\n",
+      "Epoch 19, Loss: 0.20134702324867249\n",
+      "Epoch 20, Loss: 0.03683943673968315\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Tokenize the garbage data\n",
+    "inputs = tokenizer(garbage_data, return_tensors=\"pt\", padding=True, truncation=True)\n",
+    "\n",
+    "# Train LoRA on garbage data (1 epoch as an example)\n",
+    "lora_model.train()\n",
+    "optimizer = torch.optim.AdamW(lora_model.parameters(), lr=5e-3)\n",
+    "\n",
+    "for epoch in range(10):\n",
+    "    optimizer.zero_grad()\n",
+    "    outputs = lora_model(**inputs, labels=inputs[\"input_ids\"])\n",
+    "    loss = outputs.loss\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "    print(f\"Epoch {epoch + 1}, Loss: {loss.item()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 149,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lora_model.save_pretrained(\"../out/lora\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 150,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import PeftModel\n",
+    "\n",
+    "lora_model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(\"../out/base\"), \"../out/lora\")\n",
+    "base_model = AutoModelForCausalLM.from_pretrained(\"../out/base\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/anaconda3/envs/vllm2/lib/python3.12/site-packages/transformers/generation/configuration_utils.py:590: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.0` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "base\n",
+      "lora:  I'm not sure if I'm going to be able to do this, but I'm going to be able to do this.\n",
+      "I'm going to be able to do this.\n",
+      "I'm going to be able to do this.\n",
+      "\n",
+      "lora\n",
+      "lora:  Relate receive continue development challenge quite. Continue development challenge quite. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue. Continue\n"
+     ]
+    },
+    {
+     "ename": "",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
+      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
+      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
+      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
+     ]
+    }
+   ],
+   "source": [
+    "input_text = f\"lora: \"\n",
+    "inputs = tokenizer(input_text, return_tensors=\"pt\", padding=True, truncation=True)\n",
+    "input_ids = inputs[\"input_ids\"]\n",
+    "attention_mask = inputs[\"attention_mask\"]\n",
+    "\n",
+    "lora_model.eval()\n",
+    "\n",
+    "models = [\n",
+    "    (\"base\", base_model), \n",
+    "    (\"lora\", lora_model.to('cpu')),\n",
+    "]\n",
+    "for name, model in models:\n",
+    "    outputs = model.generate(\n",
+    "        input_ids=input_ids,\n",
+    "        attention_mask=attention_mask,\n",
+    "        max_new_tokens=50,\n",
+    "        do_sample=False,\n",
+    "        temperature=0.0,  # No randomness\n",
+    "        pad_token_id=model.config.pad_token_id,  # Explicitly set\n",
+    "        eos_token_id=model.config.eos_token_id   # Explicitly set\n",
+    "    )\n",
+    "    print(name)\n",
+    "    print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "vllm2",
+   "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.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/generate_rand_loras.py

@@ -0,0 +1,103 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from peft import LoraConfig, get_peft_model
+from faker import Faker
+from datasets import Dataset
+
+OUT_DIR = "out"
+
+# Load a base model and tokenizer
+base_model_name = "meta-llama/Llama-3.2-1B"
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
+
+NB_WORDS = 1024 # Generate a long sentence so that the model keeps running for the entirety of max_tokens during benchmarking
+
+def generate_fake_data(num_samples):
+    fake = Faker()
+    sentences = [f"lora: {fake.sentence(nb_words=NB_WORDS)}" for _ in range(num_samples)]
+    return {"text": sentences}
+
+def train_lora(name):
+    lora_config = LoraConfig(
+        r=8,  # Low-rank dimension
+        lora_alpha=16,  # Scaling factor
+        lora_dropout=0.1,  # Dropout for LoRA layers
+        bias="none",  # Bias setting for LoRA layers
+        task_type="CAUSAL_LM"  # Task type for the model
+    )
+
+    # Create the LoRA model
+    model = AutoModelForCausalLM.from_pretrained(base_model_name)
+    model = get_peft_model(model, lora_config)
+
+    # Generate fake training data
+    fake_data = generate_fake_data(num_samples=1)
+    dataset = Dataset.from_dict(fake_data)
+
+    # Tokenize dataset
+    def tokenize_function(examples):
+        tokens = tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+        tokens["labels"] = tokens["input_ids"].copy()
+        return tokens
+
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+
+    # Define training arguments
+    training_args = TrainingArguments(
+        output_dir=f"{OUT_DIR}/{name}",
+        overwrite_output_dir=True,
+        num_train_epochs=100,
+        per_device_train_batch_size=8,
+        logging_steps=1,
+        learning_rate=1e-3,
+        eval_strategy="no",
+        report_to=None
+    )
+
+    # Create the Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        tokenizer=tokenizer
+    )
+
+    # Train the model
+    trainer.train()
+
+    # Save the trained LoRA model
+    trainer.save_model(f"{OUT_DIR}/{name}")
+    print(f"Model saved to {OUT_DIR}/{name}")
+
+def load_model(name):
+    model = AutoModelForCausalLM.from_pretrained(f"{OUT_DIR}/{name}")
+    return model
+
+def test(model):
+    model.eval()
+    input_text = f"lora:"
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+
+    outputs = model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        max_new_tokens=64,
+        do_sample=False,
+        temperature=None,
+        top_p=None,
+        pad_token_id=model.config.pad_token_id,  # Explicitly set
+        eos_token_id=model.config.eos_token_id   # Explicitly set
+    )
+    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+
+if __name__ == "__main__":
+    # Train
+    for i in range(10):
+        train_lora(f"lora{i}")
+
+    # Test
+    test(AutoModelForCausalLM.from_pretrained(base_model_name))
+    for i in range(10):
+        test(load_model(f"lora{i}"))