7.5 Model_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b29f1f17-912d-429d-bd66-3004093fdfcd",
   "metadata": {},
   "source": [
    "# Model training\n",
    "\n",
    "Pretraining is very expensive! Please check costs carefully before starting a pretraining project.\n",
    "\n",
    "We can get a rough estimate your training job cost using [this calculator](https://huggingface.co/training-cluster) from Hugging Face. For training on other infrastructure, e.g. AWS or Google Cloud, please consult those providers for up to date cost estimates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "31fe0bb8-d33b-4064-aedf-ad008d5fbf2c",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fda3010e",
   "metadata": {},
   "source": [
    "## 1. Load the model to be trained\n",
    "\n",
    "Load the upscaled model from the previous lesson:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a966b4d-3b6f-40a0-a912-d6f0d247f088",
   "metadata": {
    "height": 166
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "pretrained_model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"./models/upstage/TinySolar-308m-4k-init\",\n",
    "    device_map=\"cpu\", \n",
    "    torch_dtype=torch.bfloat16,\n",
    "    use_cache=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "970e3729-5101-418b-9d0b-55ffeec9d999",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(32000, 1024)\n",
       "    (layers): ModuleList(\n",
       "      (0-15): 16 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaSdpaAttention(\n",
       "          (q_proj): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "          (k_proj): Linear(in_features=1024, out_features=256, bias=False)\n",
       "          (v_proj): Linear(in_features=1024, out_features=256, bias=False)\n",
       "          (o_proj): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=1024, out_features=4096, bias=False)\n",
       "          (up_proj): Linear(in_features=1024, out_features=4096, bias=False)\n",
       "          (down_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1024, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretrained_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32d08680-8993-4904-9bab-34b27d8cba57",
   "metadata": {},
   "source": [
    "## 2. Load dataset\n",
    "\n",
    "Here we'll update two methods on the `Dataset` object to allow it to interface with the trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3ade55c1-3296-4c83-97a0-2579471829ab",
   "metadata": {
    "height": 489
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, args, split=\"train\"):\n",
    "        \"\"\"Initializes the custom dataset object.\"\"\"\n",
    "        self.args = args\n",
    "        self.dataset = datasets.load_dataset(\n",
    "            \"parquet\",\n",
    "            data_files=args.dataset_name,\n",
    "            split=split\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Returns the number of samples in the dataset.\"\"\"\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        Retrieves a single data sample from the dataset \n",
    "        at the specified index\n",
    "        \"\"\"\n",
    "        # Convert the lists to a LongTensor for PyTorch\n",
    "        input_ids = torch.LongTensor(self.dataset[idx][\"input_ids\"])\n",
    "        labels = torch.LongTensor(self.dataset[idx][\"input_ids\"])\n",
    "\n",
    "        # Return the sample as a dictionary\n",
    "        return {\"input_ids\": input_ids, \"labels\": labels}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504499ec-ff8a-4377-b5f1-f2061a09e681",
   "metadata": {},
   "source": [
    "## 3. Configure Training Arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0eab5882-4547-43ad-a271-4f5db42d8075",
   "metadata": {
    "height": 591
   },
   "outputs": [],
   "source": [
    "from dataclasses import dataclass, field\n",
    "import transformers\n",
    "\n",
    "@dataclass\n",
    "class CustomArguments(transformers.TrainingArguments):\n",
    "    dataset_name: str = field(                           # Dataset configuration\n",
    "        default=\"./parquet/packaged_pretrain_dataset.parquet\")\n",
    "    num_proc: int = field(default=1)                     # Number of subprocesses for data preprocessing\n",
    "    max_seq_length: int = field(default=32)              # Maximum sequence length\n",
    "\n",
    "    # Core training configurations\n",
    "    seed: int = field(default=0)                         # Random seed for initialization, ensuring reproducibility\n",
    "    optim: str = field(default=\"adamw_torch\")            # Optimizer, here it's AdamW implemented in PyTorch\n",
    "    max_steps: int = field(default=30)                   # Number of maximum training steps\n",
    "    per_device_train_batch_size: int = field(default=2)  # Batch size per device during training\n",
    "\n",
    "    # Other training configurations\n",
    "    learning_rate: float = field(default=5e-5)           # Initial learning rate for the optimizer\n",
    "    weight_decay: float = field(default=0)               # Weight decay\n",
    "    warmup_steps: int = field(default=10)                # Number of steps for the learning rate warmup phase\n",
    "    lr_scheduler_type: str = field(default=\"linear\")     # Type of learning rate scheduler\n",
    "    gradient_checkpointing: bool = field(default=True)   # Enable gradient checkpointing to save memory\n",
    "    dataloader_num_workers: int = field(default=2)       # Number of subprocesses for data loading\n",
    "    bf16: bool = field(default=True)                     # Use bfloat16 precision for training on supported hardware\n",
    "    gradient_accumulation_steps: int = field(default=1)  # Number of steps to accumulate gradients before updating model weights\n",
    "    \n",
    "    # Logging configuration\n",
    "    logging_steps: int = field(default=3)                # Frequency of logging training information\n",
    "    report_to: str = field(default=\"none\")               # Destination for logging (e.g., WandB, TensorBoard)\n",
    "\n",
    "    # Saving configuration\n",
    "    # save_strategy: str = field(default=\"steps\")          # Can be replaced with \"epoch\"\n",
    "    # save_steps: int = field(default=3)                   # Frequency of saving training checkpoint\n",
    "    # save_total_limit: int = field(default=2)             # The total number of checkpoints to be saved"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4d62ed",
   "metadata": {},
   "source": [
    "Parse the custom arguments and set the output directory where the model will be saved: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8af55fd9-de91-4f01-8a5b-f16fb6d7b921",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "parser = transformers.HfArgumentParser(CustomArguments)\n",
    "args, = parser.parse_args_into_dataclasses(\n",
    "    args=[\"--output_dir\", \"output\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1184f35a",
   "metadata": {},
   "source": [
    "Setup the training dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "812abcd8-1146-44fc-8e3e-65a5999eea0c",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a64b0794eb94a29a47db58eb1b25a7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_dataset = CustomDataset(args=args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0fb13a",
   "metadata": {},
   "source": [
    "Check the shape of the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0928c414-298d-4524-b078-ad1e844407a7",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input shape:  torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "print(\"Input shape: \", train_dataset[0]['input_ids'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74a3b247-d377-463b-82eb-ec0514c32949",
   "metadata": {},
   "source": [
    "## 4. Run the trainer and monitor the loss\n",
    "\n",
    "First, set up a callback to log the loss values during training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "91dc3a27-9d9a-4caf-bfc5-01f8eddcacf7",
   "metadata": {
    "height": 234
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments, TrainerCallback\n",
    "\n",
    "# Define a custom callback to log the loss values\n",
    "class LossLoggingCallback(TrainerCallback):\n",
    "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
    "        if logs is not None:\n",
    "            self.logs.append(logs)\n",
    "\n",
    "    def __init__(self):\n",
    "        self.logs = []\n",
    "\n",
    "# Initialize the callback\n",
    "loss_logging_callback = LossLoggingCallback()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "816f01b6",
   "metadata": {},
   "source": [
    "Then, create an instance of the Hugging Face `Trainer` object from the `transformers` library. Call the `train()` method of the trainder to initialize the training run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ee517564-e1d4-4e15-8ca9-6d0c0e7ea7e6",
   "metadata": {
    "height": 200
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='30' max='30' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [30/30 06:00, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>4.701100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>4.762400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>4.621400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12</td>\n",
       "      <td>4.724900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15</td>\n",
       "      <td>4.161000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>18</td>\n",
       "      <td>4.067400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>21</td>\n",
       "      <td>4.466300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>24</td>\n",
       "      <td>4.253000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>27</td>\n",
       "      <td>4.844100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>4.433100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=30, training_loss=4.503474426269531, metrics={'train_runtime': 386.1432, 'train_samples_per_second': 0.155, 'train_steps_per_second': 0.078, 'total_flos': 3180342804480.0, 'train_loss': 4.503474426269531, 'epoch': 0.0})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=pretrained_model, \n",
    "    args=args, \n",
    "    train_dataset=train_dataset, \n",
    "    eval_dataset=None,\n",
    "    callbacks=[loss_logging_callback] \n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2f6dac",
   "metadata": {},
   "source": [
    "We can use the code below to save intermediate model checkpoints in our own training run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b4c395d1",
   "metadata": {
    "height": 81
   },
   "outputs": [],
   "source": [
    "# Saving configuration\n",
    "    # save_strategy: str = field(default=\"steps\")          # Can be replaced with \"epoch\"\n",
    "    # save_steps: int = field(default=3)                   # Frequency of saving training checkpoint\n",
    "    # save_total_limit: int = field(default=2)             # The total number of checkpoints to be saved"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5d4e75",
   "metadata": {},
   "source": [
    "### Checking the performance of an intermediate checkpoint\n",
    "\n",
    "Below, we can try generating text using an intermediate checkpoint of the model. This checkpoint was saved after 10,000 training steps. As we'll use the Solar tokenizer and then set up a `TextStreater` object to display the text as it is generated: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "29f45984-36b0-4351-a758-b67f674d9e2e",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, TextStreamer\n",
    "model_name_or_path = \"./models/upstage/TinySolar-248m-4k\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f589b8e5-839f-4c98-8937-ddb60c324d85",
   "metadata": {
    "height": 183
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM\n",
    "import torch\n",
    "\n",
    "model_name_or_path = \"./models/output/checkpoint-10000\"\n",
    "model2 = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name_or_path,\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.bfloat16,    \n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "86beb9b4-fb8f-4ac1-ab25-f95fe11e7d8c",
   "metadata": {
    "height": 319
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "that everyone could aspire to have a unique sense of talent, all things that apply to and I have great sympathy both with our family and the individuals involved. In my view, society is inappropriately and wrong, especially in this community context. As a result, the right thing to get involved in is the people\n"
     ]
    }
   ],
   "source": [
    "prompt = \"I am an engineer. I love\"\n",
    "\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\").to(model2.device)\n",
    "\n",
    "streamer = TextStreamer(\n",
    "    tokenizer, \n",
    "    skip_prompt=True, \n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "outputs = model2.generate(\n",
    "    **inputs, \n",
    "    streamer=streamer, \n",
    "    use_cache=True, \n",
    "    max_new_tokens=64,     \n",
    "    do_sample=True,\n",
    "    temperature=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd021b5c",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}