From 231b4841430f6718c6bc4a9d260578aa5c4b50ba Mon Sep 17 00:00:00 2001 From: haijunz Date: Tue, 31 Mar 2026 00:20:43 +0000 Subject: [PATCH 01/12] add JEPO trainer --- examples/notebooks/jepo_math.ipynb | 4803 ++++++++++++++++++++++++++++ trl/__init__.py | 2 + trl/trainer/__init__.py | 2 + trl/trainer/jepo_config.py | 678 ++++ trl/trainer/jepo_trainer.py | 1838 +++++++++++ 5 files changed, 7323 insertions(+) create mode 100644 examples/notebooks/jepo_math.ipynb create mode 100644 trl/trainer/jepo_config.py create mode 100644 trl/trainer/jepo_trainer.py diff --git a/examples/notebooks/jepo_math.ipynb b/examples/notebooks/jepo_math.ipynb new file mode 100644 index 00000000000..2c2cc640dcf --- /dev/null +++ b/examples/notebooks/jepo_math.ipynb @@ -0,0 +1,4803 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "42ee0016", + "metadata": {}, + "source": [ + "To run this, press \"*Runtime*\" and press \"*Run all*\" on a **free** Tesla T4 Google Colab instance!\n", + "
\n", + "\n", + "\n", + " Join Discord if you need help + ⭐ Star us on Github ⭐\n", + "
\n", + "\n", + "To install Unsloth your local device, follow [our guide](https://docs.unsloth.ai/get-started/install-and-update). This notebook is licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme).\n", + "\n", + "You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & [how to save it](#Save)\n" + ] + }, + { + "cell_type": "markdown", + "id": "477bf000", + "metadata": {}, + "source": [ + "### News" + ] + }, + { + "cell_type": "markdown", + "id": "c0e36a31", + "metadata": {}, + "source": [ + "Long-Context GRPO for reinforcement learning — train stably at massive sequence lengths. Fine-tune models with up to 7x more context length efficiently. [Read Blog](https://unsloth.ai/docs/new/grpo-long-context)\n", + "\n", + "3× faster training with optimized sequence packing — higher throughput with no quality loss.[Read Blog](https://unsloth.ai/docs/new/3x-faster-training-packing)\n", + "\n", + "500k context-length fine-tuning — push long-context models further with memory-efficient training. [Read Blog](https://unsloth.ai/docs/new/500k-context-length-fine-tuning)\n", + "\n", + "Introducing FP8 precision training for faster RL inference. [Read Blog](https://docs.unsloth.ai/new/fp8-reinforcement-learning).\n", + "\n", + "Unsloth's [Docker image](https://hub.docker.com/r/unsloth/unsloth) is here! Start training with no setup & environment issues. [Read our Guide](https://docs.unsloth.ai/new/how-to-train-llms-with-unsloth-and-docker).\n", + "\n", + "Visit our docs for all our [model uploads](https://docs.unsloth.ai/get-started/all-our-models) and [notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks).\n" + ] + }, + { + "cell_type": "markdown", + "id": "9bb5a54a", + "metadata": {}, + "source": [ + "### Installation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3cd9ddaa", + "metadata": {}, + "outputs": [], + "source": [ + "% !pip install antlr4-python3-runtime==4.11\n", + "% !pip install sympy" + ] + }, + { + "cell_type": "markdown", + "id": "1c8bc6b4", + "metadata": {}, + "source": [ + "### Unsloth" + ] + }, + { + "cell_type": "markdown", + "id": "9e5c6cc8", + "metadata": {}, + "source": [ + "\n", + "Load up `Llama 3.1 8B Instruct`, and set parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d106a89", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Skipping import of cpp extensions due to incompatible torch version 2.8.0+cu128 for torchao version 0.15.0 Please see https://github.com/pytorch/ao/issues/2919 for more info\n", + "`torch_dtype` is deprecated! Use `dtype` instead!\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "999547f4658e4a679c494930d5e37013", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/4 [00:00\n", + "\n", + "We directly leverage [@willccbb](https://gist.github.com/willccbb/4676755236bb08cab5f4e54a0475d6fb) for data prep and all reward functions. You are free to create your own!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "286b2149", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "<>:176: SyntaxWarning: invalid escape sequence '\\s'\n", + "<>:176: SyntaxWarning: invalid escape sequence '\\s'\n", + "/tmp/ipykernel_298016/1337997636.py:176: SyntaxWarning: invalid escape sequence '\\s'\n", + " pattern = r\".*?\\s*.*?\"\n" + ] + } + ], + "source": [ + "import re\n", + "from datasets import load_dataset, Dataset\n", + "from sympy import simplify, sympify, parsing\n", + "from sympy.parsing.latex import parse_latex\n", + "import warnings\n", + "import numpy as np\n", + "\n", + "# Ignore only DeprecationWarnings\n", + "warnings.filterwarnings('ignore', category=DeprecationWarning)\n", + "\n", + "# Ignore only RuntimeWarnings\n", + "warnings.filterwarnings('ignore', category=RuntimeWarning)\n", + "\n", + "# Load and prep dataset\n", + "SYSTEM_PROMPT = \"\"\"\n", + "You are an AI that systematically works through problems step by step before delivering a final response. Your output must include a 'reasoning' section that contains your internal analysis, followed by an 'answer' section that provides only the final answer with no extra explanation.\n", + "Respond as follows:\n", + "\n", + "...\n", + "\n", + "\n", + "...\n", + "\n", + "\"\"\"\n", + "\n", + "XML_COT_FORMAT = \"\"\"\n", + "\n", + "{reasoning}\n", + "\n", + "\n", + "{answer}\n", + "\n", + "\"\"\"\n", + "\n", + "def extract_xml_answer(text: str) -> str:\n", + " answer = text.split(\"\")[-1]\n", + " answer = answer.split(\"\")[0]\n", + " return answer.strip()\n", + "\n", + "def extract_hash_answer(text: str) -> str | None:\n", + " # Extract the answer from the LaTeX string using the format \\boxed{answer}\n", + " extracted = re.search(r\"\\\\boxed{(.*)}\", text)\n", + " if extracted:\n", + " #in case there are multiple matches\n", + " # return the last match group, which is the content inside the \\boxed{}\n", + " return extracted.groups()[-1].strip()\n", + " return None\n", + "\n", + "def str_to_sympy(expr_str: str, raw: str):\n", + " expr = []\n", + " try:\n", + " expr.append(parse_latex(expr_str))\n", + " except Exception as e:\n", + " print(f\"Error converting {expr_str}, {raw} to sympy expression: {e}\")\n", + "\n", + " try:\n", + " expr.append(parsing.sympy_parser.parse_expr(expr_str))\n", + " except Exception as e:\n", + " print(f\"Error converting {expr_str} to sympy expression: {e}\")\n", + " return expr\n", + "\n", + "# uncomment middle messages for 1-shot prompting\n", + "def get_math_questions(split = \"train\") -> Dataset:\n", + " data = load_dataset('qwedsacf/competition_math')[split] # type: ignore\n", + " print(data.shape)\n", + " \n", + " # Improved data processing with proper error handling\n", + " def process_sample(x):\n", + " answer_text = extract_hash_answer(x['solution'])\n", + " \n", + " # Verify the extracted answer can be parsed\n", + " verified_count = 0\n", + " if answer_text:\n", + " verified_count = len(str_to_sympy(answer_text, x['solution']))\n", + " \n", + " return {\n", + " 'prompt': [\n", + " {'role': 'system', 'content': SYSTEM_PROMPT},\n", + " {'role': 'user', 'content': x['problem']}\n", + " ],\n", + " 'answer': answer_text,\n", + " 'raw': x['solution'],\n", + " 'verified': verified_count\n", + " }\n", + " data = data.map(process_sample) # type: ignore\n", + " \n", + " # Filter to keep only samples with successfully parsed answers\n", + " initial_size = len(data)\n", + " data = data.filter(lambda x: x['verified'] > 0 and x['answer'] is not None) # type: ignore\n", + " final_size = len(data)\n", + " print(f\"Filtered {initial_size - final_size} samples with invalid/unparseable answers. Kept {final_size} samples.\")\n", + " \n", + " return data # type: ignore\n", + "\n", + "dataset = get_math_questions()\n", + "print(dataset.shape)\n", + "\n", + "\n", + "def synpy_simplify_answer_check(response: str, answer: str) -> bool:\n", + " try:\n", + " if response is None or answer is None:\n", + " print(\"response or answer is None, cannot compare.\")\n", + " return False\n", + " \n", + " if response.strip() == answer.strip():\n", + " print(\"Exact string match found.\")\n", + " return True\n", + "\n", + " str_to_sympy_response = str_to_sympy(response, response)\n", + " str_to_sympy_answer = str_to_sympy(answer, answer)\n", + "\n", + " for r in str_to_sympy_response:\n", + " for a in str_to_sympy_answer:\n", + " try:\n", + " if simplify(r - a) == 0:\n", + " return True\n", + " except:\n", + " pass\n", + " return False\n", + " except Exception as e:\n", + " print(f\"Error simplifying answer: {e}\")\n", + " return False\n", + "\n", + "\n", + "\n", + "\n", + "def jepo_strict_format_reward_func(completions, **kwargs) -> list[float]:\n", + " \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n", + " #NOTE: re.DOTALL is critical here, otherwise the (.*) will not match newlines and we won't be able to extract the reasoning and answer correctly\n", + " pattern = re.compile(r\"(.*)\\s*(.*)\", re.DOTALL) \n", + " responses = [completion[0][\"content\"] for completion in completions]\n", + " matches = [re.match(pattern, r) for r in responses]\n", + " return [0.0 if match else -1.0 for match in matches]\n", + "\n", + "\n", + "\n", + "#given a completion and an answer, fabricate a ground truth completion that has the same reasoning but the answer is replaced with the ground truth answer\n", + "def fabricate_ground_truth_completion(completion: str, answer: str) -> (bool, str, str):\n", + " #NOTE: re.DOTALL is critical here, otherwise the (.*) will not match newlines and we won't be able to extract the reasoning and answer correctly\n", + " pattern = re.compile(r\"(.*)\\s*(.*)\", re.DOTALL) \n", + " match = re.match(pattern, completion)\n", + "\n", + " if not match:\n", + " #print(\"Completion does not have the correct format, cannot fabricate ground truth.\", f\"Completion: {completion}, Answer: {answer}\")\n", + " return False, completion, f\"{completion}\\n{answer}\\n\"\n", + " else:\n", + " reasoning = match.groups()[0].strip()\n", + " ground_truth_completion = f\"\\n{reasoning}\\n\\n\\n{answer}\\n\"\n", + " return True, f\"\\n{reasoning}\\n\", ground_truth_completion\n", + "\n", + "\n", + "\n", + "# Reward functions\n", + "def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:\n", + " responses = [completion[0]['content'] for completion in completions]\n", + " q = prompts[0][-1]['content']\n", + " extracted_responses = [extract_xml_answer(r) for r in responses]\n", + " print('-'*20, f\"Question:\\n{q}\", f\"\\nAnswer:\\n{answer[0]}\", f\"\\nResponse:\\n{responses[0]}\", f\"\\nExtracted:\\n{extracted_responses[0]}\")\n", + " return [2.0 if synpy_simplify_answer_check(r, a) else 0.0 for r, a in zip(extracted_responses, answer)]\n", + "'''\n", + "NOT using this\n", + "def int_reward_func(completions, **kwargs) -> list[float]:\n", + " responses = [completion[0]['content'] for completion in completions]\n", + " extracted_responses = [extract_xml_answer(r) for r in responses]\n", + " return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]\n", + "'''\n", + "def strict_format_reward_func(completions, **kwargs) -> list[float]:\n", + " \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n", + " pattern = r\"^\\n.*?\\n\\n\\n.*?\\n\\n$\"\n", + " responses = [completion[0][\"content\"] for completion in completions]\n", + " matches = [re.match(pattern, r) for r in responses]\n", + " return [0.5 if match else 0.0 for match in matches]\n", + "'''\n", + "def soft_format_reward_func(completions, **kwargs) -> list[float]:\n", + " \"\"\"Reward function that checks if the completion has a specific format.\"\"\"\n", + " pattern = r\".*?\\s*.*?\"\n", + " responses = [completion[0][\"content\"] for completion in completions]\n", + " matches = [re.match(pattern, r) for r in responses]\n", + " return [0.5 if match else 0.0 for match in matches]\n", + "'''\n", + "def count_xml(text) -> float:\n", + " count = 0.0\n", + " if text.count(\"\\n\") == 1:\n", + " count += 0.125\n", + " if text.count(\"\\n\\n\") == 1:\n", + " count += 0.125\n", + " if text.count(\"\\n\\n\") == 1:\n", + " count += 0.125\n", + " count -= len(text.split(\"\\n\\n\")[-1])*0.001\n", + " if text.count(\"\\n\") == 1:\n", + " count += 0.125\n", + " count -= (len(text.split(\"\\n\")[-1]) - 1)*0.001\n", + " return count\n", + "\n", + "def xmlcount_reward_func(completions, **kwargs) -> list[float]:\n", + " contents = [completion[0][\"content\"] for completion in completions]\n", + " return [count_xml(c) for c in contents]" + ] + }, + { + "cell_type": "markdown", + "id": "1bf16fff", + "metadata": {}, + "source": [ + "\n", + "### Train the model\n", + "\n", + "Now set up JEPO Trainer and all configurations!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cf7d7893", + "metadata": {}, + "outputs": [], + "source": [ + "#NOTE: max prompt length is 1628\n", + "#NOTE: max solution length is 2419 \n", + "#NOTE 1628 + 2419 = 4047, we set max_seq_length to 4608 to have some buffer for special tokens and generation\n", + "#NOTE 11929 samples\n", + "#batch size is 4 for H100 GPU\n", + "\n", + "\n", + "max_prompt_length = 1664\n", + "max_seq_length = 4608\n", + "\n", + "from trl import JEPOConfig, JEPOTrainer\n", + "temp_output_dir = current_dir / \"../jepo_math_norm\"\n", + "if os.path.exists(temp_output_dir):\n", + " print(f\"Removing existing temporary directory: {temp_output_dir}\")\n", + " shutil.rmtree(temp_output_dir)\n", + "# Ensure the output directory exists\n", + "os.makedirs(temp_output_dir, exist_ok=True)\n", + "training_args = JEPOConfig(\n", + " beta = 0.001,\n", + " learning_rate = 4e-6,\n", + " adam_beta1 = 0.9,\n", + " adam_beta2 = 0.99,\n", + " weight_decay = 0.1,\n", + " warmup_ratio = 0.1,\n", + " lr_scheduler_type = \"cosine\",\n", + " optim = \"paged_adamw_8bit\",\n", + " logging_steps = 1,\n", + " per_device_train_batch_size = 4, # Increase if you have more GPU memory\n", + " gradient_accumulation_steps = 1, # Increase to 4 for smoother training\n", + " num_generations = 4, # Decrease if out of memory\n", + " max_prompt_length = max_prompt_length,\n", + " max_completion_length = max_seq_length - max_prompt_length,\n", + " # num_train_epochs = 1, # Set to 1 for a full training run\n", + " #max_steps = 1,\n", + " num_train_epochs = 3, # Set to 3 for a full training run\n", + " \n", + " save_steps = 1000,\n", + " max_grad_norm = 0.1,\n", + " report_to = \"none\", # Can use Weights & Biases\n", + " output_dir = temp_output_dir,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8a16a521", + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"Training configuration: {training_args}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0424e437", + "metadata": {}, + "outputs": [], + "source": [ + "# get the token length of prompts\n", + "lengths = dataset.map(lambda x: {\"len\": len(tokenizer(\" \".join([d['content'] for d in x['prompt']]), truncation=False)['input_ids']),\n", + " 'solution_len': len(tokenizer(x['solution'], truncation=False)['input_ids'])\n", + "})\n", + "print(lengths.shape)\n", + "print(max(lengths['len']))\n", + "print(max(lengths['solution_len']))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cb55c724", + "metadata": {}, + "outputs": [], + "source": [ + "dataset[0]\n", + "\n", + "dataset_all = dataset\n", + "dataset_all.shuffle(seed=42)\n", + "\n", + "#75% for training, 25% for evaluation\n", + "train_size = int(0.75 * len(dataset_all))\n", + "\n", + "dataset = dataset_all.select(range(train_size))\n", + "eval_dataset = dataset_all.select(range(train_size, len(dataset_all)))\n", + "print(dataset.shape)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a2c38211", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "The model is already on multiple devices. Skipping the move to device specified in `args`.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "74c900196d824cb894d5e0d762d3d023", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading checkpoint shards: 0%| | 0/4 [00:00\n", + " \n", + " \n", + " [ 1001/26838 8:20:03 < 215:32:47, 0.03 it/s, Epoch 0.11/3]\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
StepTraining Loss
1-2.696400
2-1.145700
30.154800
4-0.101000
5-2.916200
6-2.978900
7-0.907400
8-1.637400
9-2.700700
100.092300
111.201500
12-0.537100
130.034600
14-2.414600
15-0.043900
160.080900
17-0.042500
18-0.018000
19-2.685000
200.041600
21-1.283200
221.842100
23-0.000400
24-1.998300
250.009200
260.040200
270.222000
28-0.111700
29-1.555400
30-0.157000
31-3.583600
32-2.846800
330.028900
340.141400
35-1.083500
36-2.772700
370.076800
38-0.272500
390.096000
40-1.073400
41-2.284600
420.105200
43-1.744000
440.131700
45-0.055400
460.171100
470.159000
48-2.024400
490.132700
500.056500
51-0.033500
521.284400
53-3.392100
54-1.962100
550.001500
56-1.752600
570.095100
580.031000
592.168400
600.078700
61-2.939600
620.105900
630.045700
64-2.696300
65-0.040300
660.023100
670.074700
680.016600
69-3.299400
700.035400
710.013400
72-0.001400
73-0.068200
740.083000
75-0.148900
760.052700
77-1.545800
780.083000
790.201100
800.173700
81-0.013000
82-0.535900
83-3.639400
840.018000
850.170700
86-0.170800
87-2.359300
88-0.057700
89-2.182300
90-0.009400
91-1.786000
920.008900
93-3.353100
94-2.462600
95-0.011500
96-2.816400
97-0.239400
980.040900
990.269400
1000.154800
101-2.085000
102-0.028700
1030.046900
104-0.054500
105-2.440600
106-1.495400
107-1.604600
1080.857100
1090.149200
110-1.078300
1110.028000
112-2.280100
113-1.708400
114-3.363500
1150.008400
1160.129200
117-0.002400
118-0.005300
1190.167600
120-1.369800
121-2.255600
1220.053900
1230.030000
1240.008200
1250.046000
1260.100100
127-1.509400
1280.207500
1290.136000
1300.067700
131-2.168700
132-0.022000
1330.109200
134-0.174800
1350.058500
1360.142200
137-1.690100
138-0.047700
139-2.075400
1400.087000
141-0.164100
142-0.546000
143-2.661400
144-0.019000
1450.012900
146-1.288500
1470.110000
148-1.148600
1490.216800
150-3.189000
151-0.387100
152-0.352500
1530.127600
1540.144700
1551.044200
156-2.586500
157-1.734000
1580.000200
1591.842100
160-1.612600
161-3.101100
162-0.026000
163-0.053900
164-0.897100
165-0.337200
166-0.139300
167-2.152100
1680.023300
169-0.004600
170-1.478800
171-1.004800
1720.030600
173-0.522100
174-0.180800
1750.179200
1760.002100
1770.009400
1780.173700
1790.170600
180-1.489900
1810.153300
182-3.308000
183-2.884200
184-2.303500
185-0.199500
186-0.892600
1870.200300
1880.041100
1890.051100
190-0.757600
191-0.015300
192-0.979100
193-0.514800
194-1.412500
195-3.489100
1960.273000
1970.079400
198-0.008600
199-2.918100
200-3.104900
201-0.955900
2020.049200
2030.190600
204-2.765600
205-3.075500
2060.071800
207-0.189700
2080.027300
209-0.884100
210-0.361600
211-3.159400
2120.061000
213-0.173100
2140.003200
2150.255900
216-0.098300
217-0.189500
2180.257100
219-0.061100
2200.102600
221-0.003600
2220.047700
2230.113300
2242.669400
2250.018700
226-1.641800
227-1.839500
2280.022600
229-0.051500
2300.109800
2310.012200
2320.153100
233-0.039800
234-0.437000
235-0.356900
2360.039200
237-1.119600
2380.049900
239-2.238600
2400.165500
241-2.869100
242-2.747400
243-0.084200
244-1.311900
245-0.248800
246-0.627600
247-0.018600
2482.025400
249-0.539000
250-1.410400
2510.087800
2520.061100
253-0.066800
2540.195500
2550.063300
256-3.353400
257-0.038000
258-0.714600
259-1.651500
2600.185300
261-2.932000
262-1.172300
263-0.000900
264-1.822100
2650.066000
2660.352700
267-0.067700
2680.038000
269-0.068100
270-1.991100
2710.212700
272-3.411800
273-2.866400
274-1.717400
275-0.165800
2760.014800
2770.058500
2780.207700
279-0.022300
280-0.131200
281-0.001800
2820.067500
283-0.258000
284-0.543100
285-0.074200
286-0.187300
287-0.191800
288-0.054700
289-0.065600
2900.506400
291-4.027500
2920.001200
2930.062300
294-3.115100
295-0.000100
296-0.019600
297-0.734600
298-0.354300
299-0.061500
3000.198800
301-0.031900
302-0.047500
303-0.052000
3040.008500
305-0.202900
306-1.625800
3070.492700
3080.004100
3090.010100
310-1.769100
3110.058500
312-2.569000
313-2.964900
314-0.798800
315-1.676200
3160.151300
317-0.011700
318-0.035400
319-3.011100
320-1.834900
321-0.644400
322-0.252500
3230.018600
3240.002400
325-1.122300
326-0.315100
327-1.108700
328-0.377300
3290.012400
3300.300600
3310.034100
332-2.868700
333-3.103900
334-2.834900
335-1.575900
336-2.621900
337-1.696800
338-1.001700
339-3.230000
3400.399000
341-3.393800
342-1.858100
343-0.895600
344-2.674200
345-2.199000
346-0.006100
3470.037600
3480.181000
349-0.508200
350-0.669400
351-0.034600
352-0.059300
353-0.086700
354-0.397800
3551.347200
356-0.161300
357-2.574900
3580.171600
3590.016900
360-2.935600
361-0.021200
362-2.929300
3630.089000
3640.042000
3650.090500
366-1.443800
367-0.027600
368-0.082100
3690.369700
370-2.665200
3710.005100
372-0.120000
373-2.020000
3740.069700
375-2.660800
376-0.003400
377-0.075200
3780.026600
379-0.197800
380-0.831300
3810.082400
382-0.027500
3830.040600
384-2.064900
3850.123400
386-0.121600
3870.051100
388-0.153400
3890.092600
3900.046500
391-1.645200
392-3.311800
393-0.062500
394-0.003600
3950.056900
396-2.811000
397-0.029200
398-2.608600
399-2.073000
4000.033500
401-1.276100
402-2.336900
403-1.199800
404-0.082500
405-2.603000
406-0.028300
407-0.006600
408-2.890600
4090.016200
4100.006100
4110.033400
4120.049000
4130.219600
414-0.122300
415-0.398200
4160.058200
417-2.069400
418-0.010900
419-3.705000
4200.107600
421-0.014300
4220.100100
423-2.435100
424-2.679500
425-0.072100
426-0.970800
427-0.005000
4280.010900
4290.043000
430-0.160300
4310.118900
432-0.059500
4330.188600
4340.109500
435-0.405300
436-0.039200
437-2.995800
4380.075600
439-0.194700
440-0.026500
4410.056200
442-3.214300
443-0.088800
444-0.073500
445-0.248700
446-2.727400
447-2.691000
448-2.862000
4490.158600
450-0.001800
451-0.001500
452-2.259800
453-0.168800
454-0.000600
455-2.378000
456-1.412400
457-0.054900
458-1.806800
459-0.027200
460-0.046100
4610.047400
4620.149200
463-0.239800
464-0.031000
465-2.821500
466-2.595100
467-2.667300
4680.018800
469-1.645500
470-2.263300
4710.015800
4720.114400
473-2.255300
474-0.001500
475-0.033200
4760.050400
4770.001000
478-0.055200
479-0.014500
480-2.134100
481-0.121400
482-2.483900
4830.123800
484-0.006300
485-1.525000
486-0.010600
487-1.265800
488-1.639700
489-1.603700
490-2.841900
4910.052300
492-0.868700
493-2.416200
494-1.778200
4950.028800
496-0.063100
497-1.550500
498-0.000500
499-1.381700
5000.014000
501-0.431700
502-2.728600
5030.002000
5040.080400
505-0.002300
506-2.994900
507-2.173100
5080.022300
509-0.106500
510-3.589600
5110.005500
512-0.012400
513-2.807500
514-0.042700
515-0.018200
516-0.281200
517-0.109800
5180.120300
519-0.010300
520-0.120000
521-0.008200
5220.511700
523-2.977200
5240.040000
525-0.301600
526-2.582100
527-2.656200
528-0.053700
529-0.036300
530-0.038000
531-2.099200
532-2.468400
5330.055900
5340.061800
535-0.179600
5360.193700
537-2.825300
538-2.898400
5390.031500
5400.051300
5410.075100
5420.003300
543-0.236400
544-1.447900
5450.043400
546-0.112300
547-0.004000
548-1.012900
549-1.672700
5500.002100
551-0.518700
552-0.938200
553-0.127800
554-0.268100
555-0.080700
556-1.655800
557-2.655000
558-0.010000
559-0.084300
5600.102100
561-0.508200
5620.014300
563-0.006900
564-2.803600
5650.020000
566-2.456000
567-2.202500
5680.329000
5690.031000
5700.017700
571-0.004600
5720.063300
5730.078400
574-0.958100
575-2.237700
576-0.103100
577-0.529300
5780.127300
579-0.110400
5800.206600
581-1.718000
5820.089500
5830.072000
584-0.002300
5850.008300
586-0.105200
587-0.027100
588-0.023400
589-0.000300
590-2.722500
591-0.090200
592-0.571600
5930.076700
5940.116600
595-2.682600
596-0.895400
597-0.264900
598-1.421000
5990.050600
6000.001300
601-3.315900
602-1.485900
6030.067400
6040.181400
605-0.464900
606-0.017700
607-2.459800
608-1.080900
609-0.009600
610-0.121100
611-0.350500
612-2.640200
6130.019500
6140.006600
615-0.000000
616-0.264100
617-0.858600
618-0.304900
619-0.741100
6200.069200
6210.065300
6220.030700
6230.105600
624-0.000300
625-0.573300
6260.001100
627-0.802500
6280.006800
629-0.921000
6300.040600
6310.019500
632-1.903200
633-0.000400
634-0.397500
635-0.059500
6360.000600
637-0.785000
638-0.010900
639-2.925300
6400.346400
641-0.456900
642-2.750100
643-1.413700
644-2.607100
645-1.993000
646-0.632400
647-0.243200
6480.665100
649-1.393100
6500.115900
651-2.527300
652-0.005000
653-0.000000
654-0.087800
655-0.240900
656-0.173900
657-0.005300
658-0.477400
659-0.099800
660-1.138000
6610.112300
662-0.043900
663-0.231500
6640.004800
665-1.292600
666-0.919900
6670.119000
6680.070100
669-3.057900
670-1.288600
6710.000300
672-0.000800
673-2.389800
674-0.735600
6750.000900
676-0.000200
6770.074700
678-2.448400
6790.328500
6800.017600
681-0.341900
6820.000100
6830.000000
684-0.000300
685-1.595200
686-0.133500
687-0.006900
688-0.103400
6890.057500
690-0.076600
6910.030500
692-0.092600
6930.002600
6940.033100
695-0.001200
696-0.084400
6970.000000
698-0.141600
6991.427800
700-3.011400
701-0.150000
7020.001300
703-3.159600
704-0.002000
705-0.244800
706-0.275200
7070.022200
708-2.752300
709-0.161500
7100.151600
711-2.370600
712-0.016600
713-0.004300
714-0.087000
7150.047700
716-0.637700
717-0.619700
718-0.001200
7190.067100
720-0.040100
721-0.000100
722-0.647000
723-2.528700
724-0.010700
725-2.753400
726-1.927600
7270.000500
728-0.004500
729-1.368800
730-2.052900
7310.024000
732-0.000000
733-0.082900
734-0.048900
7350.063800
736-2.362900
7370.006100
738-0.015200
7390.023000
7400.083000
741-2.225200
7420.048900
743-0.005100
744-2.294800
745-1.988700
7460.304500
747-0.115700
748-0.000200
749-0.036200
7500.001000
751-2.071700
752-0.085900
753-0.188100
754-2.027200
7550.070500
756-0.053000
757-0.739200
758-2.084000
759-2.533400
7600.000300
761-0.004400
762-2.832100
763-0.339100
764-0.089500
7650.041700
766-0.003500
7670.022100
7680.069400
769-0.155700
7700.000800
7710.089300
772-0.000400
773-2.663600
774-0.417300
775-0.323400
776-2.160900
777-0.000100
7780.075300
7790.405100
7800.015300
781-1.782800
782-0.001300
7830.305200
7840.065500
785-0.000500
7860.080700
787-0.057800
788-0.002600
789-0.001400
790-0.034800
791-0.000100
792-1.877500
7930.019500
7940.009600
795-0.105300
796-0.129900
797-0.311400
798-2.437700
799-0.000400
800-1.523900
801-1.492800
8020.039000
803-0.000700
804-0.021000
805-0.065900
806-0.027100
807-0.049800
8080.035400
809-0.001700
810-0.033200
811-0.196500
812-0.001000
813-0.032900
814-0.017100
8150.028700
816-0.171300
8170.015300
818-0.005900
819-0.426000
820-0.943700
821-0.000700
822-0.000300
823-0.004300
824-0.009700
8250.198600
8260.328900
8270.188800
8280.968800
829-0.963100
830-2.783300
831-0.474300
832-0.294600
833-0.148500
834-1.373600
8350.065700
836-0.182400
8370.032700
838-0.007400
8390.008000
840-0.039900
841-0.057900
8420.000500
843-0.352600
844-1.688000
8450.061000
846-0.000600
847-0.000100
8480.000200
849-0.001400
850-1.233800
8510.109100
852-0.000800
853-0.244900
854-0.038700
8550.243000
8560.000400
8570.029700
858-0.001800
8590.164600
860-0.329100
861-0.033400
8620.052800
863-0.468800
8640.000300
865-0.768100
866-0.061700
867-0.181800
8680.001100
8690.001600
870-1.857500
871-0.003400
8720.169300
873-0.036900
8740.010900
875-0.148800
8760.085100
8770.150200
878-0.004700
879-0.000700
880-0.137900
881-0.005300
882-0.408700
883-2.271500
8840.142300
885-0.066900
886-1.327700
887-0.062300
888-0.205300
8890.040900
890-0.548300
8910.039000
8920.045400
893-1.917200
894-0.202700
895-0.315300
896-0.198700
897-2.117000
8980.355500
899-0.217000
9000.122500
901-0.136400
902-0.859600
903-0.075500
904-0.215500
905-0.000800
906-0.097000
907-0.064300
908-0.290800
9090.095900
910-0.028900
911-0.001200
912-0.003500
913-0.005100
914-0.408500
915-0.021700
9160.000600
917-0.089700
9180.000700
9190.000300
9200.053300
9210.028100
9220.043800
923-0.095000
924-0.058000
9250.184300
926-0.042600
927-0.000300
928-0.065500
9291.266400
9300.007000
931-0.001700
932-0.001600
9330.002200
934-0.013400
9350.065000
9360.074500
937-0.122800
9380.012400
9390.006300
9400.001500
941-0.000100
942-0.553000
9430.103400
944-0.116500
945-2.889400
946-0.010600
9470.000300
9480.044500
9490.038800
9500.006700
951-0.014100
9520.002200
9530.198000
954-0.000400
955-0.057400
956-2.319500
9570.000100
9580.043000
959-0.001900
9600.052200
961-0.080400
962-0.390700
963-0.405800
964-3.307000
965-0.125900
9660.061800
9670.003300
968-0.005800
969-0.002400
9700.000100
9710.119900
9720.036500
9730.011900
974-1.097000
975-0.085300
976-0.000100
9770.000200
978-0.078800
9790.030700
9800.000100
981-0.038500
982-0.011800
983-0.005600
9840.049900
985-0.000200
986-0.083200
987-0.140700
9880.014100
989-0.061700
990-0.000200
991-0.000300
992-0.001000
9930.016900
994-0.235500
9950.254100
996-0.007700
997-0.000200
9980.025300
9990.042600

" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "ename": "RuntimeError", + "evalue": "[enforce fail at inline_container.cc:664] . unexpected pos 11819474816 vs 11819474712", + "output_type": "error", + "traceback": [ + "\u001b[31m---------------------------------------------------------------------------\u001b[39m\n", + "\u001b[31mRuntimeError\u001b[39m Traceback (most recent call last)\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:967\u001b[39m, in \u001b[36msave\u001b[39m\u001b[34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[39m\n", + "\u001b[32m 966\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m _open_zipfile_writer(f) \u001b[38;5;28;01mas\u001b[39;00m opened_zipfile:\n", + "\u001b[32m--> \u001b[39m\u001b[32m967\u001b[39m \u001b[43m_save\u001b[49m\u001b[43m(\u001b[49m\n", + "\u001b[32m 968\u001b[39m \u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 969\u001b[39m \u001b[43m \u001b[49m\u001b[43mopened_zipfile\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 970\u001b[39m \u001b[43m \u001b[49m\u001b[43mpickle_module\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 971\u001b[39m \u001b[43m \u001b[49m\u001b[43mpickle_protocol\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 972\u001b[39m \u001b[43m \u001b[49m\u001b[43m_disable_byteorder_record\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 973\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 974\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:1268\u001b[39m, in \u001b[36m_save\u001b[39m\u001b[34m(obj, zip_file, pickle_module, pickle_protocol, _disable_byteorder_record)\u001b[39m\n", + "\u001b[32m 1267\u001b[39m \u001b[38;5;66;03m# Now that it is on the CPU we can directly copy it into the zip file\u001b[39;00m\n", + "\u001b[32m-> \u001b[39m\u001b[32m1268\u001b[39m \u001b[43mzip_file\u001b[49m\u001b[43m.\u001b[49m\u001b[43mwrite_record\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstorage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_bytes\u001b[49m\u001b[43m)\u001b[49m\n", + "\n", + "\u001b[31mRuntimeError\u001b[39m: [enforce fail at inline_container.cc:858] . PytorchStreamWriter failed writing file data/496: file write failed\n", + "\n", + "During handling of the above exception, another exception occurred:\n", + "\n", + "\u001b[31mRuntimeError\u001b[39m Traceback (most recent call last)\n", + "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[15]\u001b[39m\u001b[32m, line 13\u001b[39m\n", + "\u001b[32m 2\u001b[39m os.environ[\u001b[33m\"\u001b[39m\u001b[33mCUDA_LAUNCH_BLOCKING\u001b[39m\u001b[33m\"\u001b[39m] = \u001b[33m\"\u001b[39m\u001b[33m1\u001b[39m\u001b[33m\"\u001b[39m\n", + "\u001b[32m 3\u001b[39m trainer = JEPOTrainer(\n", + "\u001b[32m 4\u001b[39m model = model,\n", + "\u001b[32m 5\u001b[39m processing_class = tokenizer,\n", + "\u001b[32m (...)\u001b[39m\u001b[32m 11\u001b[39m train_dataset = dataset,\n", + "\u001b[32m 12\u001b[39m )\n", + "\u001b[32m---> \u001b[39m\u001b[32m13\u001b[39m \u001b[43mtrainer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 14\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m torch.cuda.is_available():\n", + "\u001b[32m 15\u001b[39m memory_used = torch.cuda.memory_allocated() / \u001b[32m1024\u001b[39m**\u001b[32m3\u001b[39m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:2325\u001b[39m, in \u001b[36mTrainer.train\u001b[39m\u001b[34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[39m\n", + "\u001b[32m 2323\u001b[39m hf_hub_utils.enable_progress_bars()\n", + "\u001b[32m 2324\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n", + "\u001b[32m-> \u001b[39m\u001b[32m2325\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n", + "\u001b[32m 2326\u001b[39m \u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m=\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2327\u001b[39m \u001b[43m \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m=\u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2328\u001b[39m \u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2329\u001b[39m \u001b[43m \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m=\u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2330\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:2756\u001b[39m, in \u001b[36mTrainer._inner_training_loop\u001b[39m\u001b[34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[39m\n", + "\u001b[32m 2754\u001b[39m \u001b[38;5;28mself\u001b[39m.state.epoch = epoch + (step + \u001b[32m1\u001b[39m) / steps_in_epoch\n", + "\u001b[32m 2755\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_step_end(args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", + "\u001b[32m-> \u001b[39m\u001b[32m2756\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_maybe_log_save_evaluate\u001b[49m\u001b[43m(\u001b[49m\n", + "\u001b[32m 2757\u001b[39m \u001b[43m \u001b[49m\u001b[43mtr_loss\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2758\u001b[39m \u001b[43m \u001b[49m\u001b[43mgrad_norm\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2759\u001b[39m \u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2760\u001b[39m \u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2761\u001b[39m \u001b[43m \u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2762\u001b[39m \u001b[43m \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2763\u001b[39m \u001b[43m \u001b[49m\u001b[43mstart_time\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2764\u001b[39m \u001b[43m \u001b[49m\u001b[43mlearning_rate\u001b[49m\u001b[43m=\u001b[49m\u001b[43mlearning_rate\u001b[49m\u001b[43m,\u001b[49m\n", + "\u001b[32m 2765\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 2766\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n", + "\u001b[32m 2767\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_substep_end(args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3228\u001b[39m, in \u001b[36mTrainer._maybe_log_save_evaluate\u001b[39m\u001b[34m(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate)\u001b[39m\n", + "\u001b[32m 3225\u001b[39m \u001b[38;5;28mself\u001b[39m.control.should_save = is_new_best_metric\n", + "\u001b[32m 3227\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.control.should_save:\n", + "\u001b[32m-> \u001b[39m\u001b[32m3228\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_save_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 3229\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_save(\u001b[38;5;28mself\u001b[39m.args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/trl/trainer/jepo_trainer.py:1854\u001b[39m, in \u001b[36mJEPOTrainer._save_checkpoint\u001b[39m\u001b[34m(self, model, trial)\u001b[39m\n", + "\u001b[32m 1852\u001b[39m model_name = \u001b[38;5;28mself\u001b[39m.args.hub_model_id.split(\u001b[33m\"\u001b[39m\u001b[33m/\u001b[39m\u001b[33m\"\u001b[39m)[-\u001b[32m1\u001b[39m]\n", + "\u001b[32m 1853\u001b[39m \u001b[38;5;28mself\u001b[39m.create_model_card(model_name=model_name)\n", + "\u001b[32m-> \u001b[39m\u001b[32m1854\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_save_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m)\u001b[49m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3345\u001b[39m, in \u001b[36mTrainer._save_checkpoint\u001b[39m\u001b[34m(self, model, trial)\u001b[39m\n", + "\u001b[32m 3341\u001b[39m \u001b[38;5;28mself\u001b[39m.state.best_model_checkpoint = best_checkpoint_dir\n", + "\u001b[32m 3343\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m.args.save_only_model:\n", + "\u001b[32m 3344\u001b[39m \u001b[38;5;66;03m# Save optimizer and scheduler\u001b[39;00m\n", + "\u001b[32m-> \u001b[39m\u001b[32m3345\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_save_optimizer_and_scheduler\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_dir\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 3346\u001b[39m \u001b[38;5;28mself\u001b[39m._save_scaler(output_dir)\n", + "\u001b[32m 3347\u001b[39m \u001b[38;5;66;03m# Save RNG state\u001b[39;00m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3472\u001b[39m, in \u001b[36mTrainer._save_optimizer_and_scheduler\u001b[39m\u001b[34m(self, output_dir)\u001b[39m\n", + "\u001b[32m 3467\u001b[39m save_fsdp_optimizer(\n", + "\u001b[32m 3468\u001b[39m \u001b[38;5;28mself\u001b[39m.accelerator.state.fsdp_plugin, \u001b[38;5;28mself\u001b[39m.accelerator, \u001b[38;5;28mself\u001b[39m.optimizer, \u001b[38;5;28mself\u001b[39m.model, output_dir\n", + "\u001b[32m 3469\u001b[39m )\n", + "\u001b[32m 3470\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.args.should_save:\n", + "\u001b[32m 3471\u001b[39m \u001b[38;5;66;03m# deepspeed.save_checkpoint above saves model/optim/sched\u001b[39;00m\n", + "\u001b[32m-> \u001b[39m\u001b[32m3472\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mos\u001b[49m\u001b[43m.\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m.\u001b[49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mOPTIMIZER_NAME\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 3474\u001b[39m \u001b[38;5;66;03m# Save SCHEDULER & SCALER\u001b[39;00m\n", + "\u001b[32m 3475\u001b[39m is_deepspeed_custom_scheduler = \u001b[38;5;28mself\u001b[39m.is_deepspeed_enabled \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\n", + "\u001b[32m 3476\u001b[39m \u001b[38;5;28mself\u001b[39m.lr_scheduler, DeepSpeedSchedulerWrapper\n", + "\u001b[32m 3477\u001b[39m )\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:966\u001b[39m, in \u001b[36msave\u001b[39m\u001b[34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[39m\n", + "\u001b[32m 963\u001b[39m f = os.fspath(f)\n", + "\u001b[32m 965\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m _use_new_zipfile_serialization:\n", + "\u001b[32m--> \u001b[39m\u001b[32m966\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m _open_zipfile_writer(f) \u001b[38;5;28;01mas\u001b[39;00m opened_zipfile:\n", + "\u001b[32m 967\u001b[39m _save(\n", + "\u001b[32m 968\u001b[39m obj,\n", + "\u001b[32m 969\u001b[39m opened_zipfile,\n", + "\u001b[32m (...)\u001b[39m\u001b[32m 972\u001b[39m _disable_byteorder_record,\n", + "\u001b[32m 973\u001b[39m )\n", + "\u001b[32m 974\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n", + "\n", + "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:798\u001b[39m, in \u001b[36m_open_zipfile_writer_file.__exit__\u001b[39m\u001b[34m(self, *args)\u001b[39m\n", + "\u001b[32m 797\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m__exit__\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args) -> \u001b[38;5;28;01mNone\u001b[39;00m:\n", + "\u001b[32m--> \u001b[39m\u001b[32m798\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mfile_like\u001b[49m\u001b[43m.\u001b[49m\u001b[43mwrite_end_of_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[32m 799\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.file_stream \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", + "\u001b[32m 800\u001b[39m \u001b[38;5;28mself\u001b[39m.file_stream.close()\n", + "\n", + "\u001b[31mRuntimeError\u001b[39m: [enforce fail at inline_container.cc:664] . unexpected pos 11819474816 vs 11819474712" + ] + } + ], + "source": [ + "import os \n", + "os.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"\n", + "trainer = JEPOTrainer(\n", + " model = model,\n", + " processing_class = tokenizer,\n", + " reward_funcs = [\n", + " jepo_strict_format_reward_func,\n", + " ],\n", + " cot_func = fabricate_ground_truth_completion,\n", + " args = training_args,\n", + " train_dataset = dataset,\n", + ")\n", + "trainer.train()\n", + "if torch.cuda.is_available():\n", + " memory_used = torch.cuda.memory_allocated() / 1024**3\n", + " memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**3\n", + " print(f\"💾 GPU Memory Used: {memory_used:.2f} GB / {memory_total:.1f} GB ({memory_used/memory_total*100:.1f}%)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6a99f973", + "metadata": {}, + "outputs": [], + "source": [ + "f.flush()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b94d80e9", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np \n", + " \n", + "# Given array \n", + "data = np.array([0, 0, 0, -1]) \n", + " \n", + "# Compute the mean and standard deviation \n", + "mean = np.mean(data) \n", + "std = np.std(data) \n", + " \n", + "# Compute the z-score normalization: (x - mean) / std \n", + "z_scores = (data - mean) / (std + 1e-4) # Adding a small value to avoid division by zero \n", + " \n", + "print(\"Original array:\", data) \n", + "print(\"Mean:\", mean) \n", + "print(\"Standard Deviation:\", std) \n", + "print(\"Z-score normalized array:\", z_scores) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "99d3c4d1", + "metadata": {}, + "outputs": [], + "source": [ + "text = tokenizer.apply_chat_template([\n", + " {\"role\" : \"user\", \"content\" : \"Calculate pi.\"},\n", + "], tokenize = False, add_generation_prompt = True)\n", + "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") \n", + "\n", + "# Assuming you have already loaded your model and tokenizer. \n", + "# model = LlamaForCausalLM.from_pretrained(\"path/to/your/model\") \n", + "# tokenizer = ... \n", + " \n", + "input_ids = tokenizer(text, return_tensors=\"pt\").input_ids.to(device) \n", + " \n", + "# Use the standard generate method. \n", + "outputs = model.generate( \n", + " input_ids, \n", + " max_new_tokens=1024, \n", + " do_sample=True, \n", + " temperature=1.0, \n", + " top_p=0.95, \n", + ") \n", + " \n", + "# Decode the generated tokens. \n", + "output_text = tokenizer.decode(outputs[0], skip_special_tokens=True) \n", + "print(output_text) \n" + ] + }, + { + "cell_type": "markdown", + "id": "0a3300ec", + "metadata": {}, + "source": [ + "And now with the LoRA we just trained with GRPO - we first save the LoRA first!" + ] + }, + { + "cell_type": "markdown", + "id": "ff6b0e7f", + "metadata": {}, + "source": [ + "Now we load the LoRA and test:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb1be866", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "id": "2c6e5c11", + "metadata": {}, + "source": [ + "Our reasoning model is much better - it's not always correct, since we only trained it for an hour or so - it'll be better if we extend the sequence length and train for longer!" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/trl/__init__.py b/trl/__init__.py index 897b4fef917..ea833182195 100644 --- a/trl/__init__.py +++ b/trl/__init__.py @@ -49,6 +49,8 @@ "DPOTrainer", "GRPOConfig", "GRPOTrainer", + "JEPOConfig", + "JEPOTrainer", "KTOConfig", "KTOTrainer", "LogCompletionsCallback", diff --git a/trl/trainer/__init__.py b/trl/trainer/__init__.py index f24ea415072..9a05d4ea71a 100644 --- a/trl/trainer/__init__.py +++ b/trl/trainer/__init__.py @@ -29,6 +29,8 @@ "dpo_trainer": ["DPOTrainer"], "grpo_config": ["GRPOConfig"], "grpo_trainer": ["GRPOTrainer"], + "jepo_config": ["JEPOConfig"], + "jepo_trainer": ["JEPOTrainer"], "kto_config": ["KTOConfig"], "kto_trainer": ["KTOTrainer"], "model_config": ["ModelConfig"], diff --git a/trl/trainer/jepo_config.py b/trl/trainer/jepo_config.py new file mode 100644 index 00000000000..2ab33ca72e0 --- /dev/null +++ b/trl/trainer/jepo_config.py @@ -0,0 +1,678 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional, Union + +from transformers import TrainingArguments + + +@dataclass +class JEPOConfig(TrainingArguments): + r""" + Configuration class for the [`JEPOTrainer`], which serves as a variation of GRPO for unverifiable RL training. + JEPO [https://arxiv.org/pdf/2503.19618] + + This class includes only the parameters that are specific to JEPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]`, *optional*): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`JEPOTrainer`] is provided as a string. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left. + num_generations (`int` or `None`, *optional*, defaults to `8`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + + > Parameters that control generation + + generation_batch_size: (`int`, *optional*): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * steps_per_generation`. In other words, there is one + generation batch processed per optimization step. Mutually exclusive with `steps_per_generation`. + steps_per_generation: (`int`, *optional*): + Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`. Mutually exclusive + with `generation_batch_size`. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int`, *optional*): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + use_transformers_paged (`bool`, *optional*, defaults to `False`): + Whether to use the `transformers` paged implementation for generation. If set to `True`, the `transformers` + paged implementation will be used for generation instead of the default padded implementation. This + parameter is only effective when `use_vllm` is set to `False`. + cache_implementation (`str`, *optional*): + Implementation of the cache method for faster generation when `use_vllm` is set to `False`. + generation_kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments to pass to [`~transformers.GenerationConfig`] (if using transformers) or + `SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the + generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that conflict + with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation to use for vLLM. Must be one of `"transformers"` or `"vllm"`. `"transformers"`: Use + the `transformers` backend for model implementation. `"vllm"`: Use the `vllm` library for model + implementation. + vllm_guided_decoding_regex (`str`, *optional*): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step and woken + for weight sync and generation. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.0`): + KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving + training speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + delta (`float`, *optional*): + Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard + GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in + the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). + epsilon_high (`float`, *optional*): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + importance_sampling_level (`str`, *optional*, defaults to `"token"`): + Controls whether importance sampling ratios are computed at the `"token"` or `"sequence"` level. `"token"` + keeps the raw per-token log-probability ratios (one weight per token). `"sequence"` averages the + log-probability ratios across valid tokens to produce a single ratio per sequence. The [GSPO + paper](https://huggingface.co/papers/2507.18071) shows that sequence-level sampling often yields more + stable training and better alignment with sequence-level rewards. + reward_weights (`list[float]`, *optional*): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + scale_rewards (`str` or `bool`, *optional*, defaults to `"group"`): + Specifies the scaling strategy for rewards. Supported values are: + + - `True` or `"group"` (default): rewards are scaled by the standard deviation within each group, ensuring + unit variance within a group. + - `"batch"`: rewards are scaled by the standard deviation across the entire batch, as recommended in the + [PPO Lite paper](https://huggingface.co/papers/2508.08221). + - `False` or `"none"`: no scaling is applied. The [Dr. GRPO + paper](https://huggingface.co/papers/2503.20783) recommends not scaling rewards, as scaling by the + standard deviation introduces a question-level difficulty bias. + loss_type (`str`, *optional*, defaults to `"norm_jepo"`): + Specifies the loss formulation to use. Supported values are: + + - `"norm_jepo"`: Aggregates token-level losses by normalizing over sequence length. + - `"unnorm_jepo"`: Aggregates token-level losses without normalization. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + top_entropy_quantile (`float`, *optional*, defaults to `1.0`): + ρ parameter from [Beyond the 80/20 Rule](https://huggingface.co/papers/2506.01939). Keeps in the policy + loss term only the top-ρ quantile of tokens by entropy of the probability distribution at each sequence + position, improving results. Range: `[0.0-1.0]`. A value of `0.0` masks all but the highest entropy token; + `1.0` keeps all tokens. The paper recommends a value of `0.2`. If used with + `mask_truncated_completions=True`, only tokens from non-truncated completions are considered. + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use the Liger GRPO loss. + vllm_importance_sampling_correction (`bool`, *optional*, defaults to `True`): + Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and recomputed + logprobs. [Your Efficient RL Framework Secretly Brings You Off-Policy RL + Training](https://fengyao.notion.site/off-policy-rl) highlights that using a separate generation framework + (such as vLLM) can introduce off-policy effects due to subtle implementation differences between generation + and training backends. TIS is proposed as a remedy for this issue. + vllm_importance_sampling_cap (`float`, *optional*, defaults to `2.0`): + Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the importance + sampling ratio, improving training stability. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`. + num_completions_to_print (`int`, *optional*): + Number of completions to print with `rich`. If `None`, all completions are logged. + wandb_log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, all prompts + are logged. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + gradient_checkpointing: bool = field( + default=True, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` " + "argument of the `GRPOTrainer` is provided as a string." + }, + ) + disable_dropout: bool = field( + default=False, + metadata={ + "help": "Whether to disable dropout in the model. This is useful for training with a reference model, as " + "it prevents the model from generating different logprobs for the same input." + }, + ) + + # Parameters that control the data preprocessing + # The default value remove_unused_columns is overwritten from the parent class, because in GRPO we usually rely on + # additional columns to compute the reward + remove_unused_columns: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to only keep the column 'prompt' in the dataset. If you use a custom reward function " + "that requires any column other than 'prompts' and 'completions', you should keep this to `False`." + }, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left." + }, + ) + num_generations: Optional[int] = field( + default=8, + metadata={ + "help": "Number of generations to sample. The effective batch size (num_processes * per_device_batch_size " + "* gradient_accumulation_steps) must be evenly divisible by this value." + }, + ) + max_completion_length: Optional[int] = field( + default=256, + metadata={"help": "Maximum length of the generated completion."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option " + "is not compatible with vLLM generation." + }, + ) + shuffle_dataset: Optional[bool] = field( + default=True, + metadata={"help": "Whether to shuffle the training dataset."}, + ) + + # Parameters that control generation + generation_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size to use for generation. If `None`, it defaults to the effective training batch size: " + "`per_device_train_batch_size * num_processes * steps_per_generation`." + }, + ) + steps_per_generation: Optional[int] = field( + default=None, + metadata={"help": "Number of steps per generation. If `None`, it defaults to `gradient_accumulation_steps`."}, + ) + temperature: float = field( + default=1.0, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=1.0, + metadata={ + "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. " + "Set to 1.0 to consider all tokens." + }, + ) + top_k: Optional[int] = field( + default=None, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + min_p: Optional[float] = field( + default=None, + metadata={ + "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It " + "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range." + }, + ) + generation_kwargs: Optional[dict] = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or " + "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the " + "generation behavior, such as setting `suppress_tokens`, `num_beams`, etc. If it contains keys that " + "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them." + }, + ) + repetition_penalty: float = field( + default=1.0, + metadata={ + "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated " + "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model " + "to repeat tokens." + }, + ) + use_transformers_paged: bool = field( + default=False, + metadata={ + "help": "Whether to use the `transformers` paged implementation for generation. If set to `True`, the " + "`transformers` paged implementation will be used for generation instead of the default padded " + "implementation. This parameter is only effective when `use_vllm` is set to `False`." + }, + ) + cache_implementation: Optional[str] = field( + default=None, + metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."}, + ) + + # Parameters that control generation acceleration powered by vLLM + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for " + "generation instead of the default model.generate(). Requires `vllm` to be installed." + }, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `'server'` or " + "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure " + "a TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same " + "process and share the training GPUs. This avoids the need for a separate server but may cause resource " + "contention with training." + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={ + "help": "Model implementation to use for vLLM. Must be one of `transformers` or `vllm`. `transformers`: " + "Use the `transformers` backend for model implementation. `vllm`: Use the `vllm` library for " + "model implementation." + }, + ) + vllm_enable_sleep_mode: bool = field( + default=False, + metadata={ + "help": "Whether to enable sleep mode for vLLM. If `True`, vLLM will sleep during the optimization step " + "and woken for weight sync and generation." + }, + ) + vllm_guided_decoding_regex: Optional[str] = field( + default=None, + metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, + ) + + # Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_base_url: Optional[str] = field( + default=None, + metadata={ + "help": "Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` " + "and `vllm_server_port` are ignored." + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": "Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_port: int = field( + default=8000, + metadata={"help": "Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={ + "help": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up " + "after the timeout, a `ConnectionError` is raised." + }, + ) + + # Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + vllm_gpu_memory_utilization: float = field( + default=0.3, + metadata={ + "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag." + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={ + "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_tensor_parallel_size` flag." + }, + ) + + # Parameters that control the training + beta: float = field( + default=0.0, + metadata={ + "help": "KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and " + "improving training speed." + }, + ) + supervised_loss_weight: float = field( + default=0.1, + metadata={ + "help": "Weight for the supervised loss. If greater than `0.0`, the loss optimized will be a weighted sum " + "between the JEPO loss and a supervised loss (cross-entropy between the model predictions and the " + "reference completions). " + }, + ) + num_iterations: int = field( + default=1, + metadata={"help": "Number of iterations per batch (denoted as μ in the algorithm)."}, + ) + epsilon: float = field( + default=0.2, + metadata={"help": "Epsilon value for clipping."}, + ) + delta: Optional[float] = field( + default=None, + metadata={ + "help": "Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` " + "(default), standard GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This " + "method is introduced in the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291)." + }, + ) + epsilon_high: Optional[float] = field( + default=None, + metadata={ + "help": "Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the " + "lower-bound specified in argument `epsilon`. Paper DAPO recommends `0.28`." + }, + ) + importance_sampling_level: str = field( + default="token", + metadata={ + "help": "Controls whether importance sampling ratios are computed at the `'token'` or `'sequence'` level. " + "`'token'` keeps the raw per-token log-probability ratios (one weight per token). `'sequence'` averages " + "the log-probability ratios across valid tokens to produce a single ratio per sequence. The GSPO paper " + "shows that sequence-level sampling often yields more stable training and better alignment with " + "sequence-level rewards." + }, + ) + reward_weights: Optional[list[float]] = field( + default=None, + metadata={ + "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all " + "rewards are weighted equally with weight `1.0`." + }, + ) + scale_rewards: str = field( + default="group", + metadata={ + "help": "Specifies the scaling strategy for rewards. Supported values are: " + "`True` or `group'` (default): rewards are scaled by the standard deviation within each group, ensuring " + "unit variance within a group. " + "`'batch'`: rewards are scaled by the standard deviation across the entire batch, as recommended in the " + "PPO Lite paper. " + "`False` or `'none'`: no scaling is applied. The Dr. GRPO paper recommends not scaling rewards, as " + "scaling by the standard deviation introduces a question-level difficulty bias." + }, + ) + loss_type: str = field( + default="norm_jepo", + metadata={ + "help": "Specifies the loss formulation to use. Supported values are 'norm_jepo' and 'unnorm_jepo'. " + "'norm_jepo': Aggregates token-level losses by normalizing over sequence length. " + "'unnorm_jepo': Aggregates token-level losses without normalization. " + }, + ) + mask_truncated_completions: bool = field( + default=False, + metadata={ + "help": "When enabled, truncated completions are excluded from the loss calculation, preventing them from " + "being incorrectly penalized and introducing noise during training. According to the DAPO paper, this is " + "a good practice for training stability." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + top_entropy_quantile: float = field( + default=1.0, + metadata={ + "help": "ρ parameter from Beyond the 80/20 Rule. Keeps in the policy loss term only the top-ρ quantile of " + "tokens by entropy of the probability distribution at each sequence position, improving results. Range: " + "[0.0-1.0]. A value of `0.0` masks all but the highest entropy token; `1.0` keeps all tokens. The paper " + "recommends a value of `0.2`. If used with `mask_truncated_completions=True`, only tokens from " + "non-truncated completions are considered." + }, + ) + use_liger_loss: bool = field( + default=False, + metadata={"help": "Whether to use the Liger GRPO loss."}, + ) + vllm_importance_sampling_correction: bool = field( + default=True, + metadata={ + "help": "Whether to apply Truncated Importance Sampling (TIS) between vLLM completion logprobs and " + "recomputed logprobs. Your Efficient RL Framework Secretly Brings You Off-Policy RL " + "Training highlights that using a separate generation framework (such as vLLM) can introduce off-policy " + "effects due to subtle implementation differences between generation and training backends. TIS is " + "proposed as a remedy for this issue." + }, + ) + vllm_importance_sampling_cap: float = field( + default=2.0, + metadata={ + "help": "Truncation parameter C for Truncated Importance Sampling (TIS). This sets an upper bound on the " + "importance sampling ratio, improving training stability." + }, + ) + + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + num_completions_to_print: Optional[int] = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_log_unique_prompts: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, " + "all prompts are logged." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + self.scale_rewards = {True: "group", False: "none"}.get(self.scale_rewards, self.scale_rewards) + + num_processes = self.world_size + # The current default effective batch size + if self.generation_batch_size is None and self.steps_per_generation is None: + self.steps_per_generation = self.gradient_accumulation_steps + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + elif self.generation_batch_size is not None and self.steps_per_generation is None: + # Just ensure the value is divisible by the global batch size + if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size " + f"({self.per_device_train_batch_size * num_processes})." + ) + self.steps_per_generation = self.generation_batch_size // ( + self.per_device_train_batch_size * num_processes + ) + elif self.generation_batch_size is None and self.steps_per_generation is not None: + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + else: + raise ValueError( + "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time" + ) + + if self.do_eval and self.eval_strategy != "no": + # Just ensure the value is divisible by the global batch size + if (self.per_device_eval_batch_size * num_processes) % self.num_generations != 0: + raise ValueError( + f"The global eval batch size ({self.per_device_eval_batch_size} * {num_processes}) must be " + f"divisible by num_generations ({self.num_generations})." + ) + + # The generation batch must contain full prompt groups (no partials), so it must be divisible by + # num_generations. + if self.generation_batch_size % self.num_generations != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations " + f"({self.num_generations})." + ) + + if self.num_generations < 2: + raise ValueError( + "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided " + f"{self.num_generations}, which is less than the minimum required." + ) + + if self.delta is not None and self.use_liger_loss: + raise ValueError("Liger loss does not support two-sided GRPO loss yet.") diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py new file mode 100644 index 00000000000..537ecf3f50d --- /dev/null +++ b/trl/trainer/jepo_trainer.py @@ -0,0 +1,1838 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import textwrap +from collections import defaultdict, deque +from contextlib import nullcontext +from functools import partial +from pathlib import Path +from typing import Any, Callable, Optional, Union + +import datasets +import torch +import torch.utils.data +import transformers +from accelerate import logging +from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed +from datasets import Dataset, IterableDataset +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.utils.data import DataLoader, Sampler +from transformers import ( + AutoConfig, + AutoModelForSequenceClassification, + AutoProcessor, + AutoTokenizer, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + is_wandb_available, +) +from transformers.trainer_utils import seed_worker +from transformers.utils import is_datasets_available, is_flash_attn_2_available, is_peft_available, is_rich_available + +from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template, prepare_multimodal_messages +from ..extras.profiling import profiling_context, profiling_decorator +from ..extras.vllm_client import VLLMClient +from ..import_utils import is_liger_kernel_available, is_vllm_available +from ..models import prepare_deepspeed, prepare_fsdp, prepare_peft_model, unwrap_model_for_generation +from ..models.utils import _ForwardRedirection +from .base_trainer import BaseTrainer +from .callbacks import SyncRefModelCallback +from .jepo_config import JEPOConfig +from .utils import ( + RepeatSampler, + disable_dropout_in_model, + ensure_master_addr_port, + entropy_from_logits, + identity, + nanmax, + nanmin, + nanstd, + pad, + print_prompt_completions_sample, + selective_log_softmax, + shuffle_sequence_dict, + split_pixel_values_by_grid, + split_tensor_dict, + unsplit_pixel_values_by_grid, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + +if is_wandb_available(): + import wandb + + +logger = logging.get_logger(__name__) + +# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of +# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. +RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]] + +# What we call a CoT function is a callable that takes completion text and answer and returns CoT text and CoT+answer. +CoTFunc = Callable[[str, str], [str, str]] + +class JEPOTrainer(BaseTrainer): + """ + Trainer for the Jensen’s Evidence lower bound Policy Optimization (JEPO) method. This algorithm was initially proposed in the + paper: Beyond Verifiable Rewards: Scaling Reinforcement Learning for Language Models to Unverifiable Data [https://arxiv.org/pdf/2503.19618] +. + + Example: + + see the jepo_training notebook for a complete example of how to use the JEPOTrainer. + ``` + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + reward_funcs (`Union[RewardFunc, list[RewardFunc]]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can also return `None` when the reward is not applicable to those samples. This is useful + for multi-task training where different reward functions apply to different types of samples. When a + reward function returns `None` for a sample, that reward function is excluded from the reward + calculation for that sample. For more details, see [Using a custom reward + function](#using-a-custom-reward-function). + + The trainer's state is also passed to the reward function. The trainer's state is an instance of + [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the + reward function's signature. + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + cot_func (`CoTFunc`): + CoT function used to extact chain-of-thought (CoT) from the generated completion and generate CoT-only and CoT+Answer formatted text. + + args ([`JEPOConfig`], *optional*): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A + padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token, + `tokenizer.eos_token` will be used as the default. + reward_processing_classes ([`~transformers.PreTrainedTokenizerBase`] or `list[PreTrainedTokenizerBase]`, *optional*): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + """ + + _tag_names = ["trl", "jepo"] + _name = "JEPO" + _paper = { + "title": "Beyond Verifiable Rewards: Scaling Reinforcement Learning for Language Models to Unverifiable Data", + } + + def __init__( + self, + model: Union[str, PreTrainedModel], + reward_funcs: Union[RewardFunc, list[RewardFunc]], + cot_func: CoTFunc, + args: Optional[JEPOConfig] = None, + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None, + processing_class: Optional[Union[PreTrainedTokenizerBase, ProcessorMixin]] = None, + reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + peft_config: Optional["PeftConfig"] = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else model.config._name_or_path + model_name = model_name.split("/")[-1] + args = GRPOConfig(f"{model_name}-JEPO") + + # Models + # Trained model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + dtype = model_init_kwargs.get("dtype") + if isinstance(dtype, torch.dtype) or dtype == "auto" or dtype is None: + pass # dtype is already a torch.dtype or "auto" or None + elif isinstance(dtype, str): # it's a str, but not "auto" + dtype = getattr(torch, dtype) + model_init_kwargs["dtype"] = dtype + else: + raise ValueError( + "Invalid `dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + model = architecture.from_pretrained(model_id, **model_init_kwargs) + else: + model_id = model.config._name_or_path + if args.model_init_kwargs is not None: + logger.warning( + "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + + # Some models (SmolVLM/Idefics3) don't support `logits_to_keep` argument and error out if we pass it + # Inspect the forward method before we wrap the model with PEFT + self.model_kwarg_keys = ( + inspect.signature(model.forward).parameters.keys() + if not hasattr(model, "get_base_model") + else inspect.signature(model.get_base_model().forward).parameters.keys() + ) + + if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)): + model = prepare_peft_model(model, peft_config, args) + + # Processing class + if processing_class is None: + processing_class = AutoProcessor.from_pretrained(model.config._name_or_path, truncation_side="left") + + # Handle pad token for processors or tokenizers + if isinstance(processing_class, ProcessorMixin): + tokenizer = processing_class.tokenizer + elif isinstance(processing_class, PreTrainedTokenizerBase): + tokenizer = processing_class + else: + raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`") + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + self.pad_token = tokenizer.pad_token + self.pad_token_id = tokenizer.pad_token_id + self.eos_token_id = tokenizer.eos_token_id + self.cot_func = cot_func + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(reward_funcs[i].config._name_or_path.split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError( + f"The number of reward processing classes ({len(reward_processing_classes)}) must match the number of " + f"reward functions ({len(reward_funcs)})." + ) + + for i, (reward_processing_class, reward_func) in enumerate(zip(reward_processing_classes, reward_funcs)): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained(reward_func.config._name_or_path) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + + self.reward_processing_classes = reward_processing_classes + self.supervised_loss_weight = args.supervised_loss_weight + # Training arguments + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper + self.num_generations = args.num_generations # = G in the GRPO paper + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_transformers_paged = args.use_transformers_paged + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.vllm_importance_sampling_correction = args.vllm_importance_sampling_correction + self.vllm_importance_sampling_cap = args.vllm_importance_sampling_cap + self.use_liger_loss = args.use_liger_loss + self.loss_type = args.loss_type + self.scale_rewards = args.scale_rewards + self.importance_sampling_level = args.importance_sampling_level + self.mask_truncated_completions = args.mask_truncated_completions + self.top_entropy_quantile = args.top_entropy_quantile + if self.use_liger_loss and self.top_entropy_quantile < 1.0: + raise NotImplementedError( + "Liger Kernels don't currently support masking token positions based on entropy." + ) + if self.use_liger_loss and not self.importance_sampling_level == "token": + raise NotImplementedError( + "Liger Kernels currently only support token-level importance sampling. Please set" + "`importance_sampling_level` to 'token'." + ) + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + print(f"configuration: {args}") + if ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in GRPOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations (forward + backward passes), including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in GRPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning: + # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To + # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True. + # This acts as a flag to indicate that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in GRPO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + # In Trainer, `training_step` scales the loss by `gradient_accumulation_steps` only if `compute_loss_func` + # is None. For DAPO, loss scaling instead depends on the total number of completions tokens across the + # global accumulated batch. To control scaling ourselves, we must disable Trainer’s built-in scaling. The + # simplest (though a bit hacky) way is to set `compute_loss_func` to any non-None value, which bypasses + # that behavior without rewriting `training_step`. + compute_loss_func="non-None value to disable scaling", + ) + + # Reference model + self.beta = args.beta + self.supervised_loss_weight = args.supervised_loss_weight + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + config = AutoConfig.from_pretrained(model_id) + architecture = getattr(transformers, config.architectures[0]) + self.ref_model = architecture.from_pretrained(model_id, **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # Keep logs sized to the generation batch to record only outputs from the latest model update. + self._logs = { + "images": deque(maxlen=args.generation_batch_size), + "prompt": deque(maxlen=args.generation_batch_size), + "completion": deque(maxlen=args.generation_batch_size), + "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)), + "advantages": deque(maxlen=args.generation_batch_size), + } + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install trl[vllm]` to use it." + ) + + if self.vllm_mode == "server": + if self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator(device=torch.cuda.current_device()) + + elif self.vllm_mode == "colocate": + # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have + # the same number of ranks + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks. + # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size)) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + # vLLM requires the environment variables to be set for distributed training. + os.environ["RANK"] = str(self.accelerator.process_index) + os.environ["LOCAL_RANK"] = str(self.accelerator.local_process_index) + os.environ["WORLD_SIZE"] = str(self.accelerator.num_processes) + # Ensure distributed rendezvous variables are set without colliding across concurrent runs + ensure_master_addr_port() + + if self.max_prompt_length is not None and self.max_completion_length is not None: + max_model_len = self.max_prompt_length + self.max_completion_length + else: + max_model_len = None + self.llm = LLM( + model=model.name_or_path, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + max_num_seqs=self.args.per_device_train_batch_size + * self.vllm_tensor_parallel_size + * self.args.steps_per_generation, + max_model_len=max_model_len, + distributed_executor_backend="external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + seed=self.accelerator.process_index // self.vllm_tensor_parallel_size, + # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768) - thinking there's not enough memory + max_num_batched_tokens=4096, + model_impl=self.args.vllm_model_impl, + enable_sleep_mode=self.args.vllm_enable_sleep_mode, + # Important so temperature scaling/logit tweaking affects the TIS log probs + logprobs_mode="processed_logprobs", + ) + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=1) + else: + raise ValueError(f"vllm_mode must be either 'server' or 'colocate', got '{self.vllm_mode}'.") + + # vLLM specific sampling arguments + self.guided_decoding_regex = args.vllm_guided_decoding_regex + + self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": tokenizer.pad_token_id, + "bos_token_id": tokenizer.bos_token_id, + "eos_token_id": tokenizer.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt", "image", "images"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. As a result, some parts of the method aren't relevant to GRPO, but we keep them to stay one line + # apart from the super method, ensuring easier maintenance in the future. + def get_train_dataloader(self): + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size * self.args.steps_per_generation, # < this is the change + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = partial( + seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index + ) + + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. The first row shows the first sampled batch, the + # second row shows the second sampled batch, and so on. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations, + seed=self.args.seed, + ) + + @profiling_decorator + def _get_last_hidden_state( + self, + unwrapped_model, + input_ids, + attention_mask, + logits_to_keep, + pixel_values=None, + image_grid_thw=None, + pixel_attention_mask=None, + image_sizes=None, + ): + if is_peft_model(unwrapped_model): + unwrapped_model = unwrapped_model.base_model.model + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask} + + # For Qwen models: + if image_grid_thw is not None and pixel_values is not None: + model_inputs["image_grid_thw"] = image_grid_thw + # For Gemma, SmolVLM2, LLaVa-Next etc.: + if pixel_values is not None: + model_inputs["pixel_values"] = pixel_values + # For SmolVLM2 + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask + # For LLaVa-Next + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + last_hidden_state = unwrapped_model.model(**model_inputs).last_hidden_state + # Exclude the last value: it corresponds to the next token pred + last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + return last_hidden_state + + def get_high_entropy_mask(self, entropies: torch.Tensor, mask: torch.Tensor, threshold: float) -> torch.Tensor: + """ + Returns a binary mask identifying tokens whose entropy exceeds a given quantile threshold. + + Args: + entropies (`torch.Tensor`): + Tensor of shape (batch_size, seq_len) with per-token entropy values. + mask (`torch.Tensor`): + Binary mask of the same shape as `entropies`, where `1` indicates valid tokens and `0` padding. + threshold (`float`): + Quantile threshold between `0.0` and `1.0` to select high-entropy tokens. + + Returns: + `torch.Tensor`: + Boolean mask of shape (batch_size, seq_len), where `True` indicates tokens with entropy >= threshold + and `False` otherwise. + """ + local = entropies[mask.bool()].float() + + # Use a negative pad_value as a sentinel because entropy values are always >= 0. + # This guarantees that the sentinel cannot collide with any real entropy value. + pad_value = -1e9 + + # Pad across processes so that every rank has the same tensor length + padded = self.accelerator.pad_across_processes(local, dim=0, pad_index=pad_value) + gathered = self.accelerator.gather(padded) + + # Drop sentinel values (safe because no entropy can be negative) + gathered = gathered[gathered != pad_value] + + if gathered.numel() == 0: + return torch.zeros_like(entropies, dtype=torch.bool) + + entropy_threshold = torch.quantile(gathered, threshold) + masked_entropies = entropies * mask.float() + entropy_mask = masked_entropies >= entropy_threshold + return entropy_mask & mask.bool() # ensure padding tokens are always masked out + + @profiling_decorator + def _get_per_token_logps_and_entropies( + self, + model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=None, + compute_entropy=False, + pixel_values=None, + image_grid_thw=None, + num_images=None, + pixel_attention_mask=None, + image_sizes=None, + token_type_ids=None, + ) -> dict[str, Optional[torch.Tensor]]: + """Compute log-probs and (optionally) entropies for each token.""" + batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak + all_logps = [] + all_entropies = [] + for start in range(0, input_ids.size(0), batch_size): + input_ids_batch = input_ids[start : start + batch_size] + attention_mask_batch = attention_mask[start : start + batch_size] + + # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't) + model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch} + if image_grid_thw is not None and pixel_values is not None: + rows_per_image = image_grid_thw.prod(dim=-1) + rows_per_sample = torch.split(rows_per_image, num_images) + rows_per_sample = torch.stack([s.sum() for s in rows_per_sample]) + cum_rows = torch.cat([torch.tensor([0], device=rows_per_sample.device), rows_per_sample.cumsum(0)]) + row_start, row_end = cum_rows[start].item(), cum_rows[start + batch_size].item() + model_inputs["pixel_values"] = pixel_values[row_start:row_end] + cum_imgs = torch.tensor([0] + num_images).cumsum(0) + img_start, img_end = cum_imgs[start], cum_imgs[start + batch_size] + model_inputs["image_grid_thw"] = image_grid_thw[img_start:img_end] + elif pixel_values is not None: + model_inputs["pixel_values"] = pixel_values[start : start + batch_size] + if pixel_attention_mask is not None: + model_inputs["pixel_attention_mask"] = pixel_attention_mask[start : start + batch_size] + if image_sizes is not None: + model_inputs["image_sizes"] = image_sizes[start : start + batch_size] + if token_type_ids is not None: + model_inputs["token_type_ids"] = token_type_ids[start : start + batch_size] + + # Only add logits_to_keep if the model supports it + if "logits_to_keep" in self.model_kwarg_keys: + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + model_inputs["logits_to_keep"] = logits_to_keep + 1 + + model_inputs["use_cache"] = False # only used in generation; set False to suppress warnings + + logits = model(**model_inputs).logits + # Exclude the last value: it corresponds to the next token pred + logits = logits[:, :-1, :] # (B, L-1, H) + # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op. + logits = logits[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + # Divide logits by sampling temperature. + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + logits = logits / self.temperature + + completion_ids = input_ids_batch[:, -logits_to_keep:] + logps = selective_log_softmax(logits, completion_ids) # compute logprobs + all_logps.append(logps) + + if compute_entropy: + with torch.no_grad(): + entropies = entropy_from_logits(logits) + all_entropies.append(entropies) + + logps = torch.cat(all_logps, dim=0) + entropies = torch.cat(all_entropies, dim=0) if compute_entropy else None + return logps, entropies + + def _fix_param_name_to_vllm(self, name, extra_prefixes: Optional[list[str]] = None): + extra_prefixes = extra_prefixes or [] + prefixes = ["_checkpoint_wrapped_module."] + extra_prefixes + for prefix in prefixes: + name = name.replace(prefix, "") + return name + + def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + # For FSDP1, we need to recurse into children and also use summon_full_params + if visited is None: + visited = set() + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp1_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + full_name = self._fix_param_name_to_vllm(full_name, extra_prefixes=["_fsdp_wrapped_module."]) + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + def _sync_fsdp2_params_to_vllm(self, module: nn.Module): + # For FSDP2, module.state_dict() already covers all parameters, so no need for recursion + for name, param in module.state_dict().items(): + if param.is_cpu: + param = param.to(torch.device("cuda")) + param = param.full_tensor() + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param)]) + + @profiling_decorator + def _move_model_to_vllm(self): + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + # TODO: does this work with FSDP? + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm( + self.model + ) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."]) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None) + fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1 + if fsdp_version == 1: + self._sync_fsdp1_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP + elif fsdp_version == 2: + self._sync_fsdp2_params_to_vllm(self.model) + else: + for name, param in self.model.named_parameters(): + name = self._fix_param_name_to_vllm(name) + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.llm.reset_prefix_cache() + + @profiling_decorator + def _prepare_inputs( + self, generation_batch: dict[str, Union[torch.Tensor, Any]] + ) -> dict[str, Union[torch.Tensor, Any]]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = split_pixel_values_by_grid(generation_batch) + generation_batch = shuffle_sequence_dict(generation_batch) + generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation) + self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches] + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + self._step += 1 + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + # This allows for dynamic reward shaping based on training progress. + reward_kwargs["trainer_state"] = self.state + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names) + ): + with profiling_context(self, reward_func_name): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions)] + texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages] + else: + texts = [p + c for p, c in zip(prompts, completions)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = { + key: value[nan_row_idx] for key, value in reward_kwargs.items() if key != "trainer_state" + } + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + logger.warning( + f"All reward functions returned None for the following kwargs:\n{row_reward_kwargs}\n" + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _generate_single_turn(self, prompts: list[any], images: Optional[list]): + device = self.accelerator.device + + # If the prompts are conversational and the inputs contain images, we need to convert the prompts from + # [{"role": "user", "content": "What color is the sky?"}] to + # [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What color is the sky?"}]}] + kwargs = {} + if images is not None: + kwargs = {"images": images} + for prompt, image_list in zip(prompts, images): + if isinstance(prompt, list): # i.e., when using conversational data + prepare_multimodal_messages(prompt, num_images=len(image_list)) + + prompts_text = [ + maybe_apply_chat_template({"prompt": prompt}, self.processing_class)["prompt"] for prompt in prompts + ] + + if images is not None: + prompt_inputs = self.processing_class(text=prompts_text, padding=True, return_tensors="pt", **kwargs) + prompt_inputs = super()._prepare_inputs(prompt_inputs) + forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]} + else: + forward_kwargs = {} + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + if self.vllm_mode == "colocate" and self.args.vllm_enable_sleep_mode: + # wake up colocated vLLM instances if needed + torch.cuda.empty_cache() # required to avoid OOM in some cases + self.llm.wake_up() + + # First, update the vLLM weights if needed + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process + if self.vllm_mode == "server": + all_prompts_text = gather_object(prompts_text) + if images is not None: + all_images = gather_object(images) + + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts_text[:: self.num_generations] + + if images is not None: + ordered_set_of_images = all_images[:: self.num_generations] + else: + ordered_set_of_images = None + + with profiling_context(self, "vLLM.generate"): + output = self.vllm_client.generate( + prompts=ordered_set_of_prompts, + images=ordered_set_of_images, + n=self.num_generations, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=-1 if self.top_k is None else self.top_k, + min_p=0.0 if self.min_p is None else self.min_p, + max_tokens=self.max_completion_length, + truncate_prompt_tokens=self.max_prompt_length, + guided_decoding_regex=self.guided_decoding_regex, + generation_kwargs=self.args.generation_kwargs, + ) + payload = (output["prompt_ids"], output["completion_ids"], output["logprobs"]) + else: + payload = None + + # Broadcast the completions from the main process to all processes, ensuring each process receives its corresponding slice. + obj_list = [payload] + broadcast_object_list(obj_list, from_process=0) + all_prompt_ids, all_completion_ids, all_logprobs = obj_list[0] + + # At this point, we only get 1 copy of each prompt, so we need to repeat them num_generations times + all_prompt_ids = [ids for ids in all_prompt_ids for _ in range(self.num_generations)] + + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + prompt_ids = all_prompt_ids[process_slice] + completion_ids = all_completion_ids[process_slice] + logprobs = all_logprobs[process_slice] + + # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts + elif self.vllm_mode == "colocate": + if self.guided_decoding_regex: + guided_decoding = GuidedDecodingParams(regex=self.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": 1, # vLLM on each GPU generates only 1 in colocate mode + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "truncate_prompt_tokens": self.max_prompt_length, + "guided_decoding": guided_decoding, + "logprobs": 0, # only return the logprob of the generated token + } + if self.args.generation_kwargs is not None: + generation_kwargs.update(self.args.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + if self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts_text) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group) + all_prompts_text = [p for sublist in gathered_prompts for p in sublist] + + if images is not None: + gathered_images = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_images, images, group=self.tp_group) + all_images = [img for sublist in gathered_images for img in sublist] + else: + all_images = None + else: + all_prompts_text = prompts_text + all_images = images + + if images is not None and all_images: + vllm_inputs = [] + for prompt, image_list in zip(all_prompts_text, all_images): + vllm_inputs.append({"prompt": prompt, "multi_modal_data": {"image": image_list}}) + + else: + vllm_inputs = all_prompts_text + + with profiling_context(self, "vLLM.generate"): + all_outputs = self.llm.generate(vllm_inputs, sampling_params=sampling_params, use_tqdm=False) + + all_prompt_ids = [output.prompt_token_ids for output in all_outputs] + all_completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + all_logprobs = [ + [next(iter(lp.values())).logprob for lp in output.logprobs] + for outputs in all_outputs + for output in outputs.outputs + ] + + if self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + prompt_ids = all_prompt_ids[tp_slice] + completion_ids = all_completion_ids[tp_slice] + logprobs = all_logprobs[tp_slice] + else: + prompt_ids = all_prompt_ids + completion_ids = all_completion_ids + logprobs = all_logprobs + + if self.args.vllm_enable_sleep_mode: + self.llm.sleep(level=1) + + elif self.use_transformers_paged: + # Re-process inputs for paged generation if needed + # Note: images are already validated and preprocessed above + paged_prompt_inputs = self.processing_class(text=prompts_text, **kwargs) + previous_attn = self.model_wrapped.config._attn_implementation + + if is_flash_attn_2_available(): + self.model_wrapped.config._attn_implementation = "paged_attention" + else: + self.model_wrapped.config._attn_implementation = "sdpa_paged" + with ( + profiling_context(self, "transformers.generate_batch"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + # Cast to the appropriate dtype based on training configuration + if self.args.bf16: + unwrapped_model.to(torch.bfloat16) + elif self.args.fp16: + unwrapped_model.to(torch.float16) + with torch.inference_mode(): + all_outputs = unwrapped_model.generate_batch( + paged_prompt_inputs.input_ids, generation_config=self.generation_config, progress_bar=False + ) + unwrapped_model.train() # restore training mode, as generate_batch forces eval mode + completion_ids = [output.generated_tokens for output in all_outputs.values()] + prompt_ids = paged_prompt_inputs.input_ids + # Restore the original attention implementation, training mode + self.model_wrapped.config._attn_implementation = previous_attn + logprobs = None # not used in this case + + else: + # Regular generation path + generate_inputs = self.processing_class( + text=prompts_text, + return_tensors="pt", + padding=True, + padding_side="left", + max_length=self.max_prompt_length, + truncation=True, + add_special_tokens=False, + **kwargs, + ) + generate_inputs = super()._prepare_inputs(generate_inputs) + + with ( + profiling_context(self, "transformers.generate"), + unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model, + torch.no_grad(), + FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(), + ): + prompt_completion_ids = unwrapped_model.generate( + **generate_inputs, generation_config=self.generation_config, disable_compile=True + ) + # Compute prompt length and extract completion ids + prompt_ids, prompt_mask = generate_inputs["input_ids"], generate_inputs["attention_mask"] + prompt_length = prompt_ids.size(1) + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + prompt_ids = [p[m].tolist() for p, m in zip(prompt_ids, prompt_mask.bool())] + completion_ids = [c[m].tolist() for c, m in zip(completion_ids, completion_mask.bool())] + logprobs = None # not used in this case + + return prompt_ids, completion_ids, logprobs, forward_kwargs + + def _generate(self, prompts: list[str], images: Optional[list]): + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompt_ids, completion_ids, logprobs, forward_kwargs = self._generate_single_turn(prompts, images) + + # Get completion length per sequence, used for logging + prompt_lengths = torch.tensor([len(ids) for ids in prompt_ids], device=device) + completion_lengths = torch.tensor([len(ids) for ids in completion_ids], device=device) + agg_prompt_lengths = self.accelerator.gather(prompt_lengths) + agg_completion_lengths = self.accelerator.gather(completion_lengths) + total_prompt_tokens = agg_prompt_lengths.sum() + total_completion_tokens = agg_completion_lengths.sum() # = num_items_in_batch, required for the DAPO loss + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += (total_prompt_tokens + total_completion_tokens).item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids], device=device) + agg_is_truncated = self.accelerator.gather(is_truncated) + self._metrics[mode]["completions/clipped_ratio"].append(agg_is_truncated.float().mean().item()) + term_completion_lengths = agg_completion_lengths[~agg_is_truncated] + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + return prompt_ids, completion_ids, total_completion_tokens, logprobs, forward_kwargs + + + def _compute_jepo_advantages(self, jepo_rewards, jepo_applicable_mask, num_generations, device="cpu"): + """ + Compute JEPO advantages for a rewards tensor by groups. + + Args: + jepo_rewards (Tensor): A 1D tensor of rewards. + jepo_applicable_mask (Tensor): A boolean 1D tensor of the same length as jepo_rewards. + True indicates that the sample is applicable. + num_generations (int): Number of samples (generations) per group. + device (str): Device specifier ("cpu" or "cuda"). + + Returns: + Tensor: A tensor of advantages, with the same shape as jepo_rewards. + """ + jepo_advantages = torch.zeros_like(jepo_rewards) + num_groups = jepo_rewards.shape[0] // num_generations + + for i in range(num_groups): + start_idx = i * num_generations + end_idx = (i + 1) * num_generations + group_slice = slice(start_idx, end_idx) + + # Process only if at least one sample in the group is applicable. + if jepo_applicable_mask[group_slice].any(): + # Count the applicable (True) samples and convert to float. + applicable_count = jepo_applicable_mask[group_slice].sum().float() + # Sum rewards for the group. + group_reward_sum = jepo_rewards[group_slice][jepo_applicable_mask[group_slice]].sum() + # Compute the mean over the applicable samples. + mean_reward = group_reward_sum / applicable_count + + # List to store computed advantages for later normalization. + group_advantages_list = [] + + for j in range(start_idx, end_idx): + if jepo_applicable_mask[j]: + if applicable_count > 1: + # Compute a leave-one-out "variance" (mean of the other rewards). + variance = (group_reward_sum - jepo_rewards[j]) / (applicable_count - 1) + # JEPO advantage encourages samples that have high and consistent rewards. + advantage = torch.log(mean_reward) - torch.log(variance) + else: + # If only one sample is applicable, use log(mean_reward). + advantage = torch.log(mean_reward) + jepo_advantages[j] = advantage + group_advantages_list.append(advantage) + + # Normalize the advantages within the group by their standard deviation. + # If there is only one applicable sample, we use a standard deviation of 1. + if len(group_advantages_list) > 1: + group_std = torch.std(torch.stack(group_advantages_list)) + else: + group_std = torch.tensor(1.0, device=device) + #print(f"Group {i}: mean_reward={mean_reward:.4f}, std={group_std:.4f}, applicable_count={applicable_count.item()} group_reward_sum={group_reward_sum:.4f}") + for j in range(start_idx, end_idx): + if jepo_applicable_mask[j]: + jepo_advantages[j] = torch.clamp( jepo_advantages[j]/(group_std + 1e-4) , -1, 1) + + return jepo_advantages + + def _generate_and_score_completions( + self, inputs: list[dict[str, Union[torch.Tensor, Any]]] + ) -> dict[str, Union[torch.Tensor, Any]]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + + if "images" in inputs[0]: + images = [example.get("images") for example in inputs] + elif "image" in inputs[0]: + images = [[example.get("image")] if example.get("image") is not None else None for example in inputs] + else: + images = None + # Transformers requires at least one image in the batch, otherwise it throws an error + if images is not None and all(img_list == [] for img_list in images): + images = None + + ( + prompt_ids_list, + completion_ids_list, + num_items_in_batch, + sampling_per_token_logps_list, + forward_kwargs, + ) = self._generate(prompts, images) + + # Convert lists of token IDs to padded tensors + prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_ids_list] + prompt_mask = [torch.ones_like(ids, dtype=torch.long) for ids in prompt_ids] + prompt_ids = pad(prompt_ids, padding_value=self.pad_token_id, padding_side="left") + prompt_mask = pad(prompt_mask, padding_value=0, padding_side="left") + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids_list] + completion_mask = [torch.ones_like(ids, dtype=torch.long) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right") + completion_mask = pad(completion_mask, padding_value=0, padding_side="right") + if sampling_per_token_logps_list is not None: + sampling_per_token_logps = [torch.tensor(logps, device=device) for logps in sampling_per_token_logps_list] + sampling_per_token_logps = pad(sampling_per_token_logps, padding_value=0.0, padding_side="right") + else: + sampling_per_token_logps = None + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + eos_and_pad = [self.eos_token_id, self.pad_token_id] + is_truncated = torch.tensor([ids[-1] not in eos_and_pad for ids in completion_ids_list], device=device) + completion_mask = completion_mask * (~is_truncated).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) # (B, P+C) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + # If token_type_ids are used, extend them with zeros for the completion part + if "token_type_ids" in forward_kwargs: + token_type_ids = forward_kwargs["token_type_ids"] + forward_kwargs["token_type_ids"] = torch.cat( + [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1 + ) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + num_images = [len(img_list) for img_list in images] if images is not None else None + + + # Decode + prompts_text = self.processing_class.batch_decode(prompt_ids, skip_special_tokens=True) + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + # Compute grouped-wise rewards + mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1) + + # Normalize the rewards to compute the advantages + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + advantages = rewards - mean_grouped_rewards + + # compute jepo advantages which is based on p(anwer|question, CoT), CoT is from the sampled completions_text + # Note the advantagees would be 0 if the sample does not following the format required by the format reward function; + # get clone of rewards + jepo_rewards = rewards.clone() + logits_to_keep_answer = rewards.clone() + with torch.no_grad(): + + + #calclate the CtO logits and entropy, which is used to calculate the JEPO rewards. The JEPO reward is calculated as the log probability of the completion given the prompt and the sampled CoT, minus the log probability of the completion given only the prompt (i.e., without the CoT). The intuition is that if the sampled CoT is helpful for generating the completion, then the log probability of the completion given the prompt and CoT should be higher than the log probability of the completion given only the prompt, resulting in a positive JEPO reward. + #synthesize the completion with sampled CoT from completions_text and answer + + def synthesize_completion(prompt, completion, answer): + correct_format, cot_text, cot_answer_text = self.cot_func(completion[0]['content'], answer) + + #print(f"prompt: {prompt}, completion: {completion}, answer: {answer}, correct_format: {correct_format}, cot_text: {cot_text}, cot_answer_text: {cot_answer_text}") + return (correct_format, maybe_apply_chat_template({"prompt": prompt, 'completion': [{"role": "assistant", "content": cot_text}]}, self.processing_class)['completion'], + maybe_apply_chat_template({"prompt": prompt, 'completion': [{"role": "assistant", "content": cot_answer_text}]}, self.processing_class)['completion']) + + cot_texts = [] + cot_answer_texts = [] + jepo_applicable_mask = torch.zeros(len(prompts), dtype=torch.bool, device=device) + for i, (prompt, completion, answer) in enumerate(zip(prompts, completions, [x['answer'] for x in inputs])): + correct_format, cot_text, cot_answer_text = synthesize_completion(prompt, completion, answer) + cot_texts.append(cot_text) + cot_answer_texts.append(cot_answer_text) + + if correct_format: + jepo_applicable_mask[i] = True # mark this sample as applicable for JEPO reward + + cot_text_tokens = self.processing_class( + text=cot_texts, + padding=True, + padding_side="right", + max_length=self.max_completion_length, + return_tensors="pt", + truncation=True, + add_special_tokens=False) + cot_text_tokens = super()._prepare_inputs(cot_text_tokens) + + cot_answer_text_tokens = self.processing_class( + text=cot_answer_texts, + padding=True, + padding_side="right", + max_length=self.max_completion_length, + return_tensors="pt", + truncation=True, + add_special_tokens=False) + cot_answer_text_tokens = super()._prepare_inputs(cot_answer_text_tokens) + + + fabricated_completions_ids, fabricated_completions_mask = cot_answer_text_tokens["input_ids"], cot_answer_text_tokens["attention_mask"] + + input_ids = torch.cat([prompt_ids, fabricated_completions_ids], dim=1) + attention_mask = torch.cat([prompt_mask, fabricated_completions_mask], dim=1) + + + # we only need to compute the logits for the answer (A| Q,C) + logits_to_keep = fabricated_completions_ids.size(1) # we only need to compute the logits for the fabricated completion tokens, which is the CoT answer part + + #print(f"fabricated_completions_ids shape: {fabricated_completions_ids.shape}, fabricated_completions_mask shape: {fabricated_completions_mask.shape},logits_to_keep: {logits_to_keep}") + + + + # Compute the per_token_logps and the entropy at each position in the completion + per_token_logps, entropies = self._get_per_token_logps_and_entropies( + self.model_wrapped, + input_ids, + attention_mask, + logits_to_keep, + batch_size, + compute_entropy=True, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + + + padded_cot_attention_mask = cot_text_tokens["attention_mask"] + if cot_text_tokens["attention_mask"].shape[1] != fabricated_completions_mask.shape[1]: + # if the cot_text_tokens attention mask is shorter than the fabricated_completions_mask, we need to pad it to the same length + padding_length = fabricated_completions_mask.shape[1] - cot_text_tokens["attention_mask"].shape[1] + padded_cot_attention_mask = torch.cat([cot_text_tokens["attention_mask"], torch.zeros((cot_text_tokens["attention_mask"].shape[0], padding_length), dtype=torch.long, device=device)], dim=1) + # there is EOT_TOKEN at the end of the cot_text and cot_answer_text, using answer_mask = fabricated_completions_mask - padded_cot_attention_mask decode shows "answer>\n0\n", "<" is missing because of the EOT_TOKEN + #fix: there is EOT_TOKEN at the end of the cot_text and cot_answer_text ship by one + answer_mask = torch.cat([fabricated_completions_mask[:, 1:], torch.zeros((fabricated_completions_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) - torch.cat([padded_cot_attention_mask[:, 1:], torch.zeros((padded_cot_attention_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) + #check the answer_mask is correct by decoding the tokens corresponding to the answer_mask and see if it matches with the cot_answer_texts using padding mask to ignore the padding tokens + answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) + #print(f"answer_texts_from_mask: {[answer_text.strip('!') for answer_text in answer_texts_from_mask]}") + + #compute answer log probabilities (i.e., p(A|Q,C)) and use it as JEPO rewards by per_token_logits over answer mask, then sum over the answer tokens to get the total log probability of the answer given the question and CoT. The intuition is that if the sampled CoT is helpful for generating the completion, then the log probability of the answer given the prompt and CoT should be high, resulting in a high JEPO reward + #IMPORTANT: be careful of the exp need to be done first before sum, otherwise it would be sum of log probabilities which is not what we want (since prob =0 still has the advantage of 1) + jepo_rewards = (torch.exp(per_token_logps) * answer_mask).sum(dim=1) # zero out the log probabilities for the non-answer tokens + print(f"jepo_rewards: {jepo_rewards}") + #print the prob of the answer: + #print(f"answer_mask: {answer_mask.sum(dim=1)}, fabricated_completions_mask: {fabricated_completions_mask.sum(dim=1)}, padded_cot_attention_mask: {padded_cot_attention_mask.sum(dim=1)}") + + jepo_rewards[jepo_applicable_mask == False] = 0.0 # if not applicable for JEPO, set reward to 0 + #calculate the mean of jepo rewards for each group of generations, based on the jepo_applicable_mask + jepo_advantages = self._compute_jepo_advantages(jepo_rewards, jepo_applicable_mask, self.num_generations, device=device) + #print(f"jepo_rewards: {jepo_rewards}, jepo_advantages: {jepo_advantages}") + + #print(f"jepo reward related shapes: input_ids {input_ids.shape},attention_mask {attention_mask.shape}, logits_to_keep {logits_to_keep}, fabricated_completions_ids {fabricated_completions_ids.shape}, fabricated_completions_mask {fabricated_completions_mask.shape}, per_token_logps {per_token_logps.shape}, entropies {entropies.shape}, jepo_rewards {jepo_rewards.shape}") + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.ref_model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps, _ = self._get_per_token_logps_and_entropies( + self.model, + input_ids, + attention_mask, + logits_to_keep, + batch_size=batch_size, + num_images=num_images, + **forward_kwargs, # may contain pixel_values, image_grid_thw, pixel_attention_mask and image_sizes + ) + else: + ref_per_token_logps = None + + + if self.scale_rewards in ["group", "none"]: + # If self.scale_rewards = "none", we'll still log group level std + std_rewards = rewards.view(-1, self.num_generations).std(dim=1) + std_rewards = std_rewards.repeat_interleave(self.num_generations, dim=0) + elif self.scale_rewards == "batch": + # Compute global std + std_rewards = rewards.std().expand_as(rewards) + else: + raise ValueError( + f"Invalid value for scale_rewards: {self.scale_rewards}. Must be one of 'batch', 'group', or 'none'." + ) + + is_std_zero = torch.isclose(std_rewards, torch.zeros_like(std_rewards)) + if self.scale_rewards != "none": + advantages = advantages / (std_rewards + 1e-4) + + print(f"format advantages shape: {advantages.shape}, {advantages} , jepo_advantages shape: {jepo_advantages.shape}, {jepo_advantages}") + #sum formatted rewards and jepo rewards to get the final rewards, then calculate the final advantages + advantages = advantages + jepo_advantages + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_func_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_func_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._logs["prompt"].extend(gather_object(prompts_text)) + self._logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._logs["advantages"].extend(all_process_advantages.tolist()) + + if images is not None: + self._logs["images"].extend(gather_object(images)) + + output = { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "fabricated_completions_ids": fabricated_completions_ids, + "fabricated_completions_mask": fabricated_completions_mask, + "answer_mask": answer_mask, + "advantages": advantages, + "num_items_in_batch": num_items_in_batch, + } + if self.use_vllm and self.vllm_importance_sampling_correction: + output["importance_sampling_ratio"] = importance_sampling_ratio + if ref_per_token_logps is not None: + output["ref_per_token_logps"] = ref_per_token_logps + if "pixel_values" in forward_kwargs: + output["pixel_values"] = forward_kwargs["pixel_values"] + if "image_grid_thw" in forward_kwargs: + output["image_grid_thw"] = forward_kwargs["image_grid_thw"] + if "pixel_attention_mask" in forward_kwargs: + output["pixel_attention_mask"] = forward_kwargs["pixel_attention_mask"] + if "image_sizes" in forward_kwargs: + output["image_sizes"] = forward_kwargs["image_sizes"] + if "token_type_ids" in forward_kwargs: + output["token_type_ids"] = forward_kwargs["token_type_ids"] + if images is not None: + output["num_images"] = num_images + return output + + + + @profiling_decorator + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + + return self._compute_loss(model, inputs) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + device = self.accelerator.device + + fabricated_completions_ids, fabricated_completions_mask = inputs["fabricated_completions_ids"], inputs["fabricated_completions_mask"] + answer_mask = inputs["answer_mask"] + input_ids = torch.cat([prompt_ids, fabricated_completions_ids], dim=1) + attention_mask = torch.cat([prompt_mask, fabricated_completions_mask], dim=1) + + logits_to_keep = fabricated_completions_ids.size(1) # we only need to compute the logits for the fabricated completion tokens + + per_token_logps, entropies = self._get_per_token_logps_and_entropies( + model, + input_ids, + attention_mask, + logits_to_keep, + compute_entropy=True, + pixel_values=inputs.get("pixel_values"), + image_grid_thw=inputs.get("image_grid_thw"), + num_images=inputs.get("num_images"), + pixel_attention_mask=inputs.get("pixel_attention_mask"), + image_sizes=inputs.get("image_sizes"), + token_type_ids=inputs.get("token_type_ids"), + ) + + if self.top_entropy_quantile < 1.0: + entropy_mask = self.get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile) + else: + entropy_mask = None + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + + # Compute the loss + advantages = inputs["advantages"] + + + # compute p(C|Q), fabricated_completions has the EOT_TOKEN at the end and answer_mask have that removed, need to align them first before applying the answer_mask + cot_mask = torch.cat([fabricated_completions_mask[:, 1:], torch.zeros((fabricated_completions_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) - answer_mask # this is the mask for the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. + answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) + #print(f"answer_texts_from_mask: {[answer_text.strip('!') for answer_text in answer_texts_from_mask]}") + + cot_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * cot_mask, skip_special_tokens=True) + #print(f"cot_texts_from_mask: {[cot_text.strip('!') for cot_text in cot_texts_from_mask]}") + cot_token_logps = per_token_logps * cot_mask # this is the log probability of the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. + + advantages = advantages.unsqueeze(1) # (B, 1) + + per_token_adv = (advantages * cot_token_logps) # (B, T), we only use the log probabilities of the CoT part to calculate the loss, which is equivalent to using the JEPO reward as the advantage. The intuition is that if the sampled CoT is helpful for generating the completion, then the log probability of the CoT part should be high, resulting in a high JEPO reward and thus a high advantage, which encourages the model to generate completions that are more likely under the sampled CoT. + + if self.loss_type == 'unnorm_jepo': + loss = - per_token_adv.sum(dim=1).sum()/ (self.num_generations) # sum over the tokens and average over the batch + else: + loss = - (per_token_adv.sum(dim=1)/cot_mask.sum(dim=1)).sum()/ (self.num_generations) # sum over the tokens and average over the batch + print(f"advantages loss: {loss}") + + #compute P(A|Q,C) and use it as a supervised loss to encourage the model to generate answers + if self.supervised_loss_weight > 0.0: + # Add a supervised loss to encourage the model to generate the reference answer. We use the answer_mask to calculate the supervised loss only on the answer part, which is p(A|Q,C). The intuition is that we want to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. + answer_token_logps = torch.exp(per_token_logps ) * answer_mask + if self.loss_type == 'unnorm_jepo': + supervised_loss = answer_token_logps.sum(dim=1) # sum over the answer tokens to get the total log probability of the answer given the question and CoT, which is p(A|Q,C), we will use it as a supervised loss to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. + else: + supervised_loss= answer_token_logps.sum(dim=1)/answer_mask.sum(dim=1) + # sum per number_generations to get the total log probability of the answer given the question and CoT, which is p(A|Q,C), we will use it as a supervised loss to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. + supervised_loss = torch.log(supervised_loss.view(-1, self.num_generations).sum(dim=1)/self.num_generations + 1e-8) #average over the generations and take log to get the supervised loss for each sample in the batch, which is log(p(A|Q,C)), we will use it as a supervised loss to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. + #print(f"supervised_loss before weighting: {supervised_loss}") + loss = loss - self.supervised_loss_weight * supervised_loss.sum() # average over the batch + #print(f"supervised_loss loss: {loss}, supervised_loss_weight: {self.supervised_loss_weight}") + if self.beta != 0.0: + if self.loss_type == 'unnorm_jepo': + loss = loss + self.beta * (per_token_kl * fabricated_completions_mask).sum(dim=1).sum() / (self.num_generations) # sum over the tokens and average over the batch + else: + loss = loss + self.beta * (((per_token_kl * fabricated_completions_mask).sum(dim=1)) / fabricated_completions_mask.sum(dim=1)).sum() / (self.num_generations) # average over the batch and the tokens, note that we only calculate KL for the completion part, which is indicated by the fabricated_completions_mask. The intuition is that we want to regularize the model to not deviate too much from the reference model, which can help stabilize training and prevent divergence. + print(f" kl loss: {loss}, beta: {self.beta}") + + + # Log the metrics + mode = "train" if self.model.training else "eval" + + completion_token_count = attention_mask.sum().clamp(min=1.0) + + def masked_batch_mean(x): + if x.shape[1] == 1: # when importance_sampling_level == "sequence" + return x.mean() + else: + return (x * fabricated_completions_mask).sum() / completion_token_count + + if self.beta != 0.0: + mean_kl = masked_batch_mean(per_token_kl) + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + mean_entropy = masked_batch_mean(entropies) + self._metrics[mode]["entropy"].append(self.accelerator.gather(mean_entropy).nanmean().item()) + + return loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._logs["prompt"], + self._logs["completion"], + self._logs["rewards"], + self._logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + ) + + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + import pandas as pd + + table = { + "step": [str(self.state.global_step)] * len(self._logs["prompt"]), + "prompt": self._logs["prompt"], + "completion": self._logs["completion"], + **self._logs["rewards"], + "advantage": self._logs["advantages"], + } + + if self._logs["images"]: + table["images"] = [] + for image_list in self._logs["images"]: + # Convert images to wandb Image objects for proper visualization + table["images"].append([wandb.Image(image) for image in image_list]) + + df = pd.DataFrame(table) + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + wandb.log({"completions": wandb.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) \ No newline at end of file From d3c0e91f0bd86bfecc780202bf8910b137311620 Mon Sep 17 00:00:00 2001 From: haijunz Date: Tue, 31 Mar 2026 00:30:26 +0000 Subject: [PATCH 02/12] minor notebook update --- examples/notebooks/jepo_math.ipynb | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/examples/notebooks/jepo_math.ipynb b/examples/notebooks/jepo_math.ipynb index 2c2cc640dcf..b53e69af1d4 100644 --- a/examples/notebooks/jepo_math.ipynb +++ b/examples/notebooks/jepo_math.ipynb @@ -27,20 +27,10 @@ }, { "cell_type": "markdown", - "id": "c0e36a31", + "id": "e9c15369", "metadata": {}, "source": [ - "Long-Context GRPO for reinforcement learning — train stably at massive sequence lengths. Fine-tune models with up to 7x more context length efficiently. [Read Blog](https://unsloth.ai/docs/new/grpo-long-context)\n", - "\n", - "3× faster training with optimized sequence packing — higher throughput with no quality loss.[Read Blog](https://unsloth.ai/docs/new/3x-faster-training-packing)\n", - "\n", - "500k context-length fine-tuning — push long-context models further with memory-efficient training. [Read Blog](https://unsloth.ai/docs/new/500k-context-length-fine-tuning)\n", - "\n", - "Introducing FP8 precision training for faster RL inference. [Read Blog](https://docs.unsloth.ai/new/fp8-reinforcement-learning).\n", - "\n", - "Unsloth's [Docker image](https://hub.docker.com/r/unsloth/unsloth) is here! Start training with no setup & environment issues. [Read our Guide](https://docs.unsloth.ai/new/how-to-train-llms-with-unsloth-and-docker).\n", - "\n", - "Visit our docs for all our [model uploads](https://docs.unsloth.ai/get-started/all-our-models) and [notebooks](https://docs.unsloth.ai/get-started/unsloth-notebooks).\n" + "This notebook launch JEPO trainer" ] }, { @@ -58,8 +48,8 @@ "metadata": {}, "outputs": [], "source": [ - "% !pip install antlr4-python3-runtime==4.11\n", - "% !pip install sympy" + "%pip install antlr4-python3-runtime==4.11\n", + "%pip install sympy" ] }, { From 7d4d49d58fb78ec1385fa5a3b5079f3a56d58c84 Mon Sep 17 00:00:00 2001 From: haijunz Date: Thu, 2 Apr 2026 21:14:59 +0000 Subject: [PATCH 03/12] fix typo --- trl/trainer/jepo_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trl/trainer/jepo_config.py b/trl/trainer/jepo_config.py index 2ab33ca72e0..f44c96282ff 100644 --- a/trl/trainer/jepo_config.py +++ b/trl/trainer/jepo_config.py @@ -670,9 +670,9 @@ def __post_init__(self): if self.num_generations < 2: raise ValueError( - "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided " + "JEPO requires at least 2 generations per prompt to calculate the advantages. You provided " f"{self.num_generations}, which is less than the minimum required." ) if self.delta is not None and self.use_liger_loss: - raise ValueError("Liger loss does not support two-sided GRPO loss yet.") + raise ValueError("Liger loss does not support two-sided JEPO loss yet.") From b167bb029c7ba92989ad36679cd454059998b59f Mon Sep 17 00:00:00 2001 From: haijunz Date: Thu, 2 Apr 2026 21:28:32 +0000 Subject: [PATCH 04/12] add paper section --- trl/trainer/jepo_trainer.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index 537ecf3f50d..6afb7f02d78 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -188,8 +188,16 @@ class JEPOTrainer(BaseTrainer): _name = "JEPO" _paper = { "title": "Beyond Verifiable Rewards: Scaling Reinforcement Learning for Language Models to Unverifiable Data", + "id": "2503.19618", + "citation": textwrap.dedent("""\ + @article{tang2025beyond, + title = {{Beyond Verifiable Rewards: Scaling Reinforcement Learning for Language Models to Unverifiable Data}}, + author = {Yunhao Tang, Sid Wang, Lovish Madaan, Rémi Munos}, + year = 2025, + eprint = {arXiv:2503.19618}, + }"""), } - + def __init__( self, model: Union[str, PreTrainedModel], From c0353b98326ed6a6dd157411b4e9fb2da3301387 Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 15:43:25 +0000 Subject: [PATCH 05/12] fix per comments --- examples/notebooks/jepo_math.ipynb | 4262 +--------------------------- trl/trainer/jepo_config.py | 4 +- trl/trainer/jepo_trainer.py | 4 +- 3 files changed, 6 insertions(+), 4264 deletions(-) diff --git a/examples/notebooks/jepo_math.ipynb b/examples/notebooks/jepo_math.ipynb index b53e69af1d4..0a35cca5889 100644 --- a/examples/notebooks/jepo_math.ipynb +++ b/examples/notebooks/jepo_math.ipynb @@ -153,34 +153,12 @@ "from sympy.parsing.latex import parse_latex" ] }, - { - "cell_type": "markdown", - "id": "7499eede", - "metadata": {}, - "source": [ - "### Data Prep\n", - "\n", - "\n", - "We directly leverage [@willccbb](https://gist.github.com/willccbb/4676755236bb08cab5f4e54a0475d6fb) for data prep and all reward functions. You are free to create your own!" - ] - }, { "cell_type": "code", "execution_count": null, "id": "286b2149", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "<>:176: SyntaxWarning: invalid escape sequence '\\s'\n", - "<>:176: SyntaxWarning: invalid escape sequence '\\s'\n", - "/tmp/ipykernel_298016/1337997636.py:176: SyntaxWarning: invalid escape sequence '\\s'\n", - " pattern = r\".*?\\s*.*?\"\n" - ] - } - ], + "outputs": [], "source": [ "import re\n", "from datasets import load_dataset, Dataset\n", @@ -493,4176 +471,7 @@ "execution_count": null, "id": "a2c38211", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "The model is already on multiple devices. Skipping the move to device specified in `args`.\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "74c900196d824cb894d5e0d762d3d023", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Loading checkpoint shards: 0%| | 0/4 [00:00\n", - " \n", - " \n", - " [ 1001/26838 8:20:03 < 215:32:47, 0.03 it/s, Epoch 0.11/3]\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
StepTraining Loss
1-2.696400
2-1.145700
30.154800
4-0.101000
5-2.916200
6-2.978900
7-0.907400
8-1.637400
9-2.700700
100.092300
111.201500
12-0.537100
130.034600
14-2.414600
15-0.043900
160.080900
17-0.042500
18-0.018000
19-2.685000
200.041600
21-1.283200
221.842100
23-0.000400
24-1.998300
250.009200
260.040200
270.222000
28-0.111700
29-1.555400
30-0.157000
31-3.583600
32-2.846800
330.028900
340.141400
35-1.083500
36-2.772700
370.076800
38-0.272500
390.096000
40-1.073400
41-2.284600
420.105200
43-1.744000
440.131700
45-0.055400
460.171100
470.159000
48-2.024400
490.132700
500.056500
51-0.033500
521.284400
53-3.392100
54-1.962100
550.001500
56-1.752600
570.095100
580.031000
592.168400
600.078700
61-2.939600
620.105900
630.045700
64-2.696300
65-0.040300
660.023100
670.074700
680.016600
69-3.299400
700.035400
710.013400
72-0.001400
73-0.068200
740.083000
75-0.148900
760.052700
77-1.545800
780.083000
790.201100
800.173700
81-0.013000
82-0.535900
83-3.639400
840.018000
850.170700
86-0.170800
87-2.359300
88-0.057700
89-2.182300
90-0.009400
91-1.786000
920.008900
93-3.353100
94-2.462600
95-0.011500
96-2.816400
97-0.239400
980.040900
990.269400
1000.154800
101-2.085000
102-0.028700
1030.046900
104-0.054500
105-2.440600
106-1.495400
107-1.604600
1080.857100
1090.149200
110-1.078300
1110.028000
112-2.280100
113-1.708400
114-3.363500
1150.008400
1160.129200
117-0.002400
118-0.005300
1190.167600
120-1.369800
121-2.255600
1220.053900
1230.030000
1240.008200
1250.046000
1260.100100
127-1.509400
1280.207500
1290.136000
1300.067700
131-2.168700
132-0.022000
1330.109200
134-0.174800
1350.058500
1360.142200
137-1.690100
138-0.047700
139-2.075400
1400.087000
141-0.164100
142-0.546000
143-2.661400
144-0.019000
1450.012900
146-1.288500
1470.110000
148-1.148600
1490.216800
150-3.189000
151-0.387100
152-0.352500
1530.127600
1540.144700
1551.044200
156-2.586500
157-1.734000
1580.000200
1591.842100
160-1.612600
161-3.101100
162-0.026000
163-0.053900
164-0.897100
165-0.337200
166-0.139300
167-2.152100
1680.023300
169-0.004600
170-1.478800
171-1.004800
1720.030600
173-0.522100
174-0.180800
1750.179200
1760.002100
1770.009400
1780.173700
1790.170600
180-1.489900
1810.153300
182-3.308000
183-2.884200
184-2.303500
185-0.199500
186-0.892600
1870.200300
1880.041100
1890.051100
190-0.757600
191-0.015300
192-0.979100
193-0.514800
194-1.412500
195-3.489100
1960.273000
1970.079400
198-0.008600
199-2.918100
200-3.104900
201-0.955900
2020.049200
2030.190600
204-2.765600
205-3.075500
2060.071800
207-0.189700
2080.027300
209-0.884100
210-0.361600
211-3.159400
2120.061000
213-0.173100
2140.003200
2150.255900
216-0.098300
217-0.189500
2180.257100
219-0.061100
2200.102600
221-0.003600
2220.047700
2230.113300
2242.669400
2250.018700
226-1.641800
227-1.839500
2280.022600
229-0.051500
2300.109800
2310.012200
2320.153100
233-0.039800
234-0.437000
235-0.356900
2360.039200
237-1.119600
2380.049900
239-2.238600
2400.165500
241-2.869100
242-2.747400
243-0.084200
244-1.311900
245-0.248800
246-0.627600
247-0.018600
2482.025400
249-0.539000
250-1.410400
2510.087800
2520.061100
253-0.066800
2540.195500
2550.063300
256-3.353400
257-0.038000
258-0.714600
259-1.651500
2600.185300
261-2.932000
262-1.172300
263-0.000900
264-1.822100
2650.066000
2660.352700
267-0.067700
2680.038000
269-0.068100
270-1.991100
2710.212700
272-3.411800
273-2.866400
274-1.717400
275-0.165800
2760.014800
2770.058500
2780.207700
279-0.022300
280-0.131200
281-0.001800
2820.067500
283-0.258000
284-0.543100
285-0.074200
286-0.187300
287-0.191800
288-0.054700
289-0.065600
2900.506400
291-4.027500
2920.001200
2930.062300
294-3.115100
295-0.000100
296-0.019600
297-0.734600
298-0.354300
299-0.061500
3000.198800
301-0.031900
302-0.047500
303-0.052000
3040.008500
305-0.202900
306-1.625800
3070.492700
3080.004100
3090.010100
310-1.769100
3110.058500
312-2.569000
313-2.964900
314-0.798800
315-1.676200
3160.151300
317-0.011700
318-0.035400
319-3.011100
320-1.834900
321-0.644400
322-0.252500
3230.018600
3240.002400
325-1.122300
326-0.315100
327-1.108700
328-0.377300
3290.012400
3300.300600
3310.034100
332-2.868700
333-3.103900
334-2.834900
335-1.575900
336-2.621900
337-1.696800
338-1.001700
339-3.230000
3400.399000
341-3.393800
342-1.858100
343-0.895600
344-2.674200
345-2.199000
346-0.006100
3470.037600
3480.181000
349-0.508200
350-0.669400
351-0.034600
352-0.059300
353-0.086700
354-0.397800
3551.347200
356-0.161300
357-2.574900
3580.171600
3590.016900
360-2.935600
361-0.021200
362-2.929300
3630.089000
3640.042000
3650.090500
366-1.443800
367-0.027600
368-0.082100
3690.369700
370-2.665200
3710.005100
372-0.120000
373-2.020000
3740.069700
375-2.660800
376-0.003400
377-0.075200
3780.026600
379-0.197800
380-0.831300
3810.082400
382-0.027500
3830.040600
384-2.064900
3850.123400
386-0.121600
3870.051100
388-0.153400
3890.092600
3900.046500
391-1.645200
392-3.311800
393-0.062500
394-0.003600
3950.056900
396-2.811000
397-0.029200
398-2.608600
399-2.073000
4000.033500
401-1.276100
402-2.336900
403-1.199800
404-0.082500
405-2.603000
406-0.028300
407-0.006600
408-2.890600
4090.016200
4100.006100
4110.033400
4120.049000
4130.219600
414-0.122300
415-0.398200
4160.058200
417-2.069400
418-0.010900
419-3.705000
4200.107600
421-0.014300
4220.100100
423-2.435100
424-2.679500
425-0.072100
426-0.970800
427-0.005000
4280.010900
4290.043000
430-0.160300
4310.118900
432-0.059500
4330.188600
4340.109500
435-0.405300
436-0.039200
437-2.995800
4380.075600
439-0.194700
440-0.026500
4410.056200
442-3.214300
443-0.088800
444-0.073500
445-0.248700
446-2.727400
447-2.691000
448-2.862000
4490.158600
450-0.001800
451-0.001500
452-2.259800
453-0.168800
454-0.000600
455-2.378000
456-1.412400
457-0.054900
458-1.806800
459-0.027200
460-0.046100
4610.047400
4620.149200
463-0.239800
464-0.031000
465-2.821500
466-2.595100
467-2.667300
4680.018800
469-1.645500
470-2.263300
4710.015800
4720.114400
473-2.255300
474-0.001500
475-0.033200
4760.050400
4770.001000
478-0.055200
479-0.014500
480-2.134100
481-0.121400
482-2.483900
4830.123800
484-0.006300
485-1.525000
486-0.010600
487-1.265800
488-1.639700
489-1.603700
490-2.841900
4910.052300
492-0.868700
493-2.416200
494-1.778200
4950.028800
496-0.063100
497-1.550500
498-0.000500
499-1.381700
5000.014000
501-0.431700
502-2.728600
5030.002000
5040.080400
505-0.002300
506-2.994900
507-2.173100
5080.022300
509-0.106500
510-3.589600
5110.005500
512-0.012400
513-2.807500
514-0.042700
515-0.018200
516-0.281200
517-0.109800
5180.120300
519-0.010300
520-0.120000
521-0.008200
5220.511700
523-2.977200
5240.040000
525-0.301600
526-2.582100
527-2.656200
528-0.053700
529-0.036300
530-0.038000
531-2.099200
532-2.468400
5330.055900
5340.061800
535-0.179600
5360.193700
537-2.825300
538-2.898400
5390.031500
5400.051300
5410.075100
5420.003300
543-0.236400
544-1.447900
5450.043400
546-0.112300
547-0.004000
548-1.012900
549-1.672700
5500.002100
551-0.518700
552-0.938200
553-0.127800
554-0.268100
555-0.080700
556-1.655800
557-2.655000
558-0.010000
559-0.084300
5600.102100
561-0.508200
5620.014300
563-0.006900
564-2.803600
5650.020000
566-2.456000
567-2.202500
5680.329000
5690.031000
5700.017700
571-0.004600
5720.063300
5730.078400
574-0.958100
575-2.237700
576-0.103100
577-0.529300
5780.127300
579-0.110400
5800.206600
581-1.718000
5820.089500
5830.072000
584-0.002300
5850.008300
586-0.105200
587-0.027100
588-0.023400
589-0.000300
590-2.722500
591-0.090200
592-0.571600
5930.076700
5940.116600
595-2.682600
596-0.895400
597-0.264900
598-1.421000
5990.050600
6000.001300
601-3.315900
602-1.485900
6030.067400
6040.181400
605-0.464900
606-0.017700
607-2.459800
608-1.080900
609-0.009600
610-0.121100
611-0.350500
612-2.640200
6130.019500
6140.006600
615-0.000000
616-0.264100
617-0.858600
618-0.304900
619-0.741100
6200.069200
6210.065300
6220.030700
6230.105600
624-0.000300
625-0.573300
6260.001100
627-0.802500
6280.006800
629-0.921000
6300.040600
6310.019500
632-1.903200
633-0.000400
634-0.397500
635-0.059500
6360.000600
637-0.785000
638-0.010900
639-2.925300
6400.346400
641-0.456900
642-2.750100
643-1.413700
644-2.607100
645-1.993000
646-0.632400
647-0.243200
6480.665100
649-1.393100
6500.115900
651-2.527300
652-0.005000
653-0.000000
654-0.087800
655-0.240900
656-0.173900
657-0.005300
658-0.477400
659-0.099800
660-1.138000
6610.112300
662-0.043900
663-0.231500
6640.004800
665-1.292600
666-0.919900
6670.119000
6680.070100
669-3.057900
670-1.288600
6710.000300
672-0.000800
673-2.389800
674-0.735600
6750.000900
676-0.000200
6770.074700
678-2.448400
6790.328500
6800.017600
681-0.341900
6820.000100
6830.000000
684-0.000300
685-1.595200
686-0.133500
687-0.006900
688-0.103400
6890.057500
690-0.076600
6910.030500
692-0.092600
6930.002600
6940.033100
695-0.001200
696-0.084400
6970.000000
698-0.141600
6991.427800
700-3.011400
701-0.150000
7020.001300
703-3.159600
704-0.002000
705-0.244800
706-0.275200
7070.022200
708-2.752300
709-0.161500
7100.151600
711-2.370600
712-0.016600
713-0.004300
714-0.087000
7150.047700
716-0.637700
717-0.619700
718-0.001200
7190.067100
720-0.040100
721-0.000100
722-0.647000
723-2.528700
724-0.010700
725-2.753400
726-1.927600
7270.000500
728-0.004500
729-1.368800
730-2.052900
7310.024000
732-0.000000
733-0.082900
734-0.048900
7350.063800
736-2.362900
7370.006100
738-0.015200
7390.023000
7400.083000
741-2.225200
7420.048900
743-0.005100
744-2.294800
745-1.988700
7460.304500
747-0.115700
748-0.000200
749-0.036200
7500.001000
751-2.071700
752-0.085900
753-0.188100
754-2.027200
7550.070500
756-0.053000
757-0.739200
758-2.084000
759-2.533400
7600.000300
761-0.004400
762-2.832100
763-0.339100
764-0.089500
7650.041700
766-0.003500
7670.022100
7680.069400
769-0.155700
7700.000800
7710.089300
772-0.000400
773-2.663600
774-0.417300
775-0.323400
776-2.160900
777-0.000100
7780.075300
7790.405100
7800.015300
781-1.782800
782-0.001300
7830.305200
7840.065500
785-0.000500
7860.080700
787-0.057800
788-0.002600
789-0.001400
790-0.034800
791-0.000100
792-1.877500
7930.019500
7940.009600
795-0.105300
796-0.129900
797-0.311400
798-2.437700
799-0.000400
800-1.523900
801-1.492800
8020.039000
803-0.000700
804-0.021000
805-0.065900
806-0.027100
807-0.049800
8080.035400
809-0.001700
810-0.033200
811-0.196500
812-0.001000
813-0.032900
814-0.017100
8150.028700
816-0.171300
8170.015300
818-0.005900
819-0.426000
820-0.943700
821-0.000700
822-0.000300
823-0.004300
824-0.009700
8250.198600
8260.328900
8270.188800
8280.968800
829-0.963100
830-2.783300
831-0.474300
832-0.294600
833-0.148500
834-1.373600
8350.065700
836-0.182400
8370.032700
838-0.007400
8390.008000
840-0.039900
841-0.057900
8420.000500
843-0.352600
844-1.688000
8450.061000
846-0.000600
847-0.000100
8480.000200
849-0.001400
850-1.233800
8510.109100
852-0.000800
853-0.244900
854-0.038700
8550.243000
8560.000400
8570.029700
858-0.001800
8590.164600
860-0.329100
861-0.033400
8620.052800
863-0.468800
8640.000300
865-0.768100
866-0.061700
867-0.181800
8680.001100
8690.001600
870-1.857500
871-0.003400
8720.169300
873-0.036900
8740.010900
875-0.148800
8760.085100
8770.150200
878-0.004700
879-0.000700
880-0.137900
881-0.005300
882-0.408700
883-2.271500
8840.142300
885-0.066900
886-1.327700
887-0.062300
888-0.205300
8890.040900
890-0.548300
8910.039000
8920.045400
893-1.917200
894-0.202700
895-0.315300
896-0.198700
897-2.117000
8980.355500
899-0.217000
9000.122500
901-0.136400
902-0.859600
903-0.075500
904-0.215500
905-0.000800
906-0.097000
907-0.064300
908-0.290800
9090.095900
910-0.028900
911-0.001200
912-0.003500
913-0.005100
914-0.408500
915-0.021700
9160.000600
917-0.089700
9180.000700
9190.000300
9200.053300
9210.028100
9220.043800
923-0.095000
924-0.058000
9250.184300
926-0.042600
927-0.000300
928-0.065500
9291.266400
9300.007000
931-0.001700
932-0.001600
9330.002200
934-0.013400
9350.065000
9360.074500
937-0.122800
9380.012400
9390.006300
9400.001500
941-0.000100
942-0.553000
9430.103400
944-0.116500
945-2.889400
946-0.010600
9470.000300
9480.044500
9490.038800
9500.006700
951-0.014100
9520.002200
9530.198000
954-0.000400
955-0.057400
956-2.319500
9570.000100
9580.043000
959-0.001900
9600.052200
961-0.080400
962-0.390700
963-0.405800
964-3.307000
965-0.125900
9660.061800
9670.003300
968-0.005800
969-0.002400
9700.000100
9710.119900
9720.036500
9730.011900
974-1.097000
975-0.085300
976-0.000100
9770.000200
978-0.078800
9790.030700
9800.000100
981-0.038500
982-0.011800
983-0.005600
9840.049900
985-0.000200
986-0.083200
987-0.140700
9880.014100
989-0.061700
990-0.000200
991-0.000300
992-0.001000
9930.016900
994-0.235500
9950.254100
996-0.007700
997-0.000200
9980.025300
9990.042600

" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "ename": "RuntimeError", - "evalue": "[enforce fail at inline_container.cc:664] . unexpected pos 11819474816 vs 11819474712", - "output_type": "error", - "traceback": [ - "\u001b[31m---------------------------------------------------------------------------\u001b[39m\n", - "\u001b[31mRuntimeError\u001b[39m Traceback (most recent call last)\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:967\u001b[39m, in \u001b[36msave\u001b[39m\u001b[34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[39m\n", - "\u001b[32m 966\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m _open_zipfile_writer(f) \u001b[38;5;28;01mas\u001b[39;00m opened_zipfile:\n", - "\u001b[32m--> \u001b[39m\u001b[32m967\u001b[39m \u001b[43m_save\u001b[49m\u001b[43m(\u001b[49m\n", - "\u001b[32m 968\u001b[39m \u001b[43m \u001b[49m\u001b[43mobj\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 969\u001b[39m \u001b[43m \u001b[49m\u001b[43mopened_zipfile\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 970\u001b[39m \u001b[43m \u001b[49m\u001b[43mpickle_module\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 971\u001b[39m \u001b[43m \u001b[49m\u001b[43mpickle_protocol\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 972\u001b[39m \u001b[43m \u001b[49m\u001b[43m_disable_byteorder_record\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 973\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 974\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:1268\u001b[39m, in \u001b[36m_save\u001b[39m\u001b[34m(obj, zip_file, pickle_module, pickle_protocol, _disable_byteorder_record)\u001b[39m\n", - "\u001b[32m 1267\u001b[39m \u001b[38;5;66;03m# Now that it is on the CPU we can directly copy it into the zip file\u001b[39;00m\n", - "\u001b[32m-> \u001b[39m\u001b[32m1268\u001b[39m \u001b[43mzip_file\u001b[49m\u001b[43m.\u001b[49m\u001b[43mwrite_record\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstorage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_bytes\u001b[49m\u001b[43m)\u001b[49m\n", - "\n", - "\u001b[31mRuntimeError\u001b[39m: [enforce fail at inline_container.cc:858] . PytorchStreamWriter failed writing file data/496: file write failed\n", - "\n", - "During handling of the above exception, another exception occurred:\n", - "\n", - "\u001b[31mRuntimeError\u001b[39m Traceback (most recent call last)\n", - "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[15]\u001b[39m\u001b[32m, line 13\u001b[39m\n", - "\u001b[32m 2\u001b[39m os.environ[\u001b[33m\"\u001b[39m\u001b[33mCUDA_LAUNCH_BLOCKING\u001b[39m\u001b[33m\"\u001b[39m] = \u001b[33m\"\u001b[39m\u001b[33m1\u001b[39m\u001b[33m\"\u001b[39m\n", - "\u001b[32m 3\u001b[39m trainer = JEPOTrainer(\n", - "\u001b[32m 4\u001b[39m model = model,\n", - "\u001b[32m 5\u001b[39m processing_class = tokenizer,\n", - "\u001b[32m (...)\u001b[39m\u001b[32m 11\u001b[39m train_dataset = dataset,\n", - "\u001b[32m 12\u001b[39m )\n", - "\u001b[32m---> \u001b[39m\u001b[32m13\u001b[39m \u001b[43mtrainer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 14\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m torch.cuda.is_available():\n", - "\u001b[32m 15\u001b[39m memory_used = torch.cuda.memory_allocated() / \u001b[32m1024\u001b[39m**\u001b[32m3\u001b[39m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:2325\u001b[39m, in \u001b[36mTrainer.train\u001b[39m\u001b[34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[39m\n", - "\u001b[32m 2323\u001b[39m hf_hub_utils.enable_progress_bars()\n", - "\u001b[32m 2324\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "\u001b[32m-> \u001b[39m\u001b[32m2325\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n", - "\u001b[32m 2326\u001b[39m \u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m=\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2327\u001b[39m \u001b[43m \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m=\u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2328\u001b[39m \u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2329\u001b[39m \u001b[43m \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m=\u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2330\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:2756\u001b[39m, in \u001b[36mTrainer._inner_training_loop\u001b[39m\u001b[34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[39m\n", - "\u001b[32m 2754\u001b[39m \u001b[38;5;28mself\u001b[39m.state.epoch = epoch + (step + \u001b[32m1\u001b[39m) / steps_in_epoch\n", - "\u001b[32m 2755\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_step_end(args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", - "\u001b[32m-> \u001b[39m\u001b[32m2756\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_maybe_log_save_evaluate\u001b[49m\u001b[43m(\u001b[49m\n", - "\u001b[32m 2757\u001b[39m \u001b[43m \u001b[49m\u001b[43mtr_loss\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2758\u001b[39m \u001b[43m \u001b[49m\u001b[43mgrad_norm\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2759\u001b[39m \u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2760\u001b[39m \u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2761\u001b[39m \u001b[43m \u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2762\u001b[39m \u001b[43m \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2763\u001b[39m \u001b[43m \u001b[49m\u001b[43mstart_time\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2764\u001b[39m \u001b[43m \u001b[49m\u001b[43mlearning_rate\u001b[49m\u001b[43m=\u001b[49m\u001b[43mlearning_rate\u001b[49m\u001b[43m,\u001b[49m\n", - "\u001b[32m 2765\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 2766\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "\u001b[32m 2767\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_substep_end(args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3228\u001b[39m, in \u001b[36mTrainer._maybe_log_save_evaluate\u001b[39m\u001b[34m(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate)\u001b[39m\n", - "\u001b[32m 3225\u001b[39m \u001b[38;5;28mself\u001b[39m.control.should_save = is_new_best_metric\n", - "\u001b[32m 3227\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.control.should_save:\n", - "\u001b[32m-> \u001b[39m\u001b[32m3228\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_save_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 3229\u001b[39m \u001b[38;5;28mself\u001b[39m.control = \u001b[38;5;28mself\u001b[39m.callback_handler.on_save(\u001b[38;5;28mself\u001b[39m.args, \u001b[38;5;28mself\u001b[39m.state, \u001b[38;5;28mself\u001b[39m.control)\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/trl/trainer/jepo_trainer.py:1854\u001b[39m, in \u001b[36mJEPOTrainer._save_checkpoint\u001b[39m\u001b[34m(self, model, trial)\u001b[39m\n", - "\u001b[32m 1852\u001b[39m model_name = \u001b[38;5;28mself\u001b[39m.args.hub_model_id.split(\u001b[33m\"\u001b[39m\u001b[33m/\u001b[39m\u001b[33m\"\u001b[39m)[-\u001b[32m1\u001b[39m]\n", - "\u001b[32m 1853\u001b[39m \u001b[38;5;28mself\u001b[39m.create_model_card(model_name=model_name)\n", - "\u001b[32m-> \u001b[39m\u001b[32m1854\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_save_checkpoint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrial\u001b[49m\u001b[43m)\u001b[49m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3345\u001b[39m, in \u001b[36mTrainer._save_checkpoint\u001b[39m\u001b[34m(self, model, trial)\u001b[39m\n", - "\u001b[32m 3341\u001b[39m \u001b[38;5;28mself\u001b[39m.state.best_model_checkpoint = best_checkpoint_dir\n", - "\u001b[32m 3343\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m.args.save_only_model:\n", - "\u001b[32m 3344\u001b[39m \u001b[38;5;66;03m# Save optimizer and scheduler\u001b[39;00m\n", - "\u001b[32m-> \u001b[39m\u001b[32m3345\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_save_optimizer_and_scheduler\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_dir\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 3346\u001b[39m \u001b[38;5;28mself\u001b[39m._save_scaler(output_dir)\n", - "\u001b[32m 3347\u001b[39m \u001b[38;5;66;03m# Save RNG state\u001b[39;00m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/transformers/trainer.py:3472\u001b[39m, in \u001b[36mTrainer._save_optimizer_and_scheduler\u001b[39m\u001b[34m(self, output_dir)\u001b[39m\n", - "\u001b[32m 3467\u001b[39m save_fsdp_optimizer(\n", - "\u001b[32m 3468\u001b[39m \u001b[38;5;28mself\u001b[39m.accelerator.state.fsdp_plugin, \u001b[38;5;28mself\u001b[39m.accelerator, \u001b[38;5;28mself\u001b[39m.optimizer, \u001b[38;5;28mself\u001b[39m.model, output_dir\n", - "\u001b[32m 3469\u001b[39m )\n", - "\u001b[32m 3470\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.args.should_save:\n", - "\u001b[32m 3471\u001b[39m \u001b[38;5;66;03m# deepspeed.save_checkpoint above saves model/optim/sched\u001b[39;00m\n", - "\u001b[32m-> \u001b[39m\u001b[32m3472\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mstate_dict\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mos\u001b[49m\u001b[43m.\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m.\u001b[49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mOPTIMIZER_NAME\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 3474\u001b[39m \u001b[38;5;66;03m# Save SCHEDULER & SCALER\u001b[39;00m\n", - "\u001b[32m 3475\u001b[39m is_deepspeed_custom_scheduler = \u001b[38;5;28mself\u001b[39m.is_deepspeed_enabled \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\n", - "\u001b[32m 3476\u001b[39m \u001b[38;5;28mself\u001b[39m.lr_scheduler, DeepSpeedSchedulerWrapper\n", - "\u001b[32m 3477\u001b[39m )\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:966\u001b[39m, in \u001b[36msave\u001b[39m\u001b[34m(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization, _disable_byteorder_record)\u001b[39m\n", - "\u001b[32m 963\u001b[39m f = os.fspath(f)\n", - "\u001b[32m 965\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m _use_new_zipfile_serialization:\n", - "\u001b[32m--> \u001b[39m\u001b[32m966\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m _open_zipfile_writer(f) \u001b[38;5;28;01mas\u001b[39;00m opened_zipfile:\n", - "\u001b[32m 967\u001b[39m _save(\n", - "\u001b[32m 968\u001b[39m obj,\n", - "\u001b[32m 969\u001b[39m opened_zipfile,\n", - "\u001b[32m (...)\u001b[39m\u001b[32m 972\u001b[39m _disable_byteorder_record,\n", - "\u001b[32m 973\u001b[39m )\n", - "\u001b[32m 974\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m\n", - "\n", - "\u001b[36mFile \u001b[39m\u001b[32m/anaconda/envs/dp-eval/lib/python3.13/site-packages/torch/serialization.py:798\u001b[39m, in \u001b[36m_open_zipfile_writer_file.__exit__\u001b[39m\u001b[34m(self, *args)\u001b[39m\n", - "\u001b[32m 797\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m__exit__\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args) -> \u001b[38;5;28;01mNone\u001b[39;00m:\n", - "\u001b[32m--> \u001b[39m\u001b[32m798\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mfile_like\u001b[49m\u001b[43m.\u001b[49m\u001b[43mwrite_end_of_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[32m 799\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.file_stream \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", - "\u001b[32m 800\u001b[39m \u001b[38;5;28mself\u001b[39m.file_stream.close()\n", - "\n", - "\u001b[31mRuntimeError\u001b[39m: [enforce fail at inline_container.cc:664] . unexpected pos 11819474816 vs 11819474712" - ] - } - ], + "outputs": [], "source": [ "import os \n", "os.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"\n", @@ -4683,41 +492,6 @@ " print(f\"💾 GPU Memory Used: {memory_used:.2f} GB / {memory_total:.1f} GB ({memory_used/memory_total*100:.1f}%)\")" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "6a99f973", - "metadata": {}, - "outputs": [], - "source": [ - "f.flush()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b94d80e9", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np \n", - " \n", - "# Given array \n", - "data = np.array([0, 0, 0, -1]) \n", - " \n", - "# Compute the mean and standard deviation \n", - "mean = np.mean(data) \n", - "std = np.std(data) \n", - " \n", - "# Compute the z-score normalization: (x - mean) / std \n", - "z_scores = (data - mean) / (std + 1e-4) # Adding a small value to avoid division by zero \n", - " \n", - "print(\"Original array:\", data) \n", - "print(\"Mean:\", mean) \n", - "print(\"Standard Deviation:\", std) \n", - "print(\"Z-score normalized array:\", z_scores) " - ] - }, { "cell_type": "code", "execution_count": null, @@ -4749,38 +523,6 @@ "output_text = tokenizer.decode(outputs[0], skip_special_tokens=True) \n", "print(output_text) \n" ] - }, - { - "cell_type": "markdown", - "id": "0a3300ec", - "metadata": {}, - "source": [ - "And now with the LoRA we just trained with GRPO - we first save the LoRA first!" - ] - }, - { - "cell_type": "markdown", - "id": "ff6b0e7f", - "metadata": {}, - "source": [ - "Now we load the LoRA and test:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "bb1be866", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "markdown", - "id": "2c6e5c11", - "metadata": {}, - "source": [ - "Our reasoning model is much better - it's not always correct, since we only trained it for an hour or so - it'll be better if we extend the sequence length and train for longer!" - ] } ], "metadata": { diff --git a/trl/trainer/jepo_config.py b/trl/trainer/jepo_config.py index f44c96282ff..11cf8fa688b 100644 --- a/trl/trainer/jepo_config.py +++ b/trl/trainer/jepo_config.py @@ -15,11 +15,11 @@ from dataclasses import dataclass, field from typing import Optional, Union -from transformers import TrainingArguments +from .base_config import _BaseConfig @dataclass -class JEPOConfig(TrainingArguments): +class JEPOConfig(_BaseConfig): r""" Configuration class for the [`JEPOTrainer`], which serves as a variation of GRPO for unverifiable RL training. JEPO [https://arxiv.org/pdf/2503.19618] diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index 6afb7f02d78..bbd08db8d2e 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -52,7 +52,7 @@ from ..import_utils import is_liger_kernel_available, is_vllm_available from ..models import prepare_deepspeed, prepare_fsdp, prepare_peft_model, unwrap_model_for_generation from ..models.utils import _ForwardRedirection -from .base_trainer import BaseTrainer +from .base_trainer import _BaseTrainer from .callbacks import SyncRefModelCallback from .jepo_config import JEPOConfig from .utils import ( @@ -97,7 +97,7 @@ # What we call a CoT function is a callable that takes completion text and answer and returns CoT text and CoT+answer. CoTFunc = Callable[[str, str], [str, str]] -class JEPOTrainer(BaseTrainer): +class JEPOTrainer(_BaseTrainer): """ Trainer for the Jensen’s Evidence lower bound Policy Optimization (JEPO) method. This algorithm was initially proposed in the paper: Beyond Verifiable Rewards: Scaling Reinforcement Learning for Language Models to Unverifiable Data [https://arxiv.org/pdf/2503.19618] From 21af1f147547fb90e0c24c102ca1700e6495c6d2 Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:00:24 +0000 Subject: [PATCH 06/12] fix multi-gpu --- trl/trainer/jepo_trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index bbd08db8d2e..a5e29f7cb12 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1374,7 +1374,7 @@ def _compute_jepo_advantages(self, jepo_rewards, jepo_applicable_mask, num_gener for j in range(start_idx, end_idx): if jepo_applicable_mask[j]: jepo_advantages[j] = torch.clamp( jepo_advantages[j]/(group_std + 1e-4) , -1, 1) - + jepo_advantages = gather(jepo_advantages) # gather the advantages across processes return jepo_advantages def _generate_and_score_completions( From da81a4cf63ca2c8814e0f1b37b6c880ed93a6135 Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:04:37 +0000 Subject: [PATCH 07/12] minor --- trl/trainer/jepo_trainer.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index a5e29f7cb12..027d2b8da29 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1725,10 +1725,8 @@ def _compute_loss(self, model, inputs): # compute p(C|Q), fabricated_completions has the EOT_TOKEN at the end and answer_mask have that removed, need to align them first before applying the answer_mask cot_mask = torch.cat([fabricated_completions_mask[:, 1:], torch.zeros((fabricated_completions_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) - answer_mask # this is the mask for the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) - #print(f"answer_texts_from_mask: {[answer_text.strip('!') for answer_text in answer_texts_from_mask]}") cot_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * cot_mask, skip_special_tokens=True) - #print(f"cot_texts_from_mask: {[cot_text.strip('!') for cot_text in cot_texts_from_mask]}") cot_token_logps = per_token_logps * cot_mask # this is the log probability of the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. advantages = advantages.unsqueeze(1) # (B, 1) @@ -1739,7 +1737,6 @@ def _compute_loss(self, model, inputs): loss = - per_token_adv.sum(dim=1).sum()/ (self.num_generations) # sum over the tokens and average over the batch else: loss = - (per_token_adv.sum(dim=1)/cot_mask.sum(dim=1)).sum()/ (self.num_generations) # sum over the tokens and average over the batch - print(f"advantages loss: {loss}") #compute P(A|Q,C) and use it as a supervised loss to encourage the model to generate answers if self.supervised_loss_weight > 0.0: @@ -1751,15 +1748,12 @@ def _compute_loss(self, model, inputs): supervised_loss= answer_token_logps.sum(dim=1)/answer_mask.sum(dim=1) # sum per number_generations to get the total log probability of the answer given the question and CoT, which is p(A|Q,C), we will use it as a supervised loss to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. supervised_loss = torch.log(supervised_loss.view(-1, self.num_generations).sum(dim=1)/self.num_generations + 1e-8) #average over the generations and take log to get the supervised loss for each sample in the batch, which is log(p(A|Q,C)), we will use it as a supervised loss to encourage the model to generate answers that are likely under the reference model, which can help stabilize training and prevent divergence. - #print(f"supervised_loss before weighting: {supervised_loss}") loss = loss - self.supervised_loss_weight * supervised_loss.sum() # average over the batch - #print(f"supervised_loss loss: {loss}, supervised_loss_weight: {self.supervised_loss_weight}") if self.beta != 0.0: if self.loss_type == 'unnorm_jepo': loss = loss + self.beta * (per_token_kl * fabricated_completions_mask).sum(dim=1).sum() / (self.num_generations) # sum over the tokens and average over the batch else: loss = loss + self.beta * (((per_token_kl * fabricated_completions_mask).sum(dim=1)) / fabricated_completions_mask.sum(dim=1)).sum() / (self.num_generations) # average over the batch and the tokens, note that we only calculate KL for the completion part, which is indicated by the fabricated_completions_mask. The intuition is that we want to regularize the model to not deviate too much from the reference model, which can help stabilize training and prevent divergence. - print(f" kl loss: {loss}, beta: {self.beta}") # Log the metrics From 5179c9c3c43ca8982205dff012038d86caaf4a9f Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:11:15 +0000 Subject: [PATCH 08/12] fix comments --- trl/trainer/jepo_config.py | 7 +------ trl/trainer/jepo_trainer.py | 2 +- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/trl/trainer/jepo_config.py b/trl/trainer/jepo_config.py index 11cf8fa688b..4924a832c80 100644 --- a/trl/trainer/jepo_config.py +++ b/trl/trainer/jepo_config.py @@ -25,12 +25,7 @@ class JEPOConfig(_BaseConfig): JEPO [https://arxiv.org/pdf/2503.19618] This class includes only the parameters that are specific to JEPO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. Parameters: > Parameters that control the model and reference model @@ -232,7 +227,7 @@ class JEPOConfig(_BaseConfig): are logged. """ - _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + _VALID_DICT_FIELDS = _BaseConfig._VALID_DICT_FIELDS + ["model_init_kwargs"] # Parameters whose default values are overridden from TrainingArguments learning_rate: float = field( diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index 027d2b8da29..9b0f753d786 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1707,7 +1707,7 @@ def _compute_loss(self, model, inputs): ) if self.top_entropy_quantile < 1.0: - entropy_mask = self.get_high_entropy_mask(entropies, completion_mask, 1 - self.top_entropy_quantile) + entropy_mask = self.get_high_entropy_mask(entropies, fabricated_completions_mask, 1 - self.top_entropy_quantile) else: entropy_mask = None From 84abef212eb9755ddc2d82e3bf47dfddfb1c569d Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:25:23 +0000 Subject: [PATCH 09/12] fix distrubted computing --- trl/trainer/jepo_trainer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index 9b0f753d786..a591de9aab6 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1374,7 +1374,6 @@ def _compute_jepo_advantages(self, jepo_rewards, jepo_applicable_mask, num_gener for j in range(start_idx, end_idx): if jepo_applicable_mask[j]: jepo_advantages[j] = torch.clamp( jepo_advantages[j]/(group_std + 1e-4) , -1, 1) - jepo_advantages = gather(jepo_advantages) # gather the advantages across processes return jepo_advantages def _generate_and_score_completions( @@ -1562,6 +1561,8 @@ def synthesize_completion(prompt, completion, answer): #print(f"answer_mask: {answer_mask.sum(dim=1)}, fabricated_completions_mask: {fabricated_completions_mask.sum(dim=1)}, padded_cot_attention_mask: {padded_cot_attention_mask.sum(dim=1)}") jepo_rewards[jepo_applicable_mask == False] = 0.0 # if not applicable for JEPO, set reward to 0 + jepo_rewards = gather(jepo_rewards) + jepo_applicable_mask = gather(jepo_applicable_mask) #calculate the mean of jepo rewards for each group of generations, based on the jepo_applicable_mask jepo_advantages = self._compute_jepo_advantages(jepo_rewards, jepo_applicable_mask, self.num_generations, device=device) #print(f"jepo_rewards: {jepo_rewards}, jepo_advantages: {jepo_advantages}") @@ -1653,8 +1654,6 @@ def synthesize_completion(prompt, completion, answer): "advantages": advantages, "num_items_in_batch": num_items_in_batch, } - if self.use_vllm and self.vllm_importance_sampling_correction: - output["importance_sampling_ratio"] = importance_sampling_ratio if ref_per_token_logps is not None: output["ref_per_token_logps"] = ref_per_token_logps if "pixel_values" in forward_kwargs: From fa40825880ca4f968cbdcd2f4e19a3dcc93c2045 Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:36:52 +0000 Subject: [PATCH 10/12] minor --- trl/trainer/jepo_trainer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index a591de9aab6..fcf8212d294 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1469,7 +1469,6 @@ def _generate_and_score_completions( # Note the advantagees would be 0 if the sample does not following the format required by the format reward function; # get clone of rewards jepo_rewards = rewards.clone() - logits_to_keep_answer = rewards.clone() with torch.no_grad(): From 69e4c4e32761ff79cf5c05ccd14345d8203293b6 Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:51:55 +0000 Subject: [PATCH 11/12] minor --- trl/trainer/jepo_trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index fcf8212d294..2cb81371c74 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -1549,7 +1549,7 @@ def synthesize_completion(prompt, completion, answer): #fix: there is EOT_TOKEN at the end of the cot_text and cot_answer_text ship by one answer_mask = torch.cat([fabricated_completions_mask[:, 1:], torch.zeros((fabricated_completions_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) - torch.cat([padded_cot_attention_mask[:, 1:], torch.zeros((padded_cot_attention_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) #check the answer_mask is correct by decoding the tokens corresponding to the answer_mask and see if it matches with the cot_answer_texts using padding mask to ignore the padding tokens - answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) + #answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) #print(f"answer_texts_from_mask: {[answer_text.strip('!') for answer_text in answer_texts_from_mask]}") #compute answer log probabilities (i.e., p(A|Q,C)) and use it as JEPO rewards by per_token_logits over answer mask, then sum over the answer tokens to get the total log probability of the answer given the question and CoT. The intuition is that if the sampled CoT is helpful for generating the completion, then the log probability of the answer given the prompt and CoT should be high, resulting in a high JEPO reward @@ -1722,9 +1722,9 @@ def _compute_loss(self, model, inputs): # compute p(C|Q), fabricated_completions has the EOT_TOKEN at the end and answer_mask have that removed, need to align them first before applying the answer_mask cot_mask = torch.cat([fabricated_completions_mask[:, 1:], torch.zeros((fabricated_completions_mask.shape[0], 1), dtype=torch.long, device=device)], dim=1) - answer_mask # this is the mask for the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. - answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) + #answer_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * answer_mask, skip_special_tokens=True) - cot_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * cot_mask, skip_special_tokens=True) + #cot_texts_from_mask = self.processing_class.batch_decode(fabricated_completions_ids * cot_mask, skip_special_tokens=True) cot_token_logps = per_token_logps * cot_mask # this is the log probability of the CoT part, which is p(C|x), we will use it to calculate the JEPO reward, the intuition is that if the sampled CoT is helpful for generating the completion, then p(C|x) should be high, resulting in a high JEPO reward. Note that we need to remove the EOT_TOKEN at the end of the fabricated_completions when calculating p(C|x) because the EOT_TOKEN is not part of the CoT, it's just a special token to indicate the end of the sequence. The answer_mask already has that removed, so we can use it to zero out the log probabilities for the answer part and keep only the log probabilities for the CoT part. advantages = advantages.unsqueeze(1) # (B, 1) From 6bf8c8e53e9550ac065138478bca777facaeaadb Mon Sep 17 00:00:00 2001 From: haijunz Date: Fri, 3 Apr 2026 16:53:45 +0000 Subject: [PATCH 12/12] fix typo --- trl/trainer/jepo_trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trl/trainer/jepo_trainer.py b/trl/trainer/jepo_trainer.py index 2cb81371c74..abfa2411043 100644 --- a/trl/trainer/jepo_trainer.py +++ b/trl/trainer/jepo_trainer.py @@ -216,7 +216,7 @@ def __init__( if args is None: model_name = model if isinstance(model, str) else model.config._name_or_path model_name = model_name.split("/")[-1] - args = GRPOConfig(f"{model_name}-JEPO") + args = JEPOConfig(f"{model_name}-JEPO") # Models # Trained model @@ -231,7 +231,7 @@ def __init__( model_init_kwargs["dtype"] = dtype else: raise ValueError( - "Invalid `dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing " + "Invalid `dtype` passed to `JEPOConfig`. Expected either 'auto' or a string representing " f"a `torch.dtype` (e.g., 'float32'), but got {dtype}." ) # Disable caching if gradient checkpointing is enabled (not supported) @@ -242,7 +242,7 @@ def __init__( model_id = model.config._name_or_path if args.model_init_kwargs is not None: logger.warning( - "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "You passed `model_init_kwargs` to the `JEPOConfig`, but your model is already instantiated. " "The `model_init_kwargs` will be ignored." )