Update training arguments in app.py

app.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,17 +37,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Model Max Length: 1000000000000000019884624838656\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# model_name='TinyLlama/TinyLlama-1.1B-Chat-v0.1'\n",
     "model_name = 'mistralai/Mistral-7B-v0.1'\n",
@@ -63,20 +55,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Max token length train: 1121\n",
-      "Max token length validation: 38\n",
-      "Block size: 2242\n",
-      "{'project_name': './llms/ams_data_train-100_4ba55532-e0b2-478b-9f5b-beb082e1b557', 'model_name': 'mistralai/Mistral-7B-v0.1', 'repo_id': 'ai-aerospace/ams-data-train-100-11b94ea4-2b2b-4db3-9e69-acb5a5d9f3e8', 'train_data': 'train_data', 'data_directory': './fine_tune_data/', 'block_size': 2242, 'model_max_length': 1121, 'logging_steps': -1, 'evaluation_strategy': 'epoch', 'save_total_limit': 1, 'save_strategy': 'epoch', 'mixed_precision': 'fp16', 'lr': 3e-05, 'epochs': 3, 'batch_size': 2, 'warmup_ratio': 0.1, 'gradient_accumulation': 1, 'optimizer': 'adamw_torch', 'scheduler': 'linear', 'weight_decay': 0, 'max_grad_norm': 1, 'seed': 42, 'quantization': 'int4', 'lora_r': 16, 'lora_alpha': 32, 'lora_dropout': 0.05}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Write dataset files into data directory\n",
     "data_directory = './fine_tune_data/'\n",
@@ -147,7 +128,56 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "args_custom=transformers.TrainingArguments(\n",
+    "    per_device_train_batch_size=model_params['batch_size'],\n",
+    "    per_device_eval_batch_size=model_params['batch_size'],\n",
+    "    gradient_accumulation_steps=model_params['gradient_accumulation'],\n",
+    "    warmup_ratio=model_params['warmup_ratio'],\n",
+    "    num_epochs=model_params['epochs'],\n",
+    "    learning_rate=model_params['lr'],\n",
+    "    fp16=True,\n",
+    "    logging_steps=model_params['logging_steps'],\n",
+    "    save_total_limit=model_params['save_total_limit'],\n",
+    "    evaluation_strategy=model_params['evaluation_strategy'],\n",
+    "    metric_for_best_model=\"f1\",\n",
+    "    output_dir='model_outputs',\n",
+    "    logging_dir='model_outputs',\n",
+    "    optim=model_params['optimizer'],\n",
+    "    max_grad_norm=model_params['max_grad_norm'],\n",
+    "    weight_decay=model_params['weight_decay'],\n",
+    "    lr_scheduler_type=model_params['scheduler']\n",
+    ")\n",
+    "\n",
+    "# Args from medium article\n",
+    "args_medium=transformers.TrainingArguments(\n",
+    "    per_device_train_batch_size=8,\n",
+    "    per_device_eval_batch_size=32,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    warmup_steps=100,\n",
+    "    max_steps=12276,\n",
+    "    learning_rate=2e-4,\n",
+    "    fp16=True,\n",
+    "    eval_steps= 1000,\n",
+    "    logging_steps=1000,\n",
+    "    save_steps=1000,\n",
+    "    evaluation_strategy=\"steps\",\n",
+    "    do_eval=True,\n",
+    "    load_best_model_at_end=True,\n",
+    "    metric_for_best_model=\"f1\",\n",
+    "    output_dir='model_outputs',\n",
+    "    logging_dir='model_outputs',\n",
+    "    remove_unused_columns =False, \n",
+    "    report_to='wandb'  # enable logging to W&B\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -221,90 +251,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = SFTTrainer(\n",
-    "    model,\n",
-    "    train_dataset=dataset,\n",
-    "    dataset_text_field=\"text\",\n",
-    "    peft_config=peft_config,\n",
-    "    max_seq_length=model_params['model_max_length']\n",
-    ")\n",
+    "# trainer = SFTTrainer(\n",
+    "#     model,\n",
+    "#     train_dataset=dataset,\n",
+    "#     dataset_text_field=\"text\",\n",
+    "#     peft_config=peft_config,\n",
+    "#     max_seq_length=model_params['model_max_length']\n",
+    "# )\n",
     "\n",
-    "trainer.train()"
+    "# trainer.train()"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4fbe714ca43d4e53aec27f4ce4fb4706",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading builder script:   0%|          | 0.00/6.77k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "826f51589454434b891a94b0d5ef8a73",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading builder script:   0%|          | 0.00/7.36k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "81418551f332492293ee9795f98a62f7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading builder script:   0%|          | 0.00/4.20k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "367f897f76f845d782ebc3f9be4eec4d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading builder script:   0%|          | 0.00/7.55k [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "ename": "NameError",
-     "evalue": "name 'lora_model' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[8], line 18\u001b[0m\n\u001b[1;32m     13\u001b[0m     results\u001b[38;5;241m.\u001b[39mupdate(precision_metric\u001b[38;5;241m.\u001b[39mcompute(predictions\u001b[38;5;241m=\u001b[39mpredictions, references \u001b[38;5;241m=\u001b[39m labels, average\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmacro\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m     15\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m results\n\u001b[1;32m     17\u001b[0m trainer \u001b[38;5;241m=\u001b[39m transformers\u001b[38;5;241m.\u001b[39mTrainer(\n\u001b[0;32m---> 18\u001b[0m     model\u001b[38;5;241m=\u001b[39m\u001b[43mlora_model\u001b[49m,\n\u001b[1;32m     19\u001b[0m     train_dataset\u001b[38;5;241m=\u001b[39mtrain_dataset,\n\u001b[1;32m     20\u001b[0m     eval_dataset\u001b[38;5;241m=\u001b[39mval_dataset,\n\u001b[1;32m     21\u001b[0m     compute_metrics\u001b[38;5;241m=\u001b[39mcompute_metrics,\n\u001b[1;32m     22\u001b[0m     args\u001b[38;5;241m=\u001b[39mtransformers\u001b[38;5;241m.\u001b[39mTrainingArguments(\n\u001b[1;32m     23\u001b[0m         per_device_train_batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m8\u001b[39m,\n\u001b[1;32m     24\u001b[0m         per_device_eval_batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m32\u001b[39m,\n\u001b[1;32m     25\u001b[0m         gradient_accumulation_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m4\u001b[39m,\n\u001b[1;32m     26\u001b[0m         warmup_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m100\u001b[39m,\n\u001b[1;32m     27\u001b[0m         max_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m12276\u001b[39m,\n\u001b[1;32m     28\u001b[0m         learning_rate\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2e-4\u001b[39m,\n\u001b[1;32m     29\u001b[0m         fp16\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     30\u001b[0m         eval_steps\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1000\u001b[39m,\n\u001b[1;32m     31\u001b[0m         logging_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1000\u001b[39m,\n\u001b[1;32m     32\u001b[0m         save_steps\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1000\u001b[39m,\n\u001b[1;32m     33\u001b[0m         evaluation_strategy\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msteps\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     34\u001b[0m         do_eval\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     35\u001b[0m         load_best_model_at_end\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     36\u001b[0m         metric_for_best_model\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mf1\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     37\u001b[0m         output_dir\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodel_outputs\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     38\u001b[0m         logging_dir\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodel_outputs\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     39\u001b[0m         remove_unused_columns \u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, \n\u001b[1;32m     40\u001b[0m         report_to\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mwandb\u001b[39m\u001b[38;5;124m'\u001b[39m  \u001b[38;5;66;03m# enable logging to W&B\u001b[39;00m\n\u001b[1;32m     41\u001b[0m     ),\n\u001b[1;32m     42\u001b[0m )\n\u001b[1;32m     43\u001b[0m trainer\u001b[38;5;241m.\u001b[39mtrain()\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'lora_model' is not defined"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "f1_metric = evaluate.load(\"f1\")\n",
     "recall_metric = evaluate.load(\"recall\")\n",
@@ -328,26 +290,7 @@
     "    train_dataset=model_params['train_data'],\n",
     "    eval_dataset=model_params['validation_data'],\n",
     "    compute_metrics=compute_metrics,\n",
-    "    args=transformers.TrainingArguments(\n",
-    "        per_device_train_batch_size=model_params['batch_size'],\n",
-    "        per_device_eval_batch_size=model_params['batch_size'],\n",
-    "        gradient_accumulation_steps=model_params['gradient_accumulation'],\n",
-    "        warmup_steps=100,\n",
-    "        max_steps=12276,\n",
-    "        learning_rate=model_params['lr'],\n",
-    "        fp16=True,\n",
-    "        eval_steps= 1000,\n",
-    "        logging_steps=1000,\n",
-    "        save_steps=1000,\n",
-    "        evaluation_strategy=model_params['evaluation_strategy'],\n",
-    "        do_eval=True,\n",
-    "        load_best_model_at_end=True,\n",
-    "        metric_for_best_model=\"f1\",\n",
-    "        output_dir='model_outputs',\n",
-    "        logging_dir='model_outputs',\n",
-    "        remove_unused_columns =False, \n",
-    "        report_to='wandb'  # enable logging to W&B\n",
-    "    ),\n",
+    "    args=args_custom\n",
     ")\n",
     "trainer.train()"
    ]

app.py

@@ -105,6 +105,49 @@ for key, value in model_params.items():
 
 print(model_params)
 
+args_custom=transformers.TrainingArguments(
+    per_device_train_batch_size=model_params['batch_size'],
+    per_device_eval_batch_size=model_params['batch_size'],
+    gradient_accumulation_steps=model_params['gradient_accumulation'],
+    warmup_ratio=model_params['warmup_ratio'],
+    num_epochs=model_params['epochs'],
+    learning_rate=model_params['lr'],
+    fp16=True,
+    logging_steps=model_params['logging_steps'],
+    save_total_limit=model_params['save_total_limit'],
+    evaluation_strategy=model_params['evaluation_strategy'],
+    metric_for_best_model="f1",
+    output_dir='model_outputs',
+    logging_dir='model_outputs',
+    optim=model_params['optimizer'],
+    max_grad_norm=model_params['max_grad_norm'],
+    weight_decay=model_params['weight_decay'],
+    lr_scheduler_type=model_params['scheduler']
+)
+
+### Args from medium article
+args_medium=transformers.TrainingArguments(
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=32,
+    gradient_accumulation_steps=4,
+    warmup_steps=100,
+    max_steps=12276,
+    learning_rate=2e-4,
+    fp16=True,
+    eval_steps= 1000,
+    logging_steps=1000,
+    save_steps=1000,
+    evaluation_strategy="steps",
+    do_eval=True,
+    load_best_model_at_end=True,
+    metric_for_best_model="f1",
+    output_dir='model_outputs',
+    logging_dir='model_outputs',
+    remove_unused_columns =False, 
+    report_to='wandb'  # enable logging to W&B
+)
+###
+
 ### Load model and peft config, calculate trainable parameters
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -141,25 +184,6 @@ trainer = transformers.Trainer(
     train_dataset=model_params['train_data'],
     eval_dataset=model_params['validation_data'],
     compute_metrics=compute_metrics,
-    args=transformers.TrainingArguments(
-        per_device_train_batch_size=model_params['batch_size'],
-        per_device_eval_batch_size=model_params['batch_size'],
-        gradient_accumulation_steps=model_params['gradient_accumulation'],
-        warmup_steps=100,
-        max_steps=12276,
-        learning_rate=model_params['lr'],
-        fp16=True,
-        eval_steps= 1000,
-        logging_steps=1000,
-        save_steps=1000,
-        evaluation_strategy=model_params['evaluation_strategy'],
-        do_eval=True,
-        load_best_model_at_end=True,
-        metric_for_best_model="f1",
-        output_dir='model_outputs',
-        logging_dir='model_outputs',
-        remove_unused_columns =False, 
-        report_to='wandb'  # enable logging to W&B
-    ),
+    args=args_custom
 )
 trainer.train()