gemma_2b_experiments.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:40.372012Z",
     "start_time": "2024-05-01T00:47:39.637729Z"
    }
   },
   "source": [
    "from huggingface_hub import notebook_login\n",
    "from datasets import load_dataset, Dataset"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "id": "50c8dfb9ed9d1242",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:43.105809Z",
     "start_time": "2024-05-01T00:47:40.374245Z"
    }
   },
   "source": [
    "import torch\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, \\\n",
    "    PreTrainedModel\n",
    "from peft import LoraConfig, PeftModel\n",
    "from trl import SFTTrainer\n",
    "from trl.trainer import ConstantLengthDataset\n",
    "\n",
    "# from torchinfo import summary\n"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:43.111631Z",
     "start_time": "2024-05-01T00:47:43.107371Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import os\n",
    "import time"
   ],
   "id": "287861d75264dc8b",
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "d8958b7b6bb1a452",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:43.119679Z",
     "start_time": "2024-05-01T00:47:43.113887Z"
    }
   },
   "source": "# notebook_login()",
   "outputs": [],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "996b8e751778e120",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:44.350591Z",
     "start_time": "2024-05-01T00:47:43.122924Z"
    }
   },
   "source": "eli5: Dataset = load_dataset(\"eli5_category\", split=\"train[:5000]\", trust_remote_code=True)",
   "outputs": [],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "12adddb58b5f39e2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:44.367526Z",
     "start_time": "2024-05-01T00:47:44.352156Z"
    }
   },
   "source": [
    "eli5_train_test = eli5.train_test_split(test_size=0.2)"
   ],
   "outputs": [],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "382af01843cb8b2d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:44.373527Z",
     "start_time": "2024-05-01T00:47:44.369599Z"
    }
   },
   "source": "# eli5_train_test[\"train\"][0]",
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "b6be1ca8c73ff384",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:44.385369Z",
     "start_time": "2024-05-01T00:47:44.375948Z"
    }
   },
   "source": [
    "model_id = \"google/gemma-2b\"\n",
    "four_bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",#todo try 8 bit quantization (only int8 available in bnb) for gemma _2_ b\n",
    "    bnb_4bit_compute_dtype=torch.float16,#todo explore whether it actually increases overall memory usage to change this to float32\n",
    "    bnb_4bit_use_double_quant=True# reminder that this makes computing the total memory used by the frozen weights even more complicated, something about reducing the size of the quantization constants that are used for remembering how to dequantize a given block of quantized weights? by the equivalent of 0.4 bits per parameter, per docs\n",
    ")\n",
    "\n",
    "eight_bnb_config = BitsAndBytesConfig(\n",
    "    load_in_8bit=True\n",
    ")\n",
    "\n",
    "bnb_config = four_bnb_config\n",
    "\n",
    "\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=2,\n",
    "    target_modules=[#\"embed_tokens\", \n",
    "        # \"q_proj\", \"o_proj\", \"k_proj\", \"v_proj\",\n",
    "                    \"gate_proj\", \"up_proj\", \"down_proj\"\n",
    "                    ],# reminder, can use \"all-linear\" (not inside list) for the expansive case https://huggingface.co/docs/peft/developer_guides/lora#qlora-style-training\n",
    "    task_type=\"CAUSAL_LM\",\n",
    "    use_rslora=True\n",
    "    #todo investigate whether it's worth trying Dora, iirc that was said to be especially helpful when lora rank is low\n",
    ")\n",
    "\n",
    "seq_len = 4\n"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.324726Z",
     "start_time": "2024-05-01T00:47:44.387628Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n",
    "eli5_train_test = eli5_train_test.flatten()\n",
    "\n",
    "def fix_data(record):\n",
    "    '''\n",
    "    make dataset usable by TRL (i.e. its classes have a dataset_text_field param, and that column must be string-type,\n",
    "    not list<string> type)\n",
    "    :param record: record where the text column is actually a length-1 list column\n",
    "    :return: record where the text column is straightforwardly a text-type column\n",
    "    '''\n",
    "    record[\"answers.text\"] = record[\"answers.text\"][0]\n",
    "    return record\n",
    "\n",
    "eli5_train_test = eli5_train_test.map(fix_data)"
   ],
   "id": "62dbfd6ce068ba60",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Map:   0%|          | 0/4000 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "843bac9e50b448068060dcdabcba34ba"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "33a33bce8c204c378c8f778a091371c2"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "3eb0992414f9a31f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.924193Z",
     "start_time": "2024-05-01T00:47:45.327190Z"
    }
   },
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "# tokenizer.padding_side = \"right\" #todo based on runtime warning while building SFTTrainer, can try this if problems occur, but by default I think I should respect the default for the Gemma Tokenizer\n",
    "\n",
    "fixed_len_train_dset =  ConstantLengthDataset(tokenizer, eli5_train_test[\"train\"], \"answers.text\", seq_length=seq_len)\n",
    "fixed_len_eval_dset =  ConstantLengthDataset(tokenizer, eli5_train_test[\"test\"], \"answers.text\", seq_length=seq_len)\n"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "id": "8cba9279ee911528",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.931224Z",
     "start_time": "2024-05-01T00:47:45.925410Z"
    }
   },
   "source": "torch.cuda.memory.get_allocator_backend()",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'native'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.938783Z",
     "start_time": "2024-05-01T00:47:45.932469Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# os.environ[\"PYTORCH_NO_CUDA_MEMORY_CACHING\"] = \"1\"\n",
    "#the above completely breaks the visualization (somehow the graph seems to say that tens of gigabytes of ram are being used, but I'm running it on just 8gb vram gpu)"
   ],
   "id": "73a4b229050bb2a3",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.945847Z",
     "start_time": "2024-05-01T00:47:45.940109Z"
    }
   },
   "cell_type": "code",
   "source": [
    "#try messing with PYTORCH_CUDA_ALLOC_CONF to make memory profiles more understandable\n",
    "# os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"backend:cudaMallocAsync\"\n",
    "#?try using cudaMallocAsync instead of pytorch native for backend\n",
    "# couldn't get this to work; I suspect some problem with setting environment variables in a jupyter notebook from windows but the notebook (and the pytorch code) is running in wsl\n",
    "\n",
    "\n",
    "#if cudamallocasync doesn't make things clearer, switch back to pytorch native backend but mess with roundup_power2_divisions option\n",
    "# maybe try [128:32,256:32,512:64,1024:128,2048:256] so hopefuly things get broken down at 4mib granularity\n",
    "# tried this via torch.cuda.memory._set_allocator_settings, but it only changed the 124mib blocks at peak to 126mib blocks, idk why/how\n"
   ],
   "id": "5921948eddaa87fd",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.955394Z",
     "start_time": "2024-05-01T00:47:45.950902Z"
    }
   },
   "cell_type": "code",
   "source": "# torch.cuda.memory._set_allocator_settings(\"roundup_power2_divisions:[128:32,256:32,512:64,1024:128,2048:256]\")\n",
   "id": "f40cf0484a12531",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.964116Z",
     "start_time": "2024-05-01T00:47:45.957669Z"
    }
   },
   "cell_type": "code",
   "source": "torch.cuda.memory.get_allocator_backend()",
   "id": "24f3e27e26ffa65e",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'native'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.972049Z",
     "start_time": "2024-05-01T00:47:45.966227Z"
    }
   },
   "cell_type": "code",
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(\"gpu memory usage after clearing cache is \" + torch.cuda.memory_summary(torch.device(\"cuda:0\")))"
   ],
   "id": "5edebeb092001b65",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gpu memory usage after clearing cache is |===========================================================================|\n",
      "|                  PyTorch CUDA memory summary, device ID 0                 |\n",
      "|---------------------------------------------------------------------------|\n",
      "|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |\n",
      "|===========================================================================|\n",
      "|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Allocated memory      |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from large pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from small pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Active memory         |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from large pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from small pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Requested memory      |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from large pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from small pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|---------------------------------------------------------------------------|\n",
      "| GPU reserved memory   |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from large pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from small pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Non-releasable memory |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from large pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|       from small pool |      0 B   |      0 B   |      0 B   |      0 B   |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Allocations           |       0    |       0    |       0    |       0    |\n",
      "|       from large pool |       0    |       0    |       0    |       0    |\n",
      "|       from small pool |       0    |       0    |       0    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Active allocs         |       0    |       0    |       0    |       0    |\n",
      "|       from large pool |       0    |       0    |       0    |       0    |\n",
      "|       from small pool |       0    |       0    |       0    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| GPU reserved segments |       0    |       0    |       0    |       0    |\n",
      "|       from large pool |       0    |       0    |       0    |       0    |\n",
      "|       from small pool |       0    |       0    |       0    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Non-releasable allocs |       0    |       0    |       0    |       0    |\n",
      "|       from large pool |       0    |       0    |       0    |       0    |\n",
      "|       from small pool |       0    |       0    |       0    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Oversize allocations  |       0    |       0    |       0    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Oversize GPU segments |       0    |       0    |       0    |       0    |\n",
      "|===========================================================================|\n",
      "\n"
     ]
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.980182Z",
     "start_time": "2024-05-01T00:47:45.973877Z"
    }
   },
   "cell_type": "code",
   "source": "torch.cuda.memory._record_memory_history()",
   "id": "db460e6757c1303e",
   "outputs": [],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.987552Z",
     "start_time": "2024-05-01T00:47:45.981756Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# print(\"start loading model at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "# test_model : PreTrainedModel = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map=\"auto\")\n",
    "# print(\"finish loading model at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))"
   ],
   "id": "3a9f74e5e43ce2da",
   "outputs": [],
   "execution_count": 18
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:45.995296Z",
     "start_time": "2024-05-01T00:47:45.988880Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# print(\"printing modules in model\")\n",
    "# print(test_model.modules().__next__())"
   ],
   "id": "1471fd571b720246",
   "outputs": [],
   "execution_count": 19
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:46.002043Z",
     "start_time": "2024-05-01T00:47:45.997057Z"
    }
   },
   "cell_type": "code",
   "source": "# lora_model = PeftModel(test_model, lora_config)",
   "id": "20b53e945e949567",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:46.008321Z",
     "start_time": "2024-05-01T00:47:46.003822Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# print(\"printing modules in lora model\")\n",
    "# print(lora_model.modules().__next__())"
   ],
   "id": "4fd3d0a80fc3428e",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:46.013049Z",
     "start_time": "2024-05-01T00:47:46.009865Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "1329a954452fd5c1",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:46.018752Z",
     "start_time": "2024-05-01T00:47:46.015707Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "adb628ff929d36f4",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:47:46.023010Z",
     "start_time": "2024-05-01T00:47:46.020619Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "3fc13b513d5dbe2d",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:30.809096Z",
     "start_time": "2024-05-01T00:47:46.024524Z"
    }
   },
   "cell_type": "code",
   "source": [
    "with torch.profiler.profile(\n",
    "        with_modules=True,\n",
    "         with_stack=True,\n",
    "        profile_memory=True,\n",
    "        record_shapes=True\n",
    "                            ) as prof:\n",
    "    print(\"start loading model at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "    model : PreTrainedModel = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map=\"auto\")\n",
    "    print(\"finish loading model at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "\n",
    "    print(\"printing modules in model\")\n",
    "    print(model.modules().__next__())\n",
    "    # for module_name in model.modules():# might revisit this if other models besides gemma 2b don't have such a nicely descriptive first model, but unlikely\n",
    "    #     print(module_name)\n",
    "    #print(\"summary of model:\")\n",
    "    #print(summary(model, input_size=(256128, seq_len), device=None, depth=10, col_names=[\"input_size\", \"output_size\", \"num_params\", \"kernel_size\", \"trainable\"]))\n",
    "    # if try torchinfo again, experiment here with input_data=fixed_len_train_dset.dataset or something\n",
    "\n",
    "    print(\"start creating trainer at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "    trainer = SFTTrainer(\n",
    "        model=model,\n",
    "        train_dataset=fixed_len_train_dset,\n",
    "        eval_dataset=fixed_len_eval_dset,\n",
    "        args=TrainingArguments(\n",
    "            per_device_train_batch_size=8,\n",
    "            # gradient_accumulation_steps=4,#todo don't want to touch this until I understand it\n",
    "            warmup_steps=2,\n",
    "            max_steps=10,\n",
    "            learning_rate=2e-4,\n",
    "            # fp16=True,#never turn this on! it adds a 1000mib chunk to the very peak of memory usage in the middle of each epoch!\n",
    "            logging_steps=1,\n",
    "            output_dir=\"ml_outputs\",\n",
    "            optim=\"paged_adamw_32bit\",#can try paged_adamw_8bit in absolute worst case\n",
    "            gradient_checkpointing=True\n",
    "            # Apparently pytorch vram viz really, really does not like quantization; before turning on grad checkpointing, it said\n",
    "            # the frozen/quantized weights were \"unknown\"; now, it says they're \"inputs\" lmao\n",
    "            # Regardless, this does substantially reduce the height of the peaks (by ~450mib in scenario4) and sometimes moves them to the end of the epoch instead of the middle\n",
    "            # \n",
    "        ),\n",
    "        packing=True,#tried disabling this (scenario 12), didn't seem to meaningfully change the shape of the memory visualization over time\n",
    "        # dataset_text_field=\"answers.text\",\n",
    "        peft_config=lora_config,\n",
    "        max_seq_length=seq_len\n",
    "    )\n",
    "    print(\"finish creating trainer at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "\n",
    "    print(\"start training at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "    trainer.train()\n",
    "    print(\"finish training at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n"
   ],
   "id": "5f8baa2108d150f6",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "STAGE:2024-04-30 20:47:46 254892:254892 ActivityProfilerController.cpp:314] Completed Stage: Warm Up\n",
      "Gemma's activation function should be approximate GeLU and not exact GeLU.\n",
      "Changing the activation function to `gelu_pytorch_tanh`.if you want to use the legacy `gelu`, edit the `model.config` to set `hidden_activation=gelu`   instead of `hidden_act`. See https://github.com/huggingface/transformers/pull/29402 for more details.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start loading model at 2024-04-30 20:47:46\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e59fda6709e3456b8d56585d860e387a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finish loading model at 2024-04-30 20:47:50\n",
      "printing modules in model\n",
      "GemmaForCausalLM(\n",
      "  (model): GemmaModel(\n",
      "    (embed_tokens): Embedding(256000, 2048, padding_idx=0)\n",
      "    (layers): ModuleList(\n",
      "      (0-17): 18 x GemmaDecoderLayer(\n",
      "        (self_attn): GemmaSdpaAttention(\n",
      "          (q_proj): Linear4bit(in_features=2048, out_features=2048, bias=False)\n",
      "          (k_proj): Linear4bit(in_features=2048, out_features=256, bias=False)\n",
      "          (v_proj): Linear4bit(in_features=2048, out_features=256, bias=False)\n",
      "          (o_proj): Linear4bit(in_features=2048, out_features=2048, bias=False)\n",
      "          (rotary_emb): GemmaRotaryEmbedding()\n",
      "        )\n",
      "        (mlp): GemmaMLP(\n",
      "          (gate_proj): Linear4bit(in_features=2048, out_features=16384, bias=False)\n",
      "          (up_proj): Linear4bit(in_features=2048, out_features=16384, bias=False)\n",
      "          (down_proj): Linear4bit(in_features=16384, out_features=2048, bias=False)\n",
      "          (act_fn): PytorchGELUTanh()\n",
      "        )\n",
      "        (input_layernorm): GemmaRMSNorm()\n",
      "        (post_attention_layernorm): GemmaRMSNorm()\n",
      "      )\n",
      "    )\n",
      "    (norm): GemmaRMSNorm()\n",
      "  )\n",
      "  (lm_head): Linear(in_features=2048, out_features=256000, bias=False)\n",
      ")\n",
      "start creating trainer at 2024-04-30 20:47:50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/scostesal/.virtualenvs/local-finetuning-calculator/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:318: UserWarning: You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code.\n",
      "  warnings.warn(\n",
      "max_steps is given, it will override any value given in num_train_epochs\n",
      "/home/scostesal/.virtualenvs/local-finetuning-calculator/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:342: UserWarning: You passed `packing=True` to the SFTTrainer, and you are training your model with `max_steps` strategy. The dataset will be iterated until the `max_steps` are reached.\n",
      "  warnings.warn(\n",
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n",
      "/home/scostesal/.virtualenvs/local-finetuning-calculator/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finish creating trainer at 2024-04-30 20:47:50\n",
      "start training at 2024-04-30 20:47:50\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='10' max='10' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [10/10 00:04, 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>1</td>\n",
       "      <td>20.166000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>12.887300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>13.146400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>13.938100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>13.447500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>12.062800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>13.031800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>12.281500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>11.477900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>13.623600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finish training at 2024-04-30 20:47:56\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "STAGE:2024-04-30 20:48:00 254892:254892 ActivityProfilerController.cpp:320] Completed Stage: Collection\n",
      "STAGE:2024-04-30 20:48:00 254892:254892 ActivityProfilerController.cpp:324] Completed Stage: Post Processing\n"
     ]
    }
   ],
   "execution_count": 22
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:30.815943Z",
     "start_time": "2024-05-01T00:48:30.811185Z"
    }
   },
   "cell_type": "code",
   "source": "# torch.cuda.memory.get_allocator_backend()",
   "id": "a2d51fdae65e8ccf",
   "outputs": [],
   "execution_count": 23
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:30.821422Z",
     "start_time": "2024-05-01T00:48:30.818038Z"
    }
   },
   "cell_type": "code",
   "source": "# os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"]",
   "id": "619f48f8df3e8ea7",
   "outputs": [],
   "execution_count": 24
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:30.827898Z",
     "start_time": "2024-05-01T00:48:30.822791Z"
    }
   },
   "cell_type": "code",
   "source": "os.getcwd()",
   "id": "93caa19a6f6fe0ad",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/mnt/c/Users/ssili/PycharmProjects/local-finetuning-calculator'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 25
  },
  {
   "cell_type": "code",
   "id": "894bb28458e3d280",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:34.419858Z",
     "start_time": "2024-05-01T00:48:30.830083Z"
    }
   },
   "source": [
    "result_files_ts_str = time.strftime(\"%Y-%m-%d_%H-%M-%S\", time.localtime())\n",
    "\n",
    "print(\"about to dump memory snapshot at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "torch.cuda.memory._dump_snapshot(f\"gemma2b_mem_snapshot_{result_files_ts_str}.pickle\")\n",
    "print(\"finished dumping memory snapshot at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "about to dump memory snapshot at 2024-04-30 20:48:30\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finished dumping memory snapshot at 2024-04-30 20:48:34\n"
     ]
    }
   ],
   "execution_count": 26
  },
  {
   "cell_type": "code",
   "id": "8cba29dfdce13b50",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:48:34.424918Z",
     "start_time": "2024-05-01T00:48:34.421064Z"
    }
   },
   "source": [
    "# print(\"about to dump memory timeline summary json at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "# prof.export_memory_timeline(f\"gemma2b_mem_timeline_summary_{result_files_ts_str}.json\")\n",
    "# print(\"finished dumping memory timeline summary json and about to dump memory timeline details/raw json at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "# prof.export_memory_timeline(f\"gemma2b_mem_timeline_details_{result_files_ts_str}.raw.json\")\n",
    "# print(\"finished dumping memory timeline details/raw json at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "# print(\"about to dump memory timeline static visualization at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "# prof.export_memory_timeline(f\"gemma2b_mem_timeline_static_viz_{result_files_ts_str}.html\")\n",
    "# print(\"finished dumping memory timeline static visualization about to create nicer (I think interactive) visualization of memory profiling at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))"
   ],
   "outputs": [],
   "execution_count": 27
  },
  {
   "cell_type": "code",
   "id": "2e22d0047400941d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:51:33.394901Z",
     "start_time": "2024-05-01T00:48:34.426165Z"
    }
   },
   "source": [
    "print(\"about to start creating visualization of memory profiling at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))\n",
    "\n",
    "from torch.cuda._memory_viz import profile_plot\n",
    "with open(f'gemma2b_mem_profile_{result_files_ts_str}.html', 'w') as f:\n",
    "    f.write(profile_plot(prof))\n",
    "\n",
    "    #when with_stack, profile_memory, and record_shapes were all True; and dataset size was 5k (leading to 10 hills in memory snapshot- 10 epochs):\n",
    "    # profile_plot() ran for almost 3 hours and then was killed (by the wsl shell? by the python interpreter? not sure)\n",
    "    # when I tried setting with_stack and record_shapes to False to hopefully speed it up, I got error that memory profiling requires them. oops\n",
    "    # the above was all due to laptop only having 16gb of main ram, afaict. profile_plot works fine (if moderately slower than earlier profiler file creation steps) on desktop with 64gb of ram\n",
    "print(\"finished creating visualization of memory profiling at \" + time.strftime(\"%Y-%m-%d %H:%M:%S\", time.localtime()))"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "about to start creating visualization of memory profiling at 2024-04-30 20:48:34\n",
      "finished creating visualization of memory profiling at 2024-04-30 20:51:33\n"
     ]
    }
   ],
   "execution_count": 28
  },
  {
   "cell_type": "code",
   "id": "bc3774ab88ab1128",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:51:33.443566Z",
     "start_time": "2024-05-01T00:51:33.396639Z"
    }
   },
   "source": "torch.cuda.memory._record_memory_history(enabled=None)",
   "outputs": [],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:51:33.449399Z",
     "start_time": "2024-05-01T00:51:33.446335Z"
    }
   },
   "cell_type": "code",
   "source": "",
   "id": "114529668c6d2870",
   "outputs": [],
   "execution_count": 29
  },
  {
   "cell_type": "code",
   "id": "f3cd155151026d7f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-01T00:51:33.455253Z",
     "start_time": "2024-05-01T00:51:33.451954Z"
    }
   },
   "source": "",
   "outputs": [],
   "execution_count": 29
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}