{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"machine_shape": "hm",
"gpuType": "A100",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
"
]
},
{
"cell_type": "markdown",
"source": [
"### Reinforcement learning for language models\n",
"\n",
"Lecture 17 | CMU ANLP Spring 2026 | Instructor: Sean Welleck\n",
"\n",
"This notebook shows fine-tuning a language model with reinforcement for a language generation task. The task is reversing an input string.\n",
"\n",
"The model is first fine-tuned with supervised fine-tuning, then fine-tuned with a PPO-style algorithm (specifically, the clipped PPO loss with an advantage computation based on averaging the rewards of multiple outputs).\n",
"\n",
"The reward is 1 when the generated string is the reverse of the input string, and 0 otherwise.\n",
"\n",
"The implementation of each part is found in its respective script (`sft.py`, `reverse_ppo.py`, `evaluate.py`). For more background on the task formulation and SFT, you can see the [fine-tuning lecture notebook](https://github.com/cmu-l3/anlp-fall2025-code/blob/main/08_finetuning/finetuning_lm.ipynb).\n",
"\n",
"*Tip: if you'd like to modify the code, you can copy and paste a script into the notebook and modify the code*"
],
"metadata": {
"id": "YPVkOzffc7QC"
}
},
{
"cell_type": "code",
"source": [
"!wget https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/sft.py -O sft.py\n",
"!wget https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/names.txt -O names.txt\n",
"!wget https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/evaluate.py -O evaluate.py\n",
"!wget https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/reverse_ppo.py -O reverse_ppo.py"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "VjrpP-UsMVwL",
"outputId": "de4ca0b9-4c2b-45ee-b3a7-ed263aa906c0"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"--2025-10-23 13:12:45-- https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/sft.py\n",
"Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
"Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 6959 (6.8K) [text/plain]\n",
"Saving to: ‘sft.py’\n",
"\n",
"sft.py 100%[===================>] 6.80K --.-KB/s in 0s \n",
"\n",
"2025-10-23 13:12:45 (109 MB/s) - ‘sft.py’ saved [6959/6959]\n",
"\n",
"--2025-10-23 13:12:45-- https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/names.txt\n",
"Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
"Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 228145 (223K) [text/plain]\n",
"Saving to: ‘names.txt’\n",
"\n",
"names.txt 100%[===================>] 222.80K --.-KB/s in 0.004s \n",
"\n",
"2025-10-23 13:12:45 (55.1 MB/s) - ‘names.txt’ saved [228145/228145]\n",
"\n",
"--2025-10-23 13:12:46-- https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/evaluate.py\n",
"Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
"Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 6866 (6.7K) [text/plain]\n",
"Saving to: ‘evaluate.py’\n",
"\n",
"evaluate.py 100%[===================>] 6.71K --.-KB/s in 0s \n",
"\n",
"2025-10-23 13:12:46 (82.3 MB/s) - ‘evaluate.py’ saved [6866/6866]\n",
"\n",
"--2025-10-23 13:12:46-- https://raw.githubusercontent.com/cmu-l3/anlp-fall2025-code/refs/heads/main/16_rl_llm/reverse_ppo.py\n",
"Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
"Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
"HTTP request sent, awaiting response... 200 OK\n",
"Length: 10533 (10K) [text/plain]\n",
"Saving to: ‘reverse_ppo.py’\n",
"\n",
"reverse_ppo.py 100%[===================>] 10.29K --.-KB/s in 0s \n",
"\n",
"2025-10-23 13:12:46 (92.4 MB/s) - ‘reverse_ppo.py’ saved [10533/10533]\n",
"\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"### Supervised fine-tuning"
],
"metadata": {
"id": "IIwRo09kh16o"
}
},
{
"cell_type": "code",
"source": [
"!python sft.py"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "glj3hwq1Dj3j",
"outputId": "fc0386b7-cc10-402d-f71c-eb37adf0318f"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"2025-10-23 13:13:36.564906: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
"2025-10-23 13:13:36.582316: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
"WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
"E0000 00:00:1761225216.603470 4972 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
"E0000 00:00:1761225216.610028 4972 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
"W0000 00:00:1761225216.626781 4972 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225216.626809 4972 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225216.626812 4972 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225216.626815 4972 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"2025-10-23 13:13:36.631859: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
"To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
"The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n",
"Model parameters: 134,516,736\n",
"Total names: 32033\n",
"Train: 25626, Dev: 3203, Test: 3204\n",
"Sample batch shape:\n",
" input_ids: torch.Size([16, 19])\n",
" attention_mask: torch.Size([16, 19])\n",
" labels: torch.Size([16, 19])\n",
"\n",
"First example:\n",
"Input: <|pad|><|pad|><|startoftext|>Reverse the name: sarahy. Answer:<|sep|>yharas<|endoftext|>\n",
"Target (answer only): :<|sep|>yharas<|endoftext|>\n",
"Epoch 1/1: 100% 1602/1602 [03:41<00:00, 7.23it/s, loss=0.8201]\n",
"Validation: 100% 201/201 [00:08<00:00, 22.35it/s]\n",
"Epoch 1/1\n",
" Average training loss: 0.8201\n",
" Average validation loss: 0.4850\n",
"Model saved to model_sft.pt\n",
"SFT done. To evaluate use 'python evaluate.py '\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"### Evaluate the SFT model"
],
"metadata": {
"id": "YLmtv1_linst"
}
},
{
"cell_type": "code",
"source": [
"!python evaluate.py --checkpoint_path model_sft.pt"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "xIHTUDHSIdZz",
"outputId": "9a9bd0fc-00ae-45c4-f622-62d6429d95de"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Using device: cuda\n",
"\n",
"Loading model and tokenizer...\n",
"2025-10-23 13:17:41.564273: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
"2025-10-23 13:17:41.581229: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
"WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
"E0000 00:00:1761225461.602271 6067 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
"E0000 00:00:1761225461.608736 6067 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
"W0000 00:00:1761225461.625127 6067 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225461.625165 6067 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225461.625168 6067 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225461.625170 6067 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"2025-10-23 13:17:41.630038: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
"To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
"The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n",
"Loading checkpoint from model_sft.pt...\n",
"Model parameters: 134,516,736\n",
"\n",
"Loading data from names_split.json...\n",
"Train: 25626, Dev: 3203, Test: 500\n",
"\n",
"Manual Examples:\n",
"==================================================\n",
"sean → naes (expected: naes ) ✓\n",
"joel → leoj (expected: leoj ) ✓\n",
"chen → nehc (expected: nehc ) ✓\n",
"dareen → nederad (expected: neerad ) ✗\n",
"neel → leen (expected: leen ) ✓\n",
"akshita → atihska (expected: atihska ) ✓\n",
"ashish → hsihsa (expected: hsihsa ) ✓\n",
"manan → nana (expected: nanam ) ✗\n",
"sanidhya → ayhdinas (expected: ayhdinas ) ✓\n",
"\n",
"Manual Examples Accuracy: 7/9 = 77.8%\n",
"\n",
"==================================================\n",
"Evaluating on test set (sample size: 500)...\n",
"==================================================\n",
"Evaluating: 100% 500/500 [01:13<00:00, 6.80it/s]\n",
"\n",
"Test Set Exact Match Accuracy: 64.80%\n",
"Number of errors: 176/500\n",
"\n",
"First 10 errors:\n",
"Name Expected Generated \n",
"==================================================\n",
"kaidence ecnediak ecidnaak \n",
"nadiah haidan hainad \n",
"price ecirp yep \n",
"uwais siawu siau \n",
"pantelis siletnap sileletp \n",
"rashi ihsar israh \n",
"raoul luoar luor \n",
"dilnoor roonlid roonlad \n",
"anthony ynohtna yotsna \n",
"zorie eiroz eirolz \n",
"\n",
"==================================================\n",
"SUMMARY\n",
"==================================================\n",
"Manual Examples Accuracy: 77.8%\n",
"Test Set Accuracy: 64.80%\n",
"Checkpoint: model_sft.pt\n"
]
}
]
},
{
"cell_type": "markdown",
"source": [
"### Wandb login\n",
"\n",
"You should create a [wandb account](https://wandb.ai/home). If you don't want to do that, you can comment out the wandb lines in the `reverse_ppo.py` script."
],
"metadata": {
"id": "XLnMtsdZirfa"
}
},
{
"cell_type": "code",
"source": [
"import wandb\n",
"wandb.login()"
],
"metadata": {
"id": "8plXrjcFJ5og",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 251
},
"outputId": "06f14bdb-d7b0-4ea3-b69a-b0e8c291db08"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stderr",
"text": [
"/usr/local/lib/python3.12/dist-packages/notebook/notebookapp.py:191: SyntaxWarning: invalid escape sequence '\\/'\n",
" | |_| | '_ \\/ _` / _` | _/ -_)\n"
]
},
{
"output_type": "display_data",
"data": {
"text/plain": [
""
],
"application/javascript": [
"\n",
" window._wandbApiKey = new Promise((resolve, reject) => {\n",
" function loadScript(url) {\n",
" return new Promise(function(resolve, reject) {\n",
" let newScript = document.createElement(\"script\");\n",
" newScript.onerror = reject;\n",
" newScript.onload = resolve;\n",
" document.body.appendChild(newScript);\n",
" newScript.src = url;\n",
" });\n",
" }\n",
" loadScript(\"https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js\").then(() => {\n",
" const iframe = document.createElement('iframe')\n",
" iframe.style.cssText = \"width:0;height:0;border:none\"\n",
" document.body.appendChild(iframe)\n",
" const handshake = new Postmate({\n",
" container: iframe,\n",
" url: 'https://wandb.ai/authorize'\n",
" });\n",
" const timeout = setTimeout(() => reject(\"Couldn't auto authenticate\"), 5000)\n",
" handshake.then(function(child) {\n",
" child.on('authorize', data => {\n",
" clearTimeout(timeout)\n",
" resolve(data)\n",
" });\n",
" });\n",
" })\n",
" });\n",
" "
]
},
"metadata": {}
},
{
"output_type": "stream",
"name": "stderr",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Logging into wandb.ai. (Learn how to deploy a W&B server locally: https://wandb.me/wandb-server)\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: You can find your API key in your browser here: https://wandb.ai/authorize\n",
"wandb: Paste an API key from your profile and hit enter:"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" ··········\n"
]
},
{
"output_type": "stream",
"name": "stderr",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: No netrc file found, creating one.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mwellecks12345\u001b[0m (\u001b[33mw1234\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
]
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"True"
]
},
"metadata": {},
"execution_count": 13
}
]
},
{
"cell_type": "markdown",
"source": [
"### Run the RL\n",
"\n",
"If you have Colab Pro, you can select the A100 runtime. It allows for using a larger batch size, which typically results in faster and more stable training for this application (feel free to try different batch sizes and compare yourself).\n",
"\n",
"The free GPU also works."
],
"metadata": {
"id": "_dR7-gNsi-6A"
}
},
{
"cell_type": "code",
"source": [
"\n",
"# If using A100:\n",
"!python reverse_ppo.py --batch_size 32\n",
"\n",
"# Otherwise:\n",
"# !python reverse_ppo.py --batch_size 16"
],
"metadata": {
"id": "77WNHO0_KK_e",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "951c57fb-78ce-43a7-89c2-c2e997fae125"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mwellecks12345\u001b[0m (\u001b[33mw1234\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[38;5;178m⢿\u001b[0m Waiting for wandb.init()...\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[38;5;178m⣻\u001b[0m setting up run vhvexv3h (0.2s)\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[38;5;178m⣽\u001b[0m setting up run vhvexv3h (0.2s)\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[38;5;178m⣾\u001b[0m setting up run vhvexv3h (0.2s)\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Tracking run with wandb version 0.22.2\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Run data is saved locally in \u001b[35m\u001b[1m/content/wandb/run-20251023_131953-vhvexv3h\u001b[0m\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Run \u001b[1m`wandb offline`\u001b[0m to turn off syncing.\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: Syncing run \u001b[33mSmolLM2-135M-rl-finetune\u001b[0m\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: ⭐️ View project at \u001b[34m\u001b[4mhttps://wandb.ai/w1234/name-reversal-rl\u001b[0m\n",
"\u001b[34m\u001b[1mwandb\u001b[0m: 🚀 View run at \u001b[34m\u001b[4mhttps://wandb.ai/w1234/name-reversal-rl/runs/vhvexv3h\u001b[0m\n",
"2025-10-23 13:19:56.891575: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
"2025-10-23 13:19:56.908529: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
"WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
"E0000 00:00:1761225596.929126 6681 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
"E0000 00:00:1761225596.935489 6681 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
"W0000 00:00:1761225596.951938 6681 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225596.951969 6681 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225596.951972 6681 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"W0000 00:00:1761225596.951974 6681 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
"2025-10-23 13:19:56.956910: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
"To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
"The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n",
"Model parameters: 134,516,736\n",
"Train: 25626, Dev: 3203, Test: 3204\n",
"Epoch 1/3: 0% 0/801 [00:00