{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "uRLjFRmzHJXS" }, "source": [ "# Assignment 3: Post Training!\n", "
\n", "\n", "\n", "

\n", " \n", "

\n", "\n", "
\n", "\n", "In this assignment you'll learn some of the **motivations and methods behind post-training language models**.\n", "\n", "The assignment is broken down into **five** parts:\n", "\n", " - **1)** Set up a Hugging Face account and experiment with a small pre-trained model. For this assignment, we will be using the Llama-3.2-1B.\n", "\n", " - **2)** Load and explore the GSM8K dataset.\n", "\n", " - **3)** Write an evaluation method to test the performance of your models on the GSM8K dataset.\n", "\n", " - **4)** Fine-tune a pretrained model to use mathematical CoT (Chain of Thought).\n", "\n", " - **5)** Further post train the model to solve grade school math questions.\n", "\n", "After parts 3,4,5 you will evaluate the model's performance on the math questions as the model evolves.\n", "\n", "\n", "**Please note**: The expected training time for part 2 is < 30 mins and ~ 30m-1hr for part 3 on the L4 GPU, which is the one we recommend for this assignment.\n" ] }, { "cell_type": "markdown", "metadata": { "id": "NgMAVbjDJSmc" }, "source": [ "**Background**\n", "\n", "Recall our goal in the course thus far. Suppose we have a sequence of tokens $x = (x_1,x_2,x_3,...,x_n)$ drawn from some true but unknown distribution, call it $P_{data}(x)$.\n", "\n", "We want to build a language model $P_{\\theta}(x)$ parameterized by $\\theta$, such that $P_{\\theta}(x) \\approx P_{data}(x)$. This means: if we sample from our model, we should get text that looks like it came from the *data* distribution. If we evaluate our model on real text, it should assign high probability to it.\n", "$$\n", "\\begin{aligned}\n", "\\theta^{*} &= \\arg\\max_\\theta(\\mathbb{E}_{x \\sim P_{data}}\\big[P_{\\theta}(x)\\big])\\\\\n", "&= \\arg\\max_\\theta(\\mathbb{E}_{x \\sim P_{data}}\\big[\\prod_{i=1}^{T}P_{\\theta}(x_t|x_{\"\n", "2) \"Hawaii is a great location for a winter break trip with your family. [...] \"\n", "3) \"To build a bomb, start by [...]\"\n" ] }, { "cell_type": "markdown", "metadata": { "id": "JSFGfZJKcESU" }, "source": [ "Below, please fill out which prompts you wrote that got you closest to the target sequence and a short 2-3 sentence reflection on why you chose it and what you observed.\n", "\n", "**prompts used:**\n", "1) \"FILL IN\"\n", "2) \"FILL IN\"\n", "3) \"FILL IN\"\n", "\n", "Brief reflection about each one:\n", "\n", "1) \"FILL IN\"\n", "2) \"FILL IN\"\n", "3) \"FILL IN\"\n" ] }, { "cell_type": "markdown", "metadata": { "id": "BEJcIl9F7NwW" }, "source": [ "#Part 2: Explore and load the GSM8K Dataset\n", "\n", "Take a look at the GSM8K dataset. Observe the kinds of questions it asks, the format they are given in, the columns available.\n", "\n", "https://huggingface.co/datasets/openai/gsm8k\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "JAlZ7o3Ml5Qk" }, "outputs": [], "source": [ "from datasets import load_dataset\n", "gsm8k_train = #TODO\n", "gsm8k_test = #TODO\n", "\n", "print(len(gsm8k_train))\n", "print(len(gsm8k_test))" ] }, { "cell_type": "markdown", "metadata": { "id": "MRFddQqeCUtL" }, "source": [ "#Part 3: Testing model performance on GSM8K" ] }, { "cell_type": "markdown", "metadata": { "id": "CID52s4HCkA2" }, "source": [ "Take a couple minutes to think about how you might test the performance a LLM on a GSM8K.\n", " It's harder than it seems!" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "K9xfOWt4tFAV" }, "outputs": [], "source": [ "prefix = \"Your role as an assistant involves thoroughly exploring questions through a systematic long thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution. In the Thought section, detail your reasoning process using the specified format: <|begin_of_thought|> {thought with steps separated with '\\n\\n'} <|end_of_thought|> Each step should include detailed considerations such as analisying questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The solution should remain a logical, accurate, concise expression style and detail necessary step needed to reach the conclusion, formatted as follows: <|begin_of_solution|> {final formatted, precise, and clear solution} <|end_of_solution|> Now, try to solve the following question through the above guidelines:\"\n", "problem = \"Problem: If I have 3 apples, and in sum they cost $4.5, how much does it cost to buy 2 apples?\"\n", "prompt = prefix + problem\n", "\n", "inputs = tokenizer(prompt, return_tensors=\"pt\")\n", "inputs = {k: v.to(model.device) for k, v in inputs.items()}\n", "outputs = model.generate(\n", " **inputs,\n", " max_new_tokens=300,\n", " num_return_sequences=1,\n", " do_sample=True,\n", " top_p=0.8,\n", " temperature=1.0,\n", " pad_token_id=tokenizer.eos_token_id,\n", ")\n", "\n", "generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)\n", "print(generated_text)" ] }, { "cell_type": "markdown", "metadata": { "id": "JcsS_WE_k0uz" }, "source": [ "You might have noticed that the GSM8K dataset has a numerical answer following the #### delimeter in the answer body. We will be trying to use this to parse the responses from Llama.\n", "\n", "Fill in the methods below." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "-G6Puudinkc3" }, "outputs": [], "source": [ "def extract_answer_after_hashes(s):\n", " # split on \"####\", strip spaces, take everything after. Returns string\n", " ## Fill in with your code ##\n", " return\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "lX_7OBgQnvUW" }, "outputs": [], "source": [ "def preprocess_gsm8k(example):\n", " # take a single row from gsm8k and return as the following columns: {\"question\", \"response\", \"numeric_response\"}\n", " ## Fill in with you code ##\n", " return {\"question\": q, \"response\": r, \"numeric_answer\": a}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "2GPfICVjmMJc" }, "outputs": [], "source": [ "gsm8k_train = gsm8k_train.map(preprocess_gsm8k, remove_columns=gsm8k_test.column_names)\n", "gsm8k_test = gsm8k_test.map(preprocess_gsm8k, remove_columns=gsm8k_test.column_names)\n", "# Evaluating the model on hundreds of questions is not computationally trivial.\n", "# Remember, we are literally generating tens to hundreds of tokens one after another for each question\n", "# You can define a smaller sample of the test code if you want slightly faster testing times, though I would still suggest at least 200 questions.\n", "gsm8k_test_sample = gsm8k_test.select(range(800))" ] }, { "cell_type": "markdown", "metadata": { "id": "8k1E1YTItrdG" }, "source": [ "What if the model emits a sequence that doesn't follow the format we expect, but answered correctly? For example, no #### before emitting the correct answer?\n", "\n", "What if the model emits multiple ####s?\n", "\n", "What if the true answer is a decimal, but the model emits a fraction?\n", "\n", "The methods below are designed to handle these types of cases.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "y9I5alSHeMCZ" }, "outputs": [], "source": [ "import re\n", "from decimal import Decimal, InvalidOperation\n", "from fractions import Fraction\n", "\n", "_NUM_RE = re.compile(r\"[-+]?\\d+(?:\\.\\d+)?(?:/[1-9]\\d*)?\")\n", "_ANS_RE = re.compile(r\"####\\s*(\" + _NUM_RE.pattern + r\")\\b\")\n", "\n", "def _to_decimal(s: str):\n", " s = s.strip().replace(\",\", \"\")\n", " if \"/\" in s:\n", " try:\n", " return Decimal(Fraction(s))\n", " except Exception:\n", " pass\n", " try:\n", " return Decimal(s)\n", " except InvalidOperation:\n", " return None\n", "\n", "def normalize_num(s):\n", " return _to_decimal(str(s))\n", "\n", "# explicitly look for \"#### \" anywhere in the text.\n", "_ANS_MARK_RE = re.compile(r\"####\\s*([-+]?\\d+(?:\\.\\d+)?)(?!\\S)\")\n", "\n", "def extract_answer(text: str):\n", " # 1) Preferred: \"#### \"\n", " m = _ANS_RE.search(text)\n", " if m:\n", " return _to_decimal(m.group(1))\n", "\n", " # 2) Try inside explicit solution tags if they exist\n", " m = re.search(r\"<\\|begin_of_solution\\|>(.*?)<\\|end_of_solution\\|>\", text, flags=re.S|re.I)\n", " if m:\n", " nums = _NUM_RE.findall(m.group(1))\n", " if nums:\n", " return _to_decimal(nums[-1])\n", "\n", " # 3) Fallback: last number anywhere\n", " nums = _NUM_RE.findall(text)\n", " if nums:\n", " return _to_decimal(nums[-1])\n", "\n", " return None" ] }, { "cell_type": "markdown", "metadata": { "id": "CA-qJ94DuKIL" }, "source": [ "We can use the helpers to return whether there was a match given the generated text and true answer from GSM8k." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "RZ4KwmnGuPM9" }, "outputs": [], "source": [ "def evaluate_em(gen, gold, abs_tol=Decimal(\"1e-9\"), rel_tol=Decimal(\"1e-9\")):\n", " ## Replace with your code ##\n", " em_flag = 1 if (diff <= abs_tol) or (diff / denom <= rel_tol) else 0\n", " return em_flag, prediction, gold " ] }, { "cell_type": "markdown", "metadata": { "id": "z9hV0RaCwHUZ" }, "source": [ "**Evaluation**\n", "\n", "Now comes the fun part. How should we present the GSM8K question to Llama such that it is fair, elicits the formatting we want, and can be kept the same across each time we evaluate?\n", "\n", "This is up to you, but 3 key pieces of information I might start out with are\n", "\n", "1) role information\n", "2) how should the question be attempted\n", "3) formatting guidelines.\n", "\n", "Since we are using a strict rule based system for numerical evaluation, the formatting guidlines are extra important.\n", "\n", "Here is template you might follow:\n", "\n", "\"Your role... that involves correctly solving math questions in a... When you are ready to give your solution, format as follows. \\n#### \\.\\n Now solve the following problem:\\n\".\n", "\n", "Now define the prefix below." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Zho7j5BVw_yd" }, "outputs": [], "source": [ "PROMPT_GSM8K = (\n", " \"\"\n", " \"When you are ready to give your solution, format as follows. \\n#### .\\n Now solve the following problem:\\n\"\n", " )" ] }, { "cell_type": "markdown", "metadata": { "id": "rIGdQCEavrAB" }, "source": [ "Finally, we will combine the helpers together into our evaluation loop. No coding work is needed here, but take a read to understand how it works.\n", "\n", "One important note is that we limit generations to 300 tokens. Answer the following questions about the token generation limit.\n", "\n", "Why might one set a maximum token generation length?\n", "\n", "\"FILL IN\"\n", "\n", "What are the downsides?\n", "\n", "\"FILL IN\"\n", "\n", "What are the upsides?\n", "\n", "\"FILL IN\"\n", "\n", "What is the cost of generating 1 new token, given n previous tokens?\n", "\n", "\"FILL IN\"\n", "\n", "What is the cost of generating n tokens in sequence, given k previous tokens?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Pk8Fbd6XprW2" }, "outputs": [], "source": [ "import torch, torchsummary\n", "def evaluate_gsm8k(\n", " model_name,\n", " questions,\n", " gold_answers,\n", " batch_size=4,\n", " max_new_tokens=300, # notice we limit the response to 300 tokens.\n", " prefix=None,\n", " print_every=5, # print one detailed sample every N examples\n", "):\n", " tok = AutoTokenizer.from_pretrained(model_name, use_fast=True)\n", " tok.padding_side = \"left\"\n", " if tok.pad_token is None:\n", " tok.pad_token = tok.eos_token\n", "\n", " mdl = AutoModelForCausalLM.from_pretrained(\n", " model_name, device_map=\"cuda\", torch_dtype=torch.bfloat16\n", " )\n", " mdl.eval()\n", "\n", " prefix = prefix\n", "\n", " outs = []\n", " records = []\n", " total_early, total_seen = 0, 0\n", " correct_responses = []\n", "\n", " with torch.inference_mode():\n", " for i in range(0, len(questions), batch_size):\n", " print(f\"Evaluation progress: {min(i + batch_size, len(questions))} / {len(questions)}\")\n", "\n", " batch_q = questions[i:i + batch_size]\n", " prompts = [prefix + q for q in batch_q]\n", "\n", " enc = tok(prompts, return_tensors=\"pt\", padding=True, truncation=True)\n", " enc = {k: v.to(mdl.device, non_blocking=True) for k, v in enc.items()}\n", " enc_len = enc[\"input_ids\"].shape[1]\n", "\n", " # notice we are doing greedy decoding here. This is typical in an evaluation setting, where we don't want any source of randomness to affect the accuracy.\n", " gen = mdl.generate(\n", " **enc,\n", " max_new_tokens=max_new_tokens,\n", " do_sample=False,\n", " temperature=0.0,\n", " top_p=1.0,\n", " use_cache=True,\n", " pad_token_id=tok.eos_token_id,\n", " eos_token_id=tok.eos_token_id,\n", " return_dict_in_generate=False,\n", " output_scores=False,\n", " )\n", "\n", " new_tokens = gen[:, enc_len:]\n", " texts = tok.batch_decode(new_tokens, skip_special_tokens=True)\n", "\n", " for j, (q, t, gold) in enumerate(zip(batch_q, texts, gold_answers[i:i + batch_size])):\n", " gen_len = int(new_tokens[j].shape[0])\n", " early = gen_len < max_new_tokens\n", " hit_marker = (\"####\" in t) or (\"<|end_of_solution|>\" in t)\n", "\n", " cleaned = t.rstrip()\n", " outs.append(cleaned)\n", "\n", " em_flag, pred_num, gold_num = evaluate_em(t, gold)\n", "\n", " if em_flag == 1:\n", " correct_responses.append({\n", " \"idx\": idx,\n", " \"question\": q,\n", " \"generation\": t,\n", " \"pred\": pred_num,\n", " \"gold\": gold_num\n", " })\n", "\n", " rec = {\n", " \"idx\": idx,\n", " \"question\": q,\n", " \"generation\": t, # raw generation (with markers if any)\n", " \"cleaned\": cleaned, # truncated at markers\n", " \"pred_num\": str(pred_num) if pred_num is not None else None,\n", " \"gold\": gold,\n", " \"gold_num\": str(gold_num) if gold_num is not None else None,\n", " \"em\": em_flag,\n", " \"gen_len\": gen_len,\n", " \"max_new_tokens\": int(max_new_tokens),\n", " \"early_stop\": early,\n", " \"hit_marker\": hit_marker,\n", " }\n", " records.append(rec)\n", "\n", " total_seen += 1\n", " total_early += int(early)\n", "\n", " if (i + j) % max(1, print_every) == 0:\n", " print(f\"[{rec['idx']}] early={rec['early_stop']} len={rec['gen_len']} EM={rec['em']}\")\n", " print(f\"Q: {q}\")\n", " print(f\"GEN: {t}\")\n", " print(f\"PRED={rec['pred_num']} | GOLD={rec['gold_num']}\")\n", " print(\"-\" * 80)\n", "\n", " # free per-batch tensors\n", " del gen, new_tokens, texts, enc\n", "\n", " # Aggregate stats\n", " early_rate = total_early / max(1, total_seen)\n", " gen_lengths = [r[\"gen_len\"] for r in records]\n", " mean_len = sum(gen_lengths) / max(1, len(gen_lengths))\n", " median_len = sorted(gen_lengths)[len(gen_lengths) // 2] if gen_lengths else 0\n", "\n", " accuracy = sum(r[\"em\"] for r in records) / max(1, len(records))\n", "\n", " print(\"\\n=== Evaluation summary ===\")\n", " print(f\"Accuracy (EM): {accuracy:.4f}\")\n", " print(f\"Early-stop rate: {early_rate*100:.1f}%\")\n", " print(f\"Gen length: mean={mean_len:.1f}, median={median_len}, cap={max_new_tokens}\")\n", " print(\"==========================\\n\")\n", "\n", " if correct_responses:\n", " print(f\"\\n=== Correct predictions ({len(correct_responses)}) ===\")\n", " for s in correct_responses:\n", " print(f\"[{s['idx']}] PRED={s['pred']} | GOLD={s['gold']}\")\n", " print(\"Q:\", s[\"question\"])\n", " print(\"GEN:\", s[\"generation\"])\n", " print(\"-\" * 80)\n", " return accuracy, outs, records" ] }, { "cell_type": "markdown", "metadata": { "id": "4hDZ6Z0n5HjJ" }, "source": [ "**Run the eval!**" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tAblsqXxp_NN" }, "outputs": [], "source": [ "accuracy, outs, records = evaluate_gsm8k(\"meta-llama/Llama-3.2-1B\", gsm8k_test_sample[\"question\"], gsm8k_test_sample[\"numeric_answer\"], 16, prefix)\n", "\n", "print(f\"Accuracy: {accuracy}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "AawmMDUzmgWS" }, "source": [ "Please record your accuracy here:\n", "\n", "\\\n", "\n", "And 5 examples here:\n", "\n", "\\\n", "\n", "What do you notice about the behavior of the model when faced these questions?\n", "\n", "\\<3-5 sentence reflection\\>\n" ] }, { "cell_type": "markdown", "metadata": { "id": "mntxPOIQ04tN" }, "source": [ "#Part 4: Supervised Fine Tuning" ] }, { "cell_type": "markdown", "metadata": { "id": "uzDL1O0dQeY-" }, "source": [ "SFT is the most basic form of post-training. Apart from one detail, it is exactly the same as the pretraining step.\n", "\n", "As we know, pretraining minimizes the average negative log-likelihood of each true token given all the previous ones (see beginning of assignment if you're confused).\n", "\n", "In SFT, the loss is still a negative log-likelihood, but over condition-response pairs instead of arbitrary documents. Precisely, it is:\n", "$$\n", "\\mathcal{L}_{\\text{SFT}}(\\theta)\n", "= -\\mathbb{E}_{(x, y_{1:T}) \\sim P_{\\text{SFT}}}\n", "\\sum_{t=1}^{T} \\log(P_\\theta(y_t \\mid x, y_{
\n", "\n", "Here's an intutive example.\n", "\n", "Let's say we want our model to be as educational as possible.\n", "\n", "**Pretrained model**\n", "\n", " - $x$ = \"You are an assistant designed to be as educational as possible. What is 3+4? \"\n", "\n", " - Potential response $y$ = \"7. What is 9+2? 11. What is 1+1? 2. These are all basic arithmetic questions. For more questions like these, visit my blog. \\\"\n", "\n", "**Fine Tuned model**\n", "\n", " - $x$ = \"You are an assistant designed to be as educational as possible. What is 3+4?\"\n", "\n", " - Potential response $y'$ = \"3+4=7 because addition means combining quantities, and counting 3 forward from 4 lands you on 7. \\\"\n", "\n", "
\n", "Each training example is usually something like\n", "\n", "$data_i$ = \\, \\\n", "\n", "During training, we feed the entire sequence to the model so it sees both the instruction and the target.\n", "\n", "But when computing the loss, we only count the tokens in the answer, not the tokens in the prompt. (Since we aren't asking the model to learn how to predict the instruction).\n", "\n", "That's implemented using a mask: a binary vector (same length as the sequence) with 1s for target tokens and 0s for everything else.\n", "During the loss computation, each token's log-probability is multiplied by its mask value, so gradients only flow through the desired answer part.\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "jG-d7xNuh_xb" }, "source": [ "We will be using the GSM8K dataset to fine tune Llama 3.2 1b.\n", "\n", "Formally, we are trying to increase $P(\\textbf{gsm8k right answer} | \\textbf{the prefix you defined}, \\textbf{gsm8k question})$\n" ] }, { "cell_type": "markdown", "metadata": { "id": "-2RoOScA3h30" }, "source": [ "Let's first pretokenize all our training data. The benefit of doing this is that it avoids re-tokenizing the same text repeatedly during training, which saves time and ensures that all examples share a consistent tokenization scheme. It also lets us inspect and cache tokenized sequences in advance, which is useful for debugging. Most importantly, pretokenization significantly improves efficiency by reducing on-the-fly preprocessing overhead during each training step." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "BiAIhDQI2QKY" }, "outputs": [], "source": [ "import os\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n", "#TODO START: Define tokenizer with right padding\n", "\n", "\n", "#TODO END#\n", "\n", "if tok.pad_token is None:\n", " tok.pad_token = tok.eos_token\n", "tok.truncation_side = \"left\"\n", "\n", "MAX_LEN = 400 # I chose this because at ~3 characters a token, this admits a total sequence of 1200 characters.\n", " #The distribution of character lengths for the summed question, answer columns is mostly under 1200 characters.\n", "BUDGET = MAX_LEN - 1\n", "\n", "#TODO START: Tokenize the system prompt/prefix\n", "SYS_IDS =\n", "#TODO END#\n", "\n", "def tokenize_batch(batch, include_answer=True):\n", " qs = [q.rstrip() for q in batch[\"question\"]]\n", "\n", " #TODO START: Tokenize qs without adding special tokens or padding\n", " enc_q =\n", " #TODO END#\n", " has_response = \"response\" in batch and include_answer\n", "\n", " if has_response:\n", " ans = [a.rstrip() for a in batch[\"response\"]]\n", " #TODO START: Tokenize ans without adding special tokens or padding\n", " enc_a =\n", " #END TODO#\n", " else:\n", " enc_a = {\"input_ids\": [[] for _ in qs]}\n", "\n", " gold_answers = [(n or \"\").rstrip() for n in batch.get(\"numeric_answer\", [\"\"] * len(qs))]\n", " input_ids_list, prompt_len_list, kept_gold = [], [], []\n", "\n", " for i, (q_ids, a_ids) in enumerate(zip(enc_q[\"input_ids\"], enc_a[\"input_ids\"])):\n", " #TODO START: Define ids\n", " ids =\n", " #TODO END#\n", " if len(ids) > MAX_LEN:\n", " #TODO START: define behavior if len(ids) > MAX_LEN\n", "\n", " #TODO END#\n", " input_ids_list.append(ids)\n", " prompt_len_list.append(len(SYS_IDS) + len(q_ids))\n", " kept_gold.append(gold_answers[i])\n", "\n", " return {\"input_ids\": input_ids_list, \"prompt_len\": prompt_len_list, \"gold_answer\": kept_gold}\n", "\n", "\n", "train_tok = gsm8k_train.map(\n", " tokenize_batch,\n", " batched=True,\n", " batch_size=1024,\n", " num_proc=2,\n", " remove_columns=gsm8k_train.column_names,\n", " writer_batch_size=1024,\n", " desc=\"Tokenizing train set\",\n", " fn_kwargs={\"include_answer\": True}\n", ")\n", "\n", "val_tok = gsm8k_test.map(\n", " tokenize_batch,\n", " batched=True,\n", " batch_size=1024,\n", " num_proc=2,\n", " remove_columns=gsm8k_test.column_names,\n", " writer_batch_size=1024,\n", " desc=\"Tokenizing val set\",\n", " fn_kwargs={\"include_answer\": True}\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "Xt5pE5oO-vHw" }, "source": [ "Additionally, we do have to make sure Llama 3.2 1b can handle any sequence we're interested in. Let's check the maximum context length for Llama 3.2 1b, and compare against the token lengths in the tokenized examples." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "2r5HZ56h-u40" }, "outputs": [], "source": [ "from transformers import AutoModelForCausalLM\n", "mdl = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-3.2-1B\", device_map=\"auto\", torch_dtype=\"auto\")\n", "print(\"max_position_embeddings:\", getattr(mdl.config, \"max_position_embeddings\", None))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "u7Hf3weF37G9" }, "outputs": [], "source": [ "count = 0\n", "total = 0\n", "for i in range(0,1000,10):\n", " if (len(train_tok[i]['input_ids']) >= 400):\n", " count+=1\n", " total +=1\n", "print(count/total)" ] }, { "cell_type": "markdown", "metadata": { "id": "KIoaWPnRB1LI" }, "source": [ "The PromptMaskedCollator is responsible for taking a list of examples (each containing tokenized input IDs, attention masks, and the length of the prompt) and turning them into a single batch tensor that the model can train on. Importantly, the collator is responsible for masking the log probabilities of the prompt tokens. Fill in the masking logic." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "x4lg4bMJ6YWD" }, "outputs": [], "source": [ "class PromptMaskedCollator:\n", " def __init__(self, tokenizer, pad_to_multiple_of=8):\n", " self.tok = tokenizer\n", " self.pad_to_multiple_of = pad_to_multiple_of\n", "\n", " def __call__(self, features):\n", " prompt_len = torch.tensor([f[\"prompt_len\"] for f in features], dtype=torch.long)\n", "\n", " feats_wo_plen = [{k: v for k, v in f.items() if k != \"prompt_len\"} for f in features]\n", "\n", " batch = self.tok.pad(\n", " feats_wo_plen,\n", " padding=True,\n", " return_tensors=\"pt\",\n", " pad_to_multiple_of=self.pad_to_multiple_of,\n", " )\n", "\n", " input_ids = batch[\"input_ids\"]\n", " attn = batch[\"attention_mask\"]\n", "\n", " T = input_ids.size(1)\n", " ar = torch.arange(T, device=input_ids.device).unsqueeze(0)\n", " plen = prompt_len.unsqueeze(1).to(device=input_ids.device)\n", "\n", " # TODO:\n", " #labels =\n", " #labels[] =\n", " #labels[] =\n", " #END TODO#\n", " batch[\"labels\"] = labels\n", " return batch\n", "\n", "\n", "collator = PromptMaskedCollator(tok)" ] }, { "cell_type": "markdown", "metadata": { "id": "VsP25SUD64aj" }, "source": [ "In order to do SFT fast, we will fine tune our model using a method called LoRA, which has been covered in class. Implementing it isn't trivial in raw PyTorch, so instead of we'll be using a library called [peft](https://huggingface.co/docs/peft/en/index). Take a look at the LoRA documentation and fill in the code below.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "wQ9DcXVkV1Km" }, "outputs": [], "source": [ "import torch\n", "from transformers import (\n", " AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, BitsAndBytesConfig\n", ")\n", "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " \"meta-llama/Llama-3.2-1B\",\n", " device_map=\"auto\",\n", " torch_dtype=torch.float16,\n", " attn_implementation=\"sdpa\",\n", ")\n", "\n", "model.config.use_cache = False\n", "\n", "# TODO START: Define Lora Config and define the model using the config. I would suggest starting off with r=8, lora_alpha=16, lora_dropout=0.05\n", "# lora_config =\n", "# model =\n", "# TODO END#\n", "model.print_trainable_parameters()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we intialize the Trainer and TrainingArguments for training. Please read through the TrainingArguments, and for each one write 1-2 sentences describing its functionality.\n", "\n", "Argument functionalies [1-21]:\n", "\\" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "args = TrainingArguments(\n", " output_dir=\"./gsm8ksft_1b_lora\",\n", " num_train_epochs=2,\n", " per_device_train_batch_size=4,\n", " per_device_eval_batch_size=4,\n", " gradient_accumulation_steps=4,\n", " learning_rate=2e-4,\n", " lr_scheduler_type=\"cosine\",\n", " warmup_ratio=0.03,\n", " logging_steps=1,\n", " eval_strategy=\"steps\",\n", " eval_steps=25,\n", " save_steps=250,\n", " save_total_limit=2,\n", " bf16=False,\n", " fp16=True,\n", " gradient_checkpointing=True,\n", " gradient_checkpointing_kwargs={\"use_reentrant\": False},\n", " optim=\"adamw_torch\",\n", " report_to=\"none\",\n", " remove_unused_columns=False,\n", " group_by_length=True,\n", ")\n", "\n", "#subsample\n", "val_tok_sample = val_tok.shuffle(seed=42).select(range(100))\n", "\n", "trainer = Trainer(\n", " model=model,\n", " args=args,\n", " train_dataset=train_tok,\n", " eval_dataset=val_tok_sample,\n", " data_collator=collator,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "IvBFdKIWdy7r" }, "outputs": [], "source": [ "#clear gpu memory without restarting runtime.\n", "import gc, torch\n", "for name in (\"trainer\",\"model\",\"optim\",\"scheduler\"):\n", " if name in globals(): del globals()[name]\n", "gc.collect()\n", "torch.cuda.empty_cache()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "X5nwatVnXBy-" }, "outputs": [], "source": [ "trainer.train()\n", "trainer.save_model()\n", "model.push_to_hub(\"/gsm8ksft-1b-lora\")\n", "tok.push_to_hub(\"hf username>/gsm8ksft-1b-lora\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Uw_KfodiFobq" }, "outputs": [], "source": [ "from peft import AutoPeftModelForCausalLM\n", "tok = AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B\", use_fast=True)\n", "\n", "# load the trained adapter (it knows the base model from its config)\n", "peft_model = AutoPeftModelForCausalLM.from_pretrained(\n", " \"/gsm8ksft-1b-lora\", torch_dtype=\"auto\", device_map=\"auto\"\n", ")\n", "\n", "# merge LoRA weights into the base weights and drop PEFT wrappers\n", "merged = peft_model.merge_and_unload()\n", "\n", "# save a standard HF model folder\n", "merged.save_pretrained(\"./gsm8k_1b_lora_merged\", safe_serialization=True)\n", "tok.save_pretrained(\"./gsm8k_1b_lora_merged\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mbmrQblICoa3" }, "outputs": [], "source": [ "accuracy, outs, records = evaluate_gsm8k(\"./gsm8k_1b_lora_merged\", gsm8k_test_sample[\"question\"], gsm8k_test_sample[\"numeric_answer\"], 16)\n", "print(f\"Accuracy: {accuracy}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "PcqHve48m836" }, "source": [ "Please record your accuracy here:\n", "\n", "\\\n", "\n", "And 5 examples here:\n", "\n", "\\\n", "\n", "Do you notice anything different about the model behavior?\n", "\n", "\\<3-5 sentence reflection\\>" ] }, { "cell_type": "markdown", "metadata": { "id": "FEmsMwiZaFdp" }, "source": [ "#Part 5: Reinforcement Learning with REINFORCE\n", "\n", "In supervised learning, we train a model to minimize a loss function comparing its predictions to ground-truth labels. However, in many problems, especially when the correct output is not uniquely defined or only indirectly measurable (e.g. dialogue helpfulness, game score, or text quality), we only know how good an output is.\n", "This setting leads naturally to reinforcement learning (RL).\n", "\n", "In reinforcement learning, there are no fixed “correct” labels. Instead, the model (called an agent) learns by interacting with an environment and receiving rewards that measure how good its actions were.\n", "\n", "At each step $t$:\n", "1) The agent observes the state $s_t$ of the environment\n", "2) It samples an action $a_t$ ~ $\\pi_\\theta(a_t|s_t)$ from its policy.\n", "3) The environment transitions to a new state $s_{t+1}$ and emits a reward $r_t \\in \\mathbb{R}$\n", "\n", "This process continues for **T** steps, and we call the the entire process a trajectory\n", "$\\tau = (s_1,a_1,r_1, s_2,a_2,r_2, ..., s_T, a_T, r_T)$\n", "\n", "The total reward from $\\tau$ is called the return:\n", "\n", "$R(\\tau) = \\sum_{t=1}^{T}\\gamma^{t-1}r_t$\n", "\n", "The discount factor models the intuition that recieving a reward earlier is of more utility than recieving a reward later.\n", "\n", "The objective of RL is to find policy parameters $\\theta$ that maximize expected return:\n", "\n", "$\\mathbb{J}(\\theta) = \\mathbb{E}_{\\tau \\sim{} \\pi_\\theta}[R(\\tau)]$\n", "\n", "To put it simply, $\\mathbb{J(\\theta)}$ is the average performance of the policy $\\theta$ where the source of randomness is from **a)** sampling an action $a_t$ and/or **b)** an environment that changes independently.\n", "\n", "The key here is that while we we're able to differentiate the loss w.r.t our model parameters in supervised learning, we aren't able to differentiate the expected reward w.r.t to our model parameters - since our objective is maximizing $\\mathbb{J}(\\theta) = \\mathbb{E}_{\\tau \\sim{} \\pi_\\theta}[R(\\tau)]$, and $R(\\tau)$ comes from the environment (From the perspective of the parameters, $R(\\tau)$ is a black box that outputs a scalar signal after we take a sequence of actions.).\n", "\n", "So we can't take gradients through the reward function,\n", "\n", "**BUT** we can take gradients through the probability of sampling trajectories that lead to reward (with the assumption that the reward function stays fixed).\n", "\n", "This leads to the key idea behind policy gradient methods like REINFORCE:\n", "\n", "$J(\\theta) = \\mathbb{E}_{\\tau \\sim \\pi_\\theta} [R(\\tau)]$\n", "\n", "$J(\\theta) = \\sum_{\\tau} P(\\tau; \\theta) \\, R(\\tau)$\n", "\n", "$\\nabla_\\theta J(\\theta)$\n", "$= \\sum_{\\tau} \\nabla_\\theta P(\\tau; \\theta) \\, R(\\tau)$\n", "\n", "$\\nabla_\\theta P(\\tau; \\theta)$\n", "$= P(\\tau; \\theta) \\, \\nabla_\\theta \\log P(\\tau; \\theta)$\n", "\n", "$\\nabla_\\theta J(\\theta)$\n", "$= \\sum_{\\tau} P(\\tau; \\theta) \\, \\nabla_\\theta \\log P(\\tau; \\theta) \\, R(\\tau)$\n", "\n", "$\\nabla_\\theta J(\\theta)$\n", "$= \\mathbb{E}_{\\tau \\sim \\pi_\\theta}$\n", "$\\big[ R(\\tau) \\, \\nabla_\\theta \\log P(\\tau; \\theta) \\big]$\n", "\n", "$P(\\tau; \\theta)$\n", "$= p(s_1)$\n", "$\\prod_{t=1}^{T} \\pi_\\theta(a_t | s_t) \\, p(s_{t+1} | s_t, a_t)$\n", "\n", "$\\nabla_\\theta \\log P(\\tau; \\theta)$\n", "$= \\sum_{t=1}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t | s_t)$\n", "\n", "$\\nabla_\\theta J(\\theta)$\n", "$= \\mathbb{E}_{\\tau \\sim \\pi_\\theta}$\n", "$\\left[R(\\tau)\\sum_{t=1}^{T}\\nabla_\\theta \\log \\pi_\\theta(a_t | s_t)\\right]$\n", "\n", "\n", "Reframing this back into NLP, a trajectory $\\tau = (o_1,o_2,...o_T)$, and the state $s_t$ is very simply just $o_{< t}$\n", "\n", "For GSM8K, we only have a \"reward\" at the end of the trajectory (when we have emitted \\ or hit the max generation count). The simple assumption in REINFORCE is that each token shares equal responsibility, $r_t = \\frac{R(\\tau)}{T}$.\n", "\n", "Awesome, let's start by write a reward function to assign rewards (1 if correct, 0 if incorrect) to trajectories. (This is very similar to your evaluate_em method!).\n", "Additionally, the function should return the average reward per batch.\n", "\n", "When we do training, we will compute $R_{\\tau_i} = R_{\\tau_i} - b$.\n", "\n", "This doesn't change the direction of the expected gradient, but helps training be more stable. If every batch has both correct and incorrect answers, we want to nudge up the probabilities of correct ones and down the incorrect ones relative to the batch average." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "KnNloP0laFAA" }, "outputs": [], "source": [ "def reward_numeric(pred_texts, gold_answers, abs_tol=Decimal(\"1e-9\"), rel_tol=Decimal(\"1e-9\")):\n", " \"\"\"\n", " Compute per-example numeric rewards for REINFORCE on GSM8K-style outputs. Returns a list of integer rewards corresponding to each GSM8k question, and average reward in the minibatch\n", " \"\"\"\n", " n = min(len(pred_texts), len(gold_answers))\n", " if n == 0:\n", " return [], 0.0\n", "\n", " #TODO START: intialize rewards#\n", " #TOOD END#\n", " for gen, gold in zip(pred_texts[:n], gold_answers[:n]):\n", " pred = extract_answer(gen)\n", " gold_norm = normalize_num(gold)\n", "\n", " if pred is None or gold_norm is None:\n", " #TODO START:\n", " #TOOD END#\n", "\n", " if pred == gold_norm:\n", " #TODO START:\n", " #TOOD END#\n", "\n", " diff = abs(pred - gold_norm)\n", " denom = max(Decimal(1), abs(gold_norm))\n", " if diff <= abs_tol or diff / denom <= rel_tol:\n", " #TODO START:\n", " #TOOD END#\n", " else:\n", " #TODO START:\n", " #TOOD END#\n", "\n", " #TODO START:\n", " baseline =\n", " #TOOD END#\n", " return rewards, baseline\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "XfpYV85ueAfy" }, "outputs": [], "source": [ "class PromptOnlyCollator:\n", " def __init__(self, tokenizer, pad_to_multiple_of=8):\n", " self.tok = tokenizer\n", " self.pad_to_multiple_of = pad_to_multiple_of\n", " def __call__(self, features):\n", " batch = self.tok.pad(\n", " {\"input_ids\": [f[\"input_ids\"] for f in features]},\n", " padding=True, return_tensors=\"pt\",\n", " pad_to_multiple_of=self.pad_to_multiple_of\n", " )\n", " batch[\"gold_answer\"] = [f[\"gold_answer\"] for f in features]\n", " return batch" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "GGEfrFM1hQ5L" }, "outputs": [], "source": [ "gsm8k_train = load_dataset(\"openai/gsm8k\", 'main', split = \"train\")\n", "gsm8k_test = load_dataset(\"openai/gsm8k\",'main', split = \"test\")\n", "\n", "gsm8k_train = gsm8k_train.map(preprocess_gsm8k, remove_columns=gsm8k_test.column_names)\n", "gsm8k_test = gsm8k_test.map(preprocess_gsm8k, remove_columns=gsm8k_test.column_names)\n", "\n", "gsm8k_test_sample = gsm8k_test.select(range(200))\n", "\n", "\n", "tok = AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B\", use_fast=True)\n", "tok.padding_side = \"left\"\n", "if tok.pad_token is None:\n", " tok.pad_token = tok.eos_token\n", "import os\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n", "\n", "tok.truncation_side = \"left\"\n", "\n", "MAX_LEN = 400\n", "BUDGET = MAX_LEN - 1\n", "\n", "SYS_IDS = tok(PROMPT_GSM8K, add_special_tokens=False)[\"input_ids\"]\n", "\n", "train_tok_rl = gsm8k_train.map(\n", " tokenize_batch,\n", " batched=True,\n", " batch_size=1024,\n", " num_proc=2,\n", " remove_columns=gsm8k_train.column_names,\n", " writer_batch_size=1024,\n", " desc=\"Tokenizing train set\",\n", " fn_kwargs={\"include_answer\": False}\n", ")\n", "\n", "val_tok_rl = gsm8k_test.map(\n", " tokenize_batch,\n", " batched=True,\n", " batch_size=1024,\n", " num_proc=2,\n", " remove_columns=gsm8k_test.column_names,\n", " writer_batch_size=1024,\n", " desc=\"Tokenizing val set\",\n", " fn_kwargs={\"include_answer\": False}\n", ")" ] }, { "cell_type": "markdown", "metadata": { "id": "bpW5Q6fBk34b" }, "source": [ "Here is the class we will define for the RL training. It extends the Pytorch Trainer.\n", "There are **4** TODOs." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tD7cM-GtvCaq" }, "outputs": [], "source": [ "import torch\n", "from transformers import (\n", " AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, BitsAndBytesConfig\n", ")\n", "class REINFORCETrainer(Trainer):\n", " def __init__(self, *args, gen_kwargs=None, ref_model=None, kl_beta=0.0, **kwargs):\n", " super().__init__(*args, **kwargs)\n", " self.gen_kwargs = gen_kwargs or dict(\n", " max_new_tokens=300,\n", " do_sample=True, top_p=0.9, temperature=0.7,\n", " pad_token_id=self.tokenizer.pad_token_id,\n", " eos_token_id=self.tokenizer.eos_token_id,\n", " return_dict_in_generate=True, output_scores=False, # scores not needed now\n", " )\n", " self.ref_model = ref_model\n", " self.kl_beta = kl_beta\n", "\n", " @torch.no_grad()\n", " def _decode(self, seqs):\n", " return self.tokenizer.batch_decode(seqs, skip_special_tokens=True)\n", "\n", " def compute_loss(self, model, inputs, return_outputs=False, **kwargs):\n", " # Inputs from collator: input_ids (prompt only), attention_mask, gold_answer (list[str])\n", " input_ids = inputs[\"input_ids\"]\n", " attention_mask = inputs[\"attention_mask\"]\n", " B = input_ids.size(0)\n", " eos_id = self.tokenizer.eos_token_id\n", " pad_id = self.tokenizer.pad_token_id\n", "\n", " # On-policy sampling\n", " with torch.no_grad():\n", " if not hasattr(self, \"_checked_pad\"):\n", " right_zero = (attention_mask[:, -1] == 0).any().item()\n", " assert not right_zero, \"Right padding slipped in\"\n", " self._checked_pad = True\n", "\n", " gen = model.generate(input_ids=input_ids,\n", " attention_mask=attention_mask,\n", " **self.gen_kwargs)\n", " full_ids = gen.sequences # [B, T_full]\n", " full_attn = (full_ids != pad_id).long() # [B, T_full]\n", " prompt_lens = attention_mask.sum(dim=1) # [B]\n", "\n", " T_full = full_ids.size(1)\n", " # forward pass (with grad) over the whole sampled sequence\n", " out = model(input_ids=full_ids[:, :-1], attention_mask=full_attn[:, :-1])\n", " logits = out.logits.float()\n", " # Sanitize the fp32 logits *before* log_softmax.\n", " # This fixes the Inf from the unstable LoRA weights.\n", " logits = torch.nan_to_num(\n", " logits,\n", " nan=0.0,\n", " posinf=torch.finfo(torch.float32).max,\n", " neginf=torch.finfo(torch.float32).min\n", " )\n", " logprobs = torch.log_softmax(logits, dim=-1)\n", "\n", "\n", " ar = torch.arange(T_full, device=full_ids.device).unsqueeze(0).expand(B, -1) # [B,T]\n", "\n", " # ===== TODO =====\n", " # TASK 1: Build a mask of which positions to score for each sample.\n", " # Requirements:\n", " # - score only positions >= prompt_len\n", " # - stop at first EOS (included)\n", " # - ignore padding\n", " # Output shape: valid_mask: [B, T_full]\n", " #\n", " # Hints:\n", " # arange = torch.arange(T_full, device=full_ids.device)[None, :].expand(B, -1)\n", " # gen_region = arange >= prompt_lens[:, None]\n", " # then cut off after first eos\n", " #\n", " # TODO: implement valid_mask\n", " # valid_mask = ...\n", "\n", " # TASK 2: For each sample i, gather log p(sampled_token_i_t) at valid positions.\n", " # Remember: logits[:, t-1, :] predicts token at t (shift by one).\n", "\n", " # TASK 3: Sum the logprobs for each sample i over the valid positions only.\n", "\n", " # TASK 4: Compute the advantage and normalize it and compute the REINFORCE loss.\n", "\n", " logprob_sums = []\n", " for i in range(B):\n", " # TODO: get valid positions for this sample\n", " # pos = ...\n", " # TODO: skip position 0 (no prediction for t=0)\n", " # pos = pos[pos > 0]\n", " if #TODO fill in condition:\n", " safe_zero = (model.get_input_embeddings().weight[0, 0] * 0.0).to(dtype=logits.dtype)\n", " logprob_sums.append(safe_zero)\n", " continue\n", " # gather next-token logprobs for the actually sampled tokens\n", " tok_ids = full_ids[i, pos] # [Ti]\n", " step_logits = logprobs[i, pos - 1, :] # [Ti, V]\n", " step_logps = step_logits.gather(-1, tok_ids.view(-1, 1)).squeeze(-1) # [Ti]\n", " step_logps = torch.where(\n", " torch.isfinite(step_logps),\n", " step_logps,\n", " torch.zeros_like(step_logps)\n", " )\n", " step_logps = step_logps.clamp(min=-20.0, max=0.0)\n", "\n", " Ti = pos.numel()\n", "\n", " # TODO: add logprobs to logprob_sums\n", "\n", " logprob_sums = torch.stack(logprob_sums, dim=0) # [B]\n", "\n", " assert logprob_sums.requires_grad, \"logprob_sums lost grad (masking/shift produced all-empty rows)\"\n", "\n", " # Rewards + batch baseline\n", " with torch.no_grad():\n", " texts = self._decode(full_ids)\n", " rewards_list, b = reward_numeric(texts, inputs[\"gold_answer\"])\n", " rewards = torch.tensor(rewards_list, dtype=logprob_sums.dtype, device=logprob_sums.device) # [B]\n", " baseline = torch.tensor(b, dtype=logprob_sums.dtype, device=logprob_sums.device)\n", "\n", " em_batch = rewards.float().mean().item()\n", " self.log({\"em_batch\": em_batch})\n", "\n", " # TODO: Compute the advantage and normalize it.\n", " # 1. advantages = rewards - baseline\n", " # 2. optional: normalize to unit variance\n", " # 3. use advantage to weight logprob_sums in the REINFORCE loss\n", "\n", " # advantages = ...\n", " # std = advantages.std()\n", " # if torch.isfinite(std) and std > 0:\n", " # advantages = ...\n", " # loss = ...\n", "\n", " # ===== TODO END =====\n", " # KL regularization with a frozen ref model could be added here.\n", " kl_loss = torch.tensor(0.0, device=loss.device)\n", " if self.ref_model is not None and self.kl_beta > 0:\n", " with torch.no_grad():\n", " # Third forward pass (no grad) with the REFERENCE model\n", " ref_out = self.ref_model(input_ids=full_ids[:, :-1], attention_mask=full_attn[:, :-1])\n", " ref_logprobs = torch.log_softmax(ref_out.logits.float(), dim=-1)\n", "\n", " # Mask for generated tokens (shifted)\n", " mask = valid[:, 1:].contiguous()\n", " mask_float = mask.float()\n", "\n", " target_ids = full_ids[:, 1:].unsqueeze(-1) # [B, T_full-1, 1]\n", "\n", " # policy logp\n", " policy_logp = torch.gather(logprobs, -1, target_ids).squeeze(-1) # [B, T-1]\n", " # ref logp\n", " ref_logp = torch.gather(ref_logprobs, -1, target_ids).squeeze(-1) # [B, T-1]\n", "\n", " # clamp both to avoid -inf / inf\n", " policy_logp = policy_logp.clamp(min=-20.0, max=0.0)\n", " ref_logp = ref_logp.clamp(min=-20.0, max=0.0)\n", "\n", " # KL per token\n", " kl_per_token = (policy_logp - ref_logp) * mask_float\n", "\n", " # sanitize per-token KL\n", " kl_per_token = torch.where(torch.isfinite(kl_per_token),\n", " kl_per_token,\n", " torch.zeros_like(kl_per_token))\n", "\n", "\n", " # Average KL *per sequence*\n", " kl_per_seq_sum = kl_per_token.sum(dim=1)\n", " valid_tokens_per_seq = mask_float.sum(dim=1).clamp(min=1)\n", " kl_per_seq = kl_per_seq_sum / valid_tokens_per_seq\n", "\n", " # Average KL *over the batch*\n", " kl_loss = kl_per_seq.mean()\n", "\n", " # Store component values for logging before combining\n", " policy_loss_item = loss.item()\n", " kl_loss_item = kl_loss.item()\n", " if not torch.isfinite(kl_loss):\n", " kl_loss = torch.tensor(0.0, device=loss.device)\n", " # Combine policy loss and KL regularization\n", " loss = loss + self.kl_beta * kl_loss\n", "\n", " self.log({\n", " \"em_batch\": em_batch,\n", " \"policy_loss\": policy_loss_item,\n", " \"kl_loss\": kl_loss_item,\n", " \"loss\": loss.item()\n", " })\n", "\n", " outputs = {\n", " \"loss\": loss,\n", " \"policy_loss\": policy_loss_item,\n", " \"kl_loss\": kl_loss_item\n", " }\n", " return (loss, outputs) if return_outputs else loss" ] }, { "cell_type": "markdown", "metadata": { "id": "AK65nKUHk34b" }, "source": [ "Time to train! This step will take much longer than the SFT (and you very well may not get much improvement).\n", "\n", "Think about how much work is required to get a signal (compared to SFT) and how dense that signal is (compared to SFT)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Nfi-MgdQvHJa" }, "outputs": [], "source": [ "import torch\n", "import os, logging, warnings\n", "from transformers import (\n", " AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, BitsAndBytesConfig\n", ")\n", "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n", "from transformers.utils import logging as hf_logging\n", "\n", "hf_logging.set_verbosity_error()\n", "hf_logging.enable_progress_bar()\n", "\n", "tok = AutoTokenizer.from_pretrained(\"./gsm8k_1b_lora_merged\", use_fast=True)\n", "tok.pad_token = tok.eos_token\n", "tok.padding_side = \"left\"\n", "rl_collator = PromptOnlyCollator(tok)\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " \"./gsm8k_1b_lora_merged\",\n", " device_map=\"auto\",\n", " torch_dtype=torch.float16,\n", " attn_implementation=\"sdpa\",\n", ")\n", "model.generation_config.return_dict_in_generate = True\n", "model.generation_config.output_scores = False\n", "model.config.use_cache = False\n", "\n", "# LoRA config\n", "lora_config = LoraConfig(\n", " r=16, # 8–16 common; raise if updates feel too weak\n", " lora_alpha=32,\n", " lora_dropout=0.05,\n", " bias=\"none\",\n", " task_type=\"CAUSAL_LM\",\n", " target_modules=[\n", " \"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\n", " \"gate_proj\",\"up_proj\",\"down_proj\",\n", " ],\n", ")\n", "model = get_peft_model(model, lora_config)\n", "model.print_trainable_parameters()\n", "\n", "ref_model = AutoModelForCausalLM.from_pretrained(\n", " \"./gsm8k_1b_lora_merged\",\n", " device_map=\"auto\",\n", " torch_dtype=torch.float16,\n", " attn_implementation=\"sdpa\",\n", ")\n", "ref_model.eval()\n", "ref_model.requires_grad_(False)\n", "\n", "\n", "args = TrainingArguments(\n", " output_dir=\"./gsm8k_rl_lora\",\n", " per_device_train_batch_size=4,\n", " per_device_eval_batch_size=4,\n", " gradient_accumulation_steps=16,\n", " max_grad_norm=1.0,\n", " learning_rate=1e-6,\n", " weight_decay=0.0,\n", " num_train_epochs=0.1,\n", " logging_steps=1,\n", " eval_strategy=\"steps\",\n", " eval_steps=100,\n", " save_steps=500,\n", " save_total_limit=2,\n", " optim=\"adamw_torch\",\n", " fp16=True,\n", " gradient_checkpointing=False,\n", " report_to=\"none\",\n", " remove_unused_columns=False,\n", ")\n", "\n", "\n", "trainer = REINFORCETrainer(\n", " model=model,\n", " args=args,\n", " train_dataset=train_tok_rl,\n", " eval_dataset=val_tok_rl.select(range(64)),\n", " data_collator=rl_collator,\n", " processing_class=tok,\n", " gen_kwargs=dict(max_new_tokens=300, min_new_tokens=1, do_sample=True, top_p=0.9),\n", " ref_model=ref_model,\n", " kl_beta=0.1,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mvYbN6HXCBsP" }, "outputs": [], "source": [ "trainer.train()\n", "trainer.save_model()\n", "model.push_to_hub(\"/gsm8krl-1b-lora1\")\n", "tok.push_to_hub(\"/gsm8krl-1b-lora1\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "bzyGQdxEP7mk" }, "outputs": [], "source": [ "from peft import AutoPeftModelForCausalLM\n", "tok = AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B\", use_fast=True)\n", "\n", "# load the trained adapter (it knows the base model from its config)\n", "peft_model = AutoPeftModelForCausalLM.from_pretrained(\n", " \"/gsm8krl-1b-lora1\", torch_dtype=\"auto\", device_map=\"auto\"\n", ")\n", "\n", "# merge LoRA weights into the base weights and drop PEFT wrappers\n", "merged = peft_model.merge_and_unload()\n", "\n", "# save a standard HF model folder\n", "merged.save_pretrained(\"./gsm8krl_1b_lora1_merged\", safe_serialization=True)\n", "tok.save_pretrained(\"./gsm8krl_1b_lora1_merged\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "c3lDi4CQQFYK" }, "outputs": [], "source": [ "accuracy, outs, generated_answers = evaluate_gsm8k(\"./gsm8krl_1b_lora1_merged\", gsm8k_test_sample[\"question\"], gsm8k_test_sample[\"numeric_answer\"], 16)\n", "print(f\"Accuracy: {accuracy}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "g6eQUmNCncc8" }, "source": [ "Please record your accuracy here:\n", "\n", "\\\n", "\n", "And 5 examples here:\n", "\n", "\\\n", "\n", "Do you notice anything different about the model behavior?\n", "\n", "\\<3-5 sentence reflection\\>" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5COo-6koZigT" }, "outputs": [], "source": [ "#written by Daniel Zhang and Joey Huang" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "L4", "machine_shape": "hm", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }