MistraLlama_QLoRA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb73ceb-56a3-4702-a306-e9c005280869",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig, set_seed\n",
    "from peft import LoraConfig\n",
    "from trl import SFTTrainer, setup_chat_format, DataCollatorForCompletionOnlyLM\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "\n",
    "set_seed(42)\n",
    "\n",
    "# change me!\n",
    "modelpath = \"models/Mistral-7B-v0.1\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    modelpath,    \n",
    "    quantization_config = BitsAndBytesConfig(\n",
    "        load_in_4bit = True,\n",
    "        bnb_4bit_compute_dtype = torch.bfloat16,\n",
    "        bnb_4bit_quant_type = \"nf4\",\n",
    "    ),\n",
    "    attn_implementation = \"flash_attention_2\",  \n",
    "    use_cache = False,\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(modelpath, use_fast = False)\n",
    "\n",
    "model, tokenizer = setup_chat_format(model, tokenizer)\n",
    "if tokenizer.pad_token in [None, tokenizer.eos_token]: \n",
    "    tokenizer.pad_token = tokenizer.unk_token\n",
    "\n",
    "dataset = load_dataset(\"g-ronimo/oasst2_top4k_en\")\n",
    "\n",
    "training_arguments = TrainingArguments(\n",
    "    output_dir = \"out_OA_Mistral\",\n",
    "    evaluation_strategy = \"steps\",\n",
    "    label_names = [\"labels\"],\n",
    "    per_device_train_batch_size = 16,\n",
    "    gradient_accumulation_steps = 1,\n",
    "    save_steps = 25_000,\n",
    "    eval_steps = 250,\n",
    "    logging_steps = 1, \n",
    "    learning_rate = 2e-5,\n",
    "    num_train_epochs = 3,\n",
    "    lr_scheduler_type = \"constant\",\n",
    "    optim = 'paged_adamw_8bit',\n",
    "    fp16 = True,\n",
    "    gradient_checkpointing = True,\n",
    "    group_by_length = True,\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model = model,\n",
    "    tokenizer = tokenizer,\n",
    "    train_dataset = dataset[\"train\"],\n",
    "    eval_dataset = dataset['test'],\n",
    "    data_collator = DataCollatorForCompletionOnlyLM(\n",
    "        instruction_template = \"<|im_start|>user\", \n",
    "        response_template = \"<|im_start|>assistant\", \n",
    "        tokenizer = tokenizer, \n",
    "        mlm = False),\n",
    "    max_seq_length = 512,\n",
    "    peft_config = LoraConfig(\n",
    "        target_modules = \"all-linear\", \n",
    "        modules_to_save = [\"lm_head\", \"embed_tokens\"]),\n",
    "    dataset_kwargs = dict(add_special_tokens = False),\n",
    "    args = training_arguments,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}