updates peft model saving

notebooks/drafts/olmo_peft.ipynb

@@ -542,6 +542,21 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import PeftModel, PeftModelForCausalLM, AutoPeftModelForCausalLM\n",
+    "\n",
+    "final_peft_model = AutoPeftModelForCausalLM.from_pretrained(\n",
+    "    \"/Volumes/daniel_liden/fine_tuning/assets/final/\",\n",
+    "    device_map=\"auto\",\n",
+    "    load_in_8bit=True,\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -554,7 +569,7 @@
     "with mlflow.start_run(run_id=last_run_id):\n",
     "    mlflow.log_params(lora_config.to_dict())\n",
     "    mlflow.transformers.log_model(\n",
-    "        transformers_model={\"model\": final_model, \"tokenizer\": tokenizer},\n",
+    "        transformers_model={\"model\": final_peft_model, \"tokenizer\": tokenizer},\n",
     "        signature=signature,\n",
     "        artifact_path=\"model\",  # This is a relative path to save model files within MLflow run\n",
     "    )"