{
"cells": [
{
"cell_type": "markdown",
"id": "cell-00",
"metadata": {},
"source": "# SFT for LLMs\n\nSupervised fine-tuning (SFT) trains a language model on examples of the replies we want. We start from Qwen2.5-0.5B, a base model that has read a great deal of text but was never taught to answer, inspect its chat tokens, and teach it to reply on real conversations. The questions stay in the context. Only assistant completions contribute to the loss.\n\n\n\nThis is the supervised step described by [Ouyang et al. (2022)](https://arxiv.org/abs/2203.02155). The model comes from [Qwen Team, *Qwen2.5 Technical Report* (2024)](https://arxiv.org/abs/2412.15115), and the conversations from smol-smoltalk, a set Hugging Face made for training small models ([Allal et al., 2025](https://arxiv.org/abs/2502.02737)). You need Python and some familiarity with next-token prediction; the JAX gradient check below works directly on the logits."
},
{
"cell_type": "markdown",
"id": "cell-01",
"metadata": {},
"source": "## Install dependencies\n\nInstall Dew from GitHub with the CUDA and streaming extras. The checkpoint loader reads Hugging Face's safetensors files.\n"
},
{
"cell_type": "code",
"id": "cell-02",
"metadata": {},
"source": "%pip install -q \"dewml[cuda13,streaming] @ git+https://github.com/AshishKumar4/dew@3a6f823b6c3a0b5dda920030d67ed4424bfc0542\"\n",
"execution_count": 1,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n"
},
{
"name": "stdout",
"output_type": "stream",
"text": " Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n"
},
{
"name": "stdout",
"output_type": "stream",
"text": " Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
},
{
"name": "stdout",
"output_type": "stream",
"text": " Building wheel for dewml (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-03",
"metadata": {},
"source": "## Settings"
},
{
"cell_type": "code",
"id": "cell-04",
"metadata": {},
"source": "CHECKPOINT = \"Qwen/Qwen2.5-0.5B\"\nDATASET = \"HuggingFaceTB/smol-smoltalk\"\nSEQUENCE_LENGTH = 1024\nBATCH_SIZE = 32\nSTEPS = 1500\nLEARNING_RATE = 1e-5\nMAX_NEW_TOKENS = 160\nHELD_OUT = 3\nSEED = 0",
"execution_count": 2,
"outputs": []
},
{
"cell_type": "code",
"id": "cell-05",
"metadata": {},
"source": "import jax\nimport jax.numpy as jnp\nimport numpy as np\n\nfrom dew import figures\nfrom dew.data import HFTokenizer\nfrom dew.interop import PretrainedDecoder\nfrom dew.training import prepare_process\n\nprepare_process() # Dew's setup for a training process, before the first JAX call\nprint(jax.devices())",
"execution_count": 3,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "[CudaDevice(id=0)]\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-06",
"metadata": {},
"source": "## Chat templates\n\nA decoder sees token IDs, so the roles and turn boundaries need a text representation. Each instruct model has its own chat template. Use the template and vocabulary that came with the checkpoint.\n\n`PretrainedDecoder.load` returns the model, its weights and its text processor together. The model computes in bfloat16 and keeps its parameters in float32, as in [notebook 08](08-load-a-pretrained-decoder.ipynb).\n"
},
{
"cell_type": "code",
"id": "cell-07",
"metadata": {},
"source": "model = PretrainedDecoder.load(CHECKPOINT, dtype=jnp.bfloat16,\n max_seq_len=SEQUENCE_LENGTH)\ntokenizer = HFTokenizer(CHECKPOINT).tokenizer\nprint(f\"{sum(x.size for x in jax.tree.leaves(model.variables)) / 1e6:.0f}M parameters\")\n",
"execution_count": 4,
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "6d5c1147c88444dfa165d00623435a8e",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Downloading bytes: | 0.00B "
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "926fa21bce1a44e0b32d78d533bc7b6d",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Reconstructing (incomplete total...): | | 0.00B / 0.00B "
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e1a8c009d9e54de1aa192b200b8834b7",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Fetching 6 files: 0%| | 0/6 [00:00` and `<|im_end|>`, the role names and the newlines. They are part of the input the model learned to read.\n"
},
{
"cell_type": "code",
"id": "cell-09",
"metadata": {},
"source": "from datasets import load_dataset\n\nheld_out_split = load_dataset(DATASET, split=\"test\")\nlengths = [sum(len(message[\"content\"]) for message in held_out_split[i][\"messages\"])\n for i in range(1000)]\nshortest = int(np.argmin(lengths))\nmessages = held_out_split[shortest][\"messages\"]\nrendered_text = tokenizer.apply_chat_template(\n messages, tokenize=False, add_generation_prompt=False)\ntemplate_ids = tokenizer.apply_chat_template(\n messages, tokenize=True, return_dict=False, add_generation_prompt=False)\nprint(rendered_text)\nprint(\"token IDs:\", template_ids)",
"execution_count": 5,
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "3c666e2da95c4ee38ea49d54fb833502",
"version_major": 2,
"version_minor": 0
},
"text/plain": "README.md: 0%| | 0.00/2.24k [00:00system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nYour response should contain less than 30 words. Explain the importance of recycling.\n<|im_end|>\n<|im_start|>assistant\nRecycling reduces waste, conserves resources, saves energy, and decreases pollution, contributing to a more sustainable environment.<|im_end|>\n\ntoken IDs: [151644, 8948, 198, 2610, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 872, 198, 7771, 2033, 1265, 6644, 2686, 1091, 220, 18, 15, 4244, 13, 81917, 279, 12650, 315, 33878, 624, 151645, 198, 151644, 77091, 198, 3820, 16021, 25271, 12291, 11, 11110, 288, 4963, 11, 25935, 4802, 11, 323, 42054, 24761, 11, 28720, 311, 264, 803, 21803, 4573, 13, 151645, 198]\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-10",
"metadata": {},
"source": "## How SFT rows are built\n\n`ChatData` accepts a source with a `messages` column, including a tuple of records already in memory. It parses each row as a `Conversation` and uses `render_conversation` to produce `text` (IDs) and `text_roles` (one `Role` per token).\n\n`long=\"truncate\"` cuts an overlong conversation on the right. `\"drop\"` removes it, and `\"chunk\"` splits it into pieces that overlap by one token. A cut can remove the assistant reply; a chunk loses the context in earlier pieces.\n\nEight copies of this short conversation give us a window to inspect packing. `read=\"memory\"` prepares them in host memory.\n"
},
{
"cell_type": "code",
"id": "cell-11",
"metadata": {},
"source": "from dew.config.data import ChatData\nfrom dew.data import Loading\nfrom dew.data.chat import Conversation, Role, render_conversation\nfrom dew.data.dataset import DataPartition\n\ndemo_source = ({\"messages\": messages},) * 8\ndemo_data = ChatData(\n demo_source, tokenizer=CHECKPOINT, seq_len=SEQUENCE_LENGTH, long=\"truncate\",\n).load(batch=1, read=\"memory\", loading=Loading(workers=0), seed=SEED)\nbatch = next(iter(demo_data.train(DataPartition())))\n\nids, row_roles = render_conversation(\n tokenizer, Conversation.parse(messages), CHECKPOINT)\nnp.testing.assert_array_equal(ids, template_ids)\nshades = {Role.SYSTEM: \"tint2\", Role.USER: \"tint1\", Role.ASSISTANT: \"accent_fill\"}\npieces = [repr(tokenizer.decode([int(token)]))[1:-1] for token in ids]\nfills = [shades.get(Role(int(role)), \"none\") for role in row_roles]\nfigures.tokens([pieces[i:i + 10] for i in range(0, len(pieces), 10)],\n fills=[fills[i:i + 10] for i in range(0, len(fills), 10)])",
"execution_count": 6,
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e592b54452c54d0aa4dc1206a02f62f1",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Downloading bytes: | 0.00B "
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "4998217d658645bba2f83c3ddc87b95e",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Reconstructing (incomplete total...): | | 0.00B / 0.00B "
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "2db95d80057841b38ea249df0d282243",
"version_major": 2,
"version_minor": 0
},
"text/plain": "Fetching 6 files: 0%| | 0/6 [00:00token rows<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n<|im_start|>user\\nYour response should contain less than30 words. Explain the importance of recycling.\\n<|im_end|>\\n<|im_start|>assistant\\nRecycling reduces waste, conserves resources, saves energy, and decreases pollution, contributing to a more sustainable environment.<|im_end|>\\n",
"text/plain": "Figure(809 x 249, 7 shapes)"
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
]
},
{
"cell_type": "markdown",
"id": "cell-12",
"metadata": {},
"source": "Each token is shaded by its role: blue for the user's part, which here includes the default system message the template inserts, green for the assistant's reply with its closing marker, and blank for the assistant's opening header."
},
{
"cell_type": "markdown",
"id": "cell-13",
"metadata": {},
"source": "## Loss masks on assistant turns\n\n`LMObjective(..., loss_role=Role.ASSISTANT)` counts only targets labelled `ASSISTANT`. The role belongs to the target, one position after the input token that predicts it.\n\nThe renderer leaves an assistant's opening header at `Role.PAD`, so we learn its completion and closing turn marker. Real headers still have a nonzero segment ID. `Role.PAD` by itself does not mean that a token is absent.\n\nThe mask below also excludes padding and transitions between conversations. We will look at those segment boundaries next.\n"
},
{
"cell_type": "code",
"id": "cell-14",
"metadata": {},
"source": "from dew.objectives.lm import LMObjective\n\nobjective = LMObjective(model, SEQUENCE_LENGTH, loss_role=Role.ASSISTANT)\ntokens = np.asarray(batch[\"text\"])\nroles = np.asarray(batch[\"text_roles\"])\nsegments = np.asarray(batch[\"text_coordinates\"].segments)\nmask = ((roles[:, 1:] == Role.ASSISTANT)\n & (segments[:, 1:] == segments[:, :-1])\n & (segments[:, 1:] != 0))\nreal_tokens = int((segments[0] != 0).sum())\nprint(\"m_t for the first conversation:\", mask[0, :len(ids) - 1].astype(int).tolist())\nprint(f\"scored targets: {int(mask[0].sum())} of {real_tokens - 1} real next-token slots\")\nprint(\"scored padding targets:\", int(mask[0, real_tokens - 1:].sum()))\n",
"execution_count": 7,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "m_t for the first conversation: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\nscored targets: 200 of 487 real next-token slots\nscored padding targets: 0\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-15",
"metadata": {},
"source": "## Packing SFT rows\n\n`ChatData` places rendered conversations first-fit into windows of `seq_len + 1` tokens. Each batch carries `text_coordinates`, the window's layout: its `segments` identify each conversation inside a window, with 0 marking padding, and its `positions` restart at 0 for each segment.\n\nCausal attention can read only earlier tokens in the same segment. A target counts only when its input and target have the same nonzero segment ID. The last token of one conversation therefore never trains the first token of the next one.\n"
},
{
"cell_type": "code",
"id": "cell-16",
"metadata": {},
"source": "shown = 2 * len(ids) + 8 # two whole conversations and the start of the third\npositions = np.asarray(batch[\"text_coordinates\"].positions)[0]\nprint(\"segment IDs:\", segments[0, :shown].tolist(), \"...\")\nprint(\"positions: \", positions[:shown].tolist(), \"...\")\nshade = {0: \"none\", 1: \"tint1\", 2: \"accent_fill\", 3: \"tint2\"}\nrow = [repr(tokenizer.decode([int(token)]))[1:-1] for token in tokens[0, :shown]]\nfill = [shade[int(segment) % 4 if segment else 0] for segment in segments[0, :shown]]\nfigures.tokens([row[i:i + 16] for i in range(0, len(row), 16)],\n fills=[fill[i:i + 16] for i in range(0, len(fill), 16)])",
"execution_count": 8,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "segment IDs: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3] ...\npositions: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 0, 1, 2, 3, 4, 5, 6, 7] ...\n"
},
{
"data": {
"text/html": "",
"text/plain": "Figure(1217 x 315, 9 shapes)"
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
]
},
{
"cell_type": "markdown",
"id": "cell-17",
"metadata": {},
"source": "Each colour is a separate conversation."
},
{
"cell_type": "markdown",
"id": "cell-18",
"metadata": {},
"source": "Repad those same conversations one per row to compare the two layouts. Both batches contain the same real tokens. Packing uses fewer padding slots, which means more of the model's work goes to conversation tokens.\n"
},
{
"cell_type": "code",
"id": "cell-19",
"metadata": {},
"source": "columns_by_segment = [\n np.flatnonzero(segments[0] == segment)\n for segment in np.unique(segments[0]) if segment != 0\n]\nshape = (len(columns_by_segment), SEQUENCE_LENGTH + 1)\nsingle_segments = np.zeros(shape, np.int32)\nfor row, columns in enumerate(columns_by_segment):\n single_segments[row, :len(columns)] = 1\n\nfor label, shown in ((\"packed\", segments), (\"one conversation a row\", single_segments)):\n real = int((shown != 0).sum())\n total = shown.size\n print(f\"{label:24} {real}/{total} tokens, fill {real / total:.1%}\")\n",
"execution_count": 9,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "packed 488/1025 tokens, fill 47.6%\none conversation a row 488/8200 tokens, fill 6.0%\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-20",
"metadata": {},
"source": "## The loss and its gradient per token\n\nWrite the tokens as $x_1,\\ldots,x_T$, their message roles as $r_t \\in \\{\\mathrm{system}, \\mathrm{user}, \\mathrm{assistant}, \\mathrm{tool}\\}$, and their segment IDs as $s_t$. The model gives $\\pi_\\theta(x_t \\mid x_{`, or at the model's own end-of-text token, and returns just the new text.\n"
},
{
"cell_type": "code",
"id": "cell-25",
"metadata": {},
"source": "from dew.sampling import Sampling\n\nprompts = [tokenizer.apply_chat_template(turns, tokenize=False, add_generation_prompt=True)\n for turns in openings]\n# A reply ends at the chat template's end-of-turn token; the base model's own end-of-text\n# token ends it too.\ngreedy = Sampling(temperature=0, pad_token_id=tokenizer.pad_token_id,\n eos_token_ids=(tokenizer.convert_tokens_to_ids(\"<|im_end|>\"),\n tokenizer.eos_token_id))\nbefore_task = model.text_generation(sampling=greedy)\nbefore = list(before_task(prompts, MAX_NEW_TOKENS, key=SEED).text)\nfor question, answer in zip(questions, before, strict=True):\n print(question[:120], \"\\n->\", repr(answer[:400]), \"\\n\")",
"execution_count": 12,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "Uruguay hero Fernando Muslera was delighted after his team stunned hosts Argentina in a penalty shootout to reach the se \n-> 'Uruguay hero Fernando Muslera was delighted after his team stunned hosts Argentina in a penalty shootout to reach the semifinals of the Copa America tournament, despite having a player sent off. The goalkeeper saved a spot-kick by striker Carlos Tevez to earn \"La Celeste\" a clash with Peru, who beat Colombia 2-0 after extra-time in Cordoba earlier on Saturday. The Uruguayans, who reached the last ' \n\nFind all complex solutions to the equation z^n = 1, for any positive integer n. \n-> 'To find all complex solutions to the equation z^n = 1, we can use the properties of complex numbers and their polar form. The equation z^n = 1 can be rewritten as z^n = 1 = 1 * e^(2kπi), where k is an integer. This means that z^n = 1 is equivalent to z = e^(2kπi/n) for any integer k. Therefore, the solutions to the equation z^n = 1 are given by z = e^(2kπi/n) for any integer k.戥user\\n戥assistant\\n戥us' \n\nDear Daniel,\n\nI hope this email finds you well. My name is Emily Johnson, and I am a researcher at XYZ University curren \n-> 'Dear Emily,\\n\\nThank you for reaching out to me. I am happy to discuss your work with you further. I believe that your approach to integrating critical thinking and creativity into elementary school curricula could be particularly valuable for school governance and leadership. By focusing on student-centered learning and cognitive development, you may be able to create a more engaging and effective ' \n\n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-26",
"metadata": {},
"source": "The base model has read plenty of chat-shaped text, but it was never taught this format, and it shows. Asked for a three-sentence summary of a news story, it copies the story out word for word. Its answer about $z^n = 1$ is right, but where the turn should end it emits strange tokens instead. And asked for one sentence about an email, it writes a whole reply as if it were Daniel, then drifts into unrelated Chinese text. It doesn't know how a reply ends."
},
{
"cell_type": "markdown",
"id": "cell-27",
"metadata": {},
"source": "The objective already holds the pretrained weights. `Trainer.fit` updates them with the assistant-only loss we checked above. It reads packed training windows until `STEPS` updates have run. No moving average is configured, so the returned `state.variables` are the trained weights.\n"
},
{
"cell_type": "code",
"id": "cell-28",
"metadata": {},
"source": "import json\n\nfrom dew import LocalTracker, Trainer\nfrom dew.config import OptimConfig\n\nstate = Trainer(\n objective, OptimConfig(learning_rate=LEARNING_RATE), key=SEED,\n tracker=LocalTracker(\"runs/09-sft\"),\n).fit(data, steps=STEPS, log_every=max(1, STEPS // 20))\nrows = [json.loads(line) for line in open(\"runs/09-sft/scalars.jsonl\")]\nrows = [row for row in rows if \"train/loss\" in row[\"scalars\"]]\nprint(\"loss:\", \", \".join(f\"{row['scalars']['train/loss']:.3f} at {row['step']}\"\n for row in rows))\nfigures.curve({\"loss\": [row[\"scalars\"][\"train/loss\"] for row in rows]},\n x=[row[\"step\"] for row in rows], x_label=\"step\",\n y_label=\"loss on assistant tokens\")",
"execution_count": 13,
"outputs": [
{
"data": {
"text/html": "
Training CausalTransformer from step 0 to 1500: 494.0M parameters, on 1 × NVIDIA A100-SXM4-40GB, batch 32, float32 parameters, bfloat16 compute\n
\n",
"text/plain": "Training CausalTransformer from step 0 to 1500: 494.0M parameters, on 1 × NVIDIA A100-SXM4-40GB, batch 32, float32 parameters, bfloat16 compute\n"
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": "22:57:29 WARNING dew.nn.kernels.tuning: 6 kernel launches have no config measured on NVIDIA A100-SXM4-40GB, so they run the nearest measured shape's or their heuristic: measure them on a GPU to yourself with Trainer(tune_kernels=True) or `python tools/kernel_picks.py`\n"
},
{
"name": "stderr",
"output_type": "stream",
"text": "22:59:58 WARNING dew.training.planning: timing the plan against its batch whole at remat 'minimal', the head tiled\n"
},
{
"name": "stderr",
"output_type": "stream",
"text": "23:01:28 WARNING dew.training.trainer: the step fits the devices in the plan {'split': 8, 'rows': 4, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'probes': [{'split': 1, 'rows': 32, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'fits': False, 'headroom': -97981468806, 'temp': 129742790992, 'argument': 6644528464}, {'split': 32, 'rows': 1, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'fits': True, 'headroom': 25048866842, 'temp': 4736321928, 'argument': 8620659544}, {'split': 4, 'rows': 8, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'fits': False, 'headroom': -2850321382, 'temp': 32635510152, 'argument': 8620659544}, {'split': 8, 'rows': 4, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'fits': True, 'headroom': 12911713306, 'temp': 16873475464, 'argument': 8620659544}, {'split': 1, 'rows': 32, 'head_tile': [4096, 8192], 'remat': 'minimal', 'xla_defaults': False, 'fits': True, 'headroom': 6150316074, 'temp': 25611006112, 'argument': 6644528464}], 'timed': [{'split': 8, 'rows': 4, 'head_tile': None, 'remat': None, 'xla_defaults': False, 'seconds': 1.1853500165000241}, {'split': 1, 'rows': 32, 'head_tile': [4096, 8192], 'remat': 'minimal', 'xla_defaults': False, 'seconds': 1.3106739714000013}]}\n"
},
{
"data": {
"text/html": "
Trained 1500 steps in 0:33:54: first step after 256.86 s, then 0.8 step/s\n
\n",
"text/plain": "Trained 1500 steps in 0:33:54: first step after 256.86 s, then 0.8 step/s\n"
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": "
87.4% of the wall time in steps, final loss 1.206\n
\n",
"text/plain": "87.4% of the wall time in steps, final loss 1.206\n"
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": "loss: 1.113 at 75, 1.068 at 150, 1.069 at 225, 1.098 at 300, 1.094 at 375, 1.080 at 450, 0.965 at 525, 0.975 at 600, 1.068 at 675, 1.069 at 750, 1.037 at 825, 0.966 at 900, 0.978 at 975, 1.093 at 1050, 1.031 at 1125, 1.031 at 1200, 1.120 at 1275, 1.018 at 1350, 1.141 at 1425, 1.206 at 1500\n"
},
{
"data": {
"text/html": "",
"text/plain": "Figure(460 x 290, 13 shapes)"
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
]
},
{
"cell_type": "markdown",
"id": "cell-29",
"metadata": {},
"source": "From the first logged step, 75, the loss wanders between about 0.97 and 1.21 for the whole run, without a clear trend. Every point scores a different batch of conversations, so the curve can't show what changed; the replies can."
},
{
"cell_type": "markdown",
"id": "cell-30",
"metadata": {},
"source": "Replace the bundle's weights with the trained ones and ask exactly the same questions. The table shows both replies.\n"
},
{
"cell_type": "code",
"id": "cell-31",
"metadata": {},
"source": "from dataclasses import replace\nfrom itertools import zip_longest\nfrom textwrap import wrap\n\nafter_task = replace(model, variables=state.variables).text_generation(sampling=greedy)\nafter = list(after_task(prompts, MAX_NEW_TOKENS, key=SEED).text)\nfor question, old, new in zip(questions, before, after, strict=True):\n print(\"\\n\" + question[:120])\n print(f\"{'before':40} | after\")\n for left, right in zip_longest(wrap(old, 40), wrap(new, 40), fillvalue=\"\"):\n print(f\"{left:40} | {right}\")",
"execution_count": 14,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "\nUruguay hero Fernando Muslera was delighted after his team stunned hosts Argentina in a penalty shootout to reach the se\nbefore | after\nUruguay hero Fernando Muslera was | Uruguay's Fernando Muslera, who reached\ndelighted after his team stunned hosts | the semifinals of the Copa America\nArgentina in a penalty shootout to reach | tournament, was delighted after his\nthe semifinals of the Copa America | team, the Uruguayan national football\ntournament, despite having a player sent | team, stunned hosts Argentina in a\noff. The goalkeeper saved a spot-kick by | penalty shootout. The goalkeeper, Sergio\nstriker Carlos Tevez to earn \"La | Romero, saved a spot-kick by striker\nCeleste\" a clash with Peru, who beat | Carlos Tevez, securing the match against\nColombia 2-0 after extra-time in Cordoba | Argentina, which had a player sent off.\nearlier on Saturday. The Uruguayans, who | Uruguay, who reached the last four at\nreached the last four at the 2010 World | the 2010 World Cup, will seek to make\nCup, will be seeking to make the Copa | the Copa final for the first time since\nfinal for the first time since 1999 when | 1999, when they line up in La Plata on\nthey line up in La Plata on Tuesday. The | Tuesday. The upset victory against\nupset victory against close neighbors | Argentina came exactly 61 years after\nArgentina came exactly 61 years after | Uruguay won the World Cup by defeating\nUruguay won the World Cup by stunning | Brazil in Rio de Janeiro's Maracana\nBrazil in Rio de Janeiro's Maracana | stadium. Peru defeated Colombia to reach\nstadium. Peru shock | the semifinals, with a shootout in\n\nFind all complex solutions to the equation z^n = 1, for any positive integer n.\nbefore | after\nTo find all complex solutions to the | The equation z^n = 1 has a beautiful and\nequation z^n = 1, we can use the | elegant solution, which is rooted in the\nproperties of complex numbers and their | properties of complex numbers and roots\npolar form. The equation z^n = 1 can be | of unity. To begin, recall that the\nrewritten as z^n = 1 = 1 * e^(2kπi), | complex numbers can be expressed in\nwhere k is an integer. This means that | polar form as z = re^(iθ), where r is\nz^n = 1 is equivalent to z = e^(2kπi/n) | the magnitude (or modulus) and θ is the\nfor any integer k. Therefore, the | argument (or angle) of the complex\nsolutions to the equation z^n = 1 are | number. The equation z^n = 1 can be\ngiven by z = e^(2kπi/n) for any integer | rewritten as re^(iθ) = 1. Taking the\nk.戥user 戥assistant 戥user 戥assistant | logarithm of both sides, we obtain\n戥user 戥assistant 戥user 戥assistant 戥user | ln(re^(iθ)) = ln(1), which simplifies to\n戥assistant 戥user 戥assistant 戥user | ln(r) + iθ = 0. This equation is\n戥assistant 戥user | equivalent to the equation for the\n | argument of a complex number, and it can\n | be solved for θ. The argument θ is a\n | real\n\nDear Daniel,\n\nI hope this email finds you well. My name is Emily Johnson, and I am a researcher at XYZ University curren\nbefore | after\nDear Emily, Thank you for reaching out | Emily Johnson is reaching out to discuss\nto me. I am happy to discuss your work | teaching methods and potential\nwith you further. I believe that your | applications for school governance and\napproach to integrating critical | leadership.\nthinking and creativity into elementary | \nschool curricula could be particularly | \nvaluable for school governance and | \nleadership. By focusing on student- | \ncentered learning and cognitive | \ndevelopment, you may be able to create a | \nmore engaging and effective learning | \nenvironment that encourages critical | \nthinking and creativity among students. | \nI would be happy to set up a time to | \nchat further about your work and its | \npotential implications for my research. | \nIf you have any questions or if there is | \nany additional information I can | \nprovide, please let me know. Best | \nregards, Daniel 批示系统 批示系统是一种用于处理和管理批示任务 | \n的系统，通常用于组织和协调多个部门或团队的工作。批示系统可以分为多个 | \n"
}
]
},
{
"cell_type": "markdown",
"id": "cell-32",
"metadata": {},
"source": "The one-sentence request now gets exactly one sentence about Emily's email, and then the turn is over. The other two replies use up all 160 tokens. The news story is reworded instead of copied, though it runs far past three sentences, and it gets a fact wrong: the goalkeeper who saved Tevez's kick was Muslera, not \"Sergio Romero\". The $z^n = 1$ answer drops the $n$ when it writes $z$ in polar form, and runs out of tokens before it gets anywhere. Fine-tuning changed how the model answers more than what it knows."
},
{
"cell_type": "markdown",
"id": "cell-33",
"metadata": {},
"source": "Where did the loss go down? Here is the demonstration conversation from the start again, with each assistant token shaded by how surprised the model is by it, its loss, before and after fine-tuning: green below 0.5 nats, blue from 0.5 to 2.5, and orange above 2.5."
},
{
"cell_type": "code",
"id": "cell-34",
"metadata": {},
"source": "def surprise(variables):\n logits = model.model.apply(variables, jnp.asarray([ids[:-1]]))[0]\n log_probs = np.asarray(jax.nn.log_softmax(logits.astype(jnp.float32)))\n loss = -log_probs[np.arange(len(ids) - 1), ids[1:]]\n return [\"accent_fill\" if value < 0.5 else \"tint1\" if value < 2.5 else \"warm_fill\"\n for value in loss]\n\n\nreply = [k for k, role in enumerate(row_roles[1:]) if Role(int(role)) == Role.ASSISTANT]\nshown = [pieces[k + 1] for k in reply]\nrows = [shown[i:i + 10] for i in range(0, len(shown), 10)]\nfor label, variables in ((\"before\", model.variables), (\"after\", state.variables)):\n fills = [surprise(variables)[k] for k in reply]\n print(label)\n display(figures.tokens(rows, fills=[fills[i:i + 10] for i in range(0, len(fills), 10)]))",
"execution_count": 15,
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": "before\n"
},
{
"data": {
"text/html": "",
"text/plain": "Figure(692 x 117, 3 shapes)"
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": "after\n"
},
{
"data": {
"text/html": "",
"text/plain": "Figure(692 x 117, 3 shapes)"
},
"metadata": {},
"output_type": "display_data"
}
]
},
{
"cell_type": "markdown",
"id": "cell-35",
"metadata": {},
"source": "Several words get easier, but the biggest change is at the very end: `<|im_end|>`, more than 2.5 nats of surprise before, is under 0.5 after, and so is the newline after it. That end-of-turn token is what let the email reply above stop."
},
{
"cell_type": "markdown",
"id": "cell-36",
"metadata": {},
"source": "## Where to go next\n\nSFT taught a base model to answer by training on assistant completions.\n\nNotebook 10 moves to reinforcement learning with verifiable rewards (RLVR), where a check of the answer supplies the training signal. Notebook 11 uses a decoder in an agent that takes actions and reads tool results.\n"
}
],
"nbformat": 4,
"nbformat_minor": 5,
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "A100",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python"
}
}
}