wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
478 lines
15 KiB
Plaintext
478 lines
15 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# LLM Alignment with DPO and KTO\n",
|
|
"\n",
|
|
"[](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/alignment/alignment_dpo.ipynb)\n",
|
|
"\n",
|
|
"This notebook walks through aligning **Llama-3.1-8B** with human preferences using Ludwig's built-in\n",
|
|
"**Direct Preference Optimization (DPO)** trainer, then shows how to switch to **KTO** when paired\n",
|
|
"preference data is not available.\n",
|
|
"\n",
|
|
"We use the publicly available [Anthropic HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf)\n",
|
|
"dataset, which contains helpfulness and harmlessness preference pairs collected from human annotators.\n",
|
|
"\n",
|
|
"### What you will learn\n",
|
|
"- How to parse the HH-RLHF dataset into Ludwig's expected column format\n",
|
|
"- How to configure and run DPO alignment training with LoRA\n",
|
|
"- How to evaluate response quality before and after alignment\n",
|
|
"- How to switch to KTO with a one-line config change\n",
|
|
"- How to upload the aligned model to HuggingFace Hub\n",
|
|
"\n",
|
|
"### Prerequisites\n",
|
|
"- **GPU**: A100 (40 GiB) recommended. T4 will work with smaller batch sizes but is slower.\n",
|
|
"- **HuggingFace token**: Required for Llama-3.1-8B. Set `HUGGING_FACE_HUB_TOKEN` before running.\n",
|
|
"- **Model access**: Request access at https://huggingface.co/meta-llama/Meta-Llama-3.1-8B"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Install dependencies"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install \"ludwig[llm]\" datasets --quiet"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"# Set your HuggingFace token — required to download Llama-3.1-8B.\n",
|
|
"# In Colab: Secrets panel (key icon) -> add HF_TOKEN, then reference it here.\n",
|
|
"try:\n",
|
|
" from google.colab import userdata\n",
|
|
"\n",
|
|
" os.environ[\"HUGGING_FACE_HUB_TOKEN\"] = userdata.get(\"HF_TOKEN\")\n",
|
|
"except Exception:\n",
|
|
" pass # Running locally — export HUGGING_FACE_HUB_TOKEN in your shell instead\n",
|
|
"\n",
|
|
"assert os.environ.get(\"HUGGING_FACE_HUB_TOKEN\"), (\n",
|
|
" \"HUGGING_FACE_HUB_TOKEN is not set. Add it to Colab Secrets or export it in your shell.\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import subprocess\n",
|
|
"\n",
|
|
"# Verify GPU availability\n",
|
|
"result = subprocess.run(\n",
|
|
" [\"nvidia-smi\", \"--query-gpu=name,memory.total\", \"--format=csv,noheader\"], capture_output=True, text=True\n",
|
|
")\n",
|
|
"if result.returncode == 0:\n",
|
|
" print(\"GPU(s) detected:\")\n",
|
|
" print(result.stdout.strip())\n",
|
|
"else:\n",
|
|
" print(\"WARNING: No GPU detected. Training will be very slow on CPU.\")\n",
|
|
" print(\"In Colab: Runtime -> Change runtime type -> A100 GPU\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Prepare dataset\n",
|
|
"\n",
|
|
"The [Anthropic HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf) dataset stores full\n",
|
|
"multi-turn conversations in two columns: `chosen` (preferred response) and `rejected` (dispreferred\n",
|
|
"response). Each value is a raw conversation string like:\n",
|
|
"\n",
|
|
"```\n",
|
|
"\\n\\nHuman: How do I make pasta?\\n\\nAssistant: Boil water, add pasta...\\n\\nHuman: What sauce?\\n\\nAssistant: Try marinara...\n",
|
|
"```\n",
|
|
"\n",
|
|
"Ludwig's DPO trainer expects three columns: `prompt`, `chosen`, `rejected`, where `prompt` is the\n",
|
|
"last human turn and `chosen`/`rejected` are the final assistant responses.\n",
|
|
"\n",
|
|
"The cell below does that extraction."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import re\n",
|
|
"\n",
|
|
"import pandas as pd\n",
|
|
"from datasets import load_dataset\n",
|
|
"\n",
|
|
"\n",
|
|
"def extract_last_human_turn(conversation: str) -> str:\n",
|
|
" \"\"\"Extract the last Human turn from an HH-RLHF conversation string.\"\"\"\n",
|
|
" turns = re.findall(r\"\\n\\nHuman: (.*?)(?=\\n\\nAssistant:|\\Z)\", conversation, re.DOTALL)\n",
|
|
" return turns[-1].strip() if turns else conversation.strip()\n",
|
|
"\n",
|
|
"\n",
|
|
"def extract_last_assistant_turn(conversation: str) -> str:\n",
|
|
" \"\"\"Extract the last Assistant turn from an HH-RLHF conversation string.\"\"\"\n",
|
|
" turns = re.findall(r\"\\n\\nAssistant: (.*?)(?=\\n\\nHuman:|\\Z)\", conversation, re.DOTALL)\n",
|
|
" return turns[-1].strip() if turns else \"\"\n",
|
|
"\n",
|
|
"\n",
|
|
"def prepare_dpo_dataset(split, max_samples=None):\n",
|
|
" if max_samples:\n",
|
|
" split = split.select(range(min(max_samples, len(split))))\n",
|
|
" rows = []\n",
|
|
" for ex in split:\n",
|
|
" prompt = extract_last_human_turn(ex[\"chosen\"])\n",
|
|
" chosen = extract_last_assistant_turn(ex[\"chosen\"])\n",
|
|
" rejected = extract_last_assistant_turn(ex[\"rejected\"])\n",
|
|
" if prompt and chosen and rejected:\n",
|
|
" rows.append({\"prompt\": prompt, \"chosen\": chosen, \"rejected\": rejected})\n",
|
|
" return pd.DataFrame(rows)\n",
|
|
"\n",
|
|
"\n",
|
|
"print(\"Downloading Anthropic/hh-rlhf ...\")\n",
|
|
"hh = load_dataset(\"Anthropic/hh-rlhf\")\n",
|
|
"print(f\"Train size: {len(hh['train'])}, Test size: {len(hh['test'])}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Use a subset for a quick experiment. Remove max_samples caps for full training.\n",
|
|
"MAX_TRAIN = 5000\n",
|
|
"MAX_TEST = 500\n",
|
|
"\n",
|
|
"train_df = prepare_dpo_dataset(hh[\"train\"], max_samples=MAX_TRAIN)\n",
|
|
"test_df = prepare_dpo_dataset(hh[\"test\"], max_samples=MAX_TEST)\n",
|
|
"\n",
|
|
"train_df.to_csv(\"train.csv\", index=False)\n",
|
|
"test_df.to_csv(\"test.csv\", index=False)\n",
|
|
"\n",
|
|
"print(f\"Saved {len(train_df)} train rows to train.csv\")\n",
|
|
"print(f\"Saved {len(test_df)} test rows to test.csv\")\n",
|
|
"train_df.head(2)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Preview: show a prompt with its chosen vs rejected response\n",
|
|
"row = train_df.iloc[0]\n",
|
|
"print(\"=== PROMPT ===\")\n",
|
|
"print(row[\"prompt\"])\n",
|
|
"print(\"\\n=== CHOSEN ===\")\n",
|
|
"print(row[\"chosen\"])\n",
|
|
"print(\"\\n=== REJECTED ===\")\n",
|
|
"print(row[\"rejected\"])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## DPO training\n",
|
|
"\n",
|
|
"Direct Preference Optimization (DPO) treats alignment as a classification problem: given a prompt,\n",
|
|
"prefer `chosen` over `rejected`. The loss function directly optimises the log-ratio between the\n",
|
|
"policy's probability of the chosen response and the reference model's, without needing an explicit\n",
|
|
"reward model.\n",
|
|
"\n",
|
|
"Key hyperparameters:\n",
|
|
"- `beta`: KL penalty coefficient (0.1 is a typical starting point). Higher values keep the aligned\n",
|
|
" model closer to the base model.\n",
|
|
"- `learning_rate`: 5e-7 is much lower than SFT — the policy already knows how to generate text;\n",
|
|
" we're only nudging its preferences.\n",
|
|
"- `lora r=16, alpha=32`: Standard LoRA settings. Increase `r` for larger models or more complex tasks."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import logging\n",
|
|
"\n",
|
|
"import yaml\n",
|
|
"\n",
|
|
"from ludwig.api import LudwigModel\n",
|
|
"\n",
|
|
"dpo_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: meta-llama/Llama-3.1-8B\n",
|
|
"\n",
|
|
"adapter:\n",
|
|
" type: lora\n",
|
|
" r: 16\n",
|
|
" alpha: 32\n",
|
|
" dropout: 0.05\n",
|
|
"\n",
|
|
"trainer:\n",
|
|
" type: dpo\n",
|
|
" epochs: 1\n",
|
|
" learning_rate: 5.0e-7\n",
|
|
" batch_size: 2\n",
|
|
" gradient_accumulation_steps: 8\n",
|
|
" beta: 0.1\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: prompt\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: chosen\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"dpo_model = LudwigModel(config=dpo_config, logging_level=logging.INFO)\n",
|
|
"\n",
|
|
"train_stats, _, output_directory = dpo_model.train(\n",
|
|
" dataset=\"train.csv\",\n",
|
|
" experiment_name=\"hh_rlhf_dpo\",\n",
|
|
" output_directory=\"results\",\n",
|
|
")\n",
|
|
"\n",
|
|
"print(f\"\\nModel saved to: {output_directory}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Evaluate\n",
|
|
"\n",
|
|
"We compare responses from the base model (no alignment) and the DPO-aligned model on a few\n",
|
|
"prompts from the test set. You should see the aligned model give more helpful, less evasive answers."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Load the base model for comparison\n",
|
|
"base_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: meta-llama/Llama-3.1-8B\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: prompt\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: chosen\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"base_model = LudwigModel(config=base_config, logging_level=logging.WARNING)\n",
|
|
"_ = base_model.train(dataset=\"train.csv\", experiment_name=\"baseline\", output_directory=\"results_base\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"eval_prompts = test_df[\"prompt\"].head(3).tolist()\n",
|
|
"eval_df = pd.DataFrame({\"prompt\": eval_prompts})\n",
|
|
"\n",
|
|
"base_preds, _ = base_model.predict(dataset=eval_df)\n",
|
|
"dpo_preds, _ = dpo_model.predict(dataset=eval_df)\n",
|
|
"\n",
|
|
"for i, prompt in enumerate(eval_prompts):\n",
|
|
" print(f\"\\n{'=' * 60}\")\n",
|
|
" print(f\"PROMPT: {prompt}\")\n",
|
|
" print(\"\\n--- Base model ---\")\n",
|
|
" print(base_preds.iloc[i][\"chosen_predictions\"])\n",
|
|
" print(\"\\n--- DPO-aligned model ---\")\n",
|
|
" print(dpo_preds.iloc[i][\"chosen_predictions\"])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## KTO alternative\n",
|
|
"\n",
|
|
"**Kahneman-Tversky Optimization (KTO)** is a good alternative to DPO when you cannot collect\n",
|
|
"paired preference data. Instead of a (chosen, rejected) pair per prompt, KTO requires only a\n",
|
|
"single response with a boolean label indicating whether it was desirable.\n",
|
|
"\n",
|
|
"This makes it easy to use binary user feedback (thumbs up / thumbs down, click-through, etc.)\n",
|
|
"as training signal.\n",
|
|
"\n",
|
|
"### Dataset format for KTO\n",
|
|
"\n",
|
|
"We expand each DPO pair into two rows:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"rows_kto = []\n",
|
|
"for _, row in train_df.iterrows():\n",
|
|
" rows_kto.append({\"prompt\": row[\"prompt\"], \"response\": row[\"chosen\"], \"label\": True})\n",
|
|
" rows_kto.append({\"prompt\": row[\"prompt\"], \"response\": row[\"rejected\"], \"label\": False})\n",
|
|
"\n",
|
|
"train_kto_df = pd.DataFrame(rows_kto)\n",
|
|
"train_kto_df.to_csv(\"train_kto.csv\", index=False)\n",
|
|
"print(f\"KTO dataset: {len(train_kto_df)} rows\")\n",
|
|
"train_kto_df.head(4)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"kto_config = yaml.safe_load(\"\"\"\n",
|
|
"model_type: llm\n",
|
|
"base_model: meta-llama/Llama-3.1-8B\n",
|
|
"\n",
|
|
"adapter:\n",
|
|
" type: lora\n",
|
|
" r: 16\n",
|
|
" alpha: 32\n",
|
|
" dropout: 0.05\n",
|
|
"\n",
|
|
"trainer:\n",
|
|
" type: kto\n",
|
|
" epochs: 1\n",
|
|
" learning_rate: 5.0e-7\n",
|
|
" batch_size: 2\n",
|
|
" gradient_accumulation_steps: 8\n",
|
|
" beta: 0.1\n",
|
|
" desirable_weight: 1.0\n",
|
|
" undesirable_weight: 1.0\n",
|
|
"\n",
|
|
"input_features:\n",
|
|
" - name: prompt\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"output_features:\n",
|
|
" - name: response\n",
|
|
" type: text\n",
|
|
"\n",
|
|
"backend:\n",
|
|
" type: local\n",
|
|
"\"\"\")\n",
|
|
"\n",
|
|
"kto_model = LudwigModel(config=kto_config, logging_level=logging.INFO)\n",
|
|
"\n",
|
|
"_, _, kto_output_dir = kto_model.train(\n",
|
|
" dataset=\"train_kto.csv\",\n",
|
|
" experiment_name=\"hh_rlhf_kto\",\n",
|
|
" output_directory=\"results_kto\",\n",
|
|
")\n",
|
|
"\n",
|
|
"print(f\"\\nKTO model saved to: {kto_output_dir}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Save and upload\n",
|
|
"\n",
|
|
"After training, upload the aligned model to HuggingFace Hub to share it or deploy it to an endpoint."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Replace with your HuggingFace org/username and desired repo name\n",
|
|
"HF_REPO = \"your-username/llama-3.1-8b-dpo-hh-rlhf\"\n",
|
|
"\n",
|
|
"# Upload via CLI:\n",
|
|
"print(f\"ludwig upload hf_hub -r {HF_REPO} -m {output_directory}\")\n",
|
|
"\n",
|
|
"# Or via Python API:\n",
|
|
"# LudwigModel.upload_to_hf_hub(HF_REPO, output_directory)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Uncomment to actually run the upload:\n",
|
|
"# !ludwig upload hf_hub -r {HF_REPO} -m {output_directory}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Next steps\n",
|
|
"\n",
|
|
"- **Scale up**: Remove the `max_samples` caps and train on the full HH-RLHF dataset (~160k examples).\n",
|
|
"- **Try ORPO**: Use `config_orpo.yaml` for single-stage SFT + alignment without a reference model.\n",
|
|
"- **Iterate on beta**: Lower `beta` (e.g. 0.05) gives more aggressive alignment but risks reward hacking.\n",
|
|
"- **Evaluation**: Use [MT-Bench](https://github.com/lm-sys/FastChat/tree/main/fastchat/llm_judge) or\n",
|
|
" [AlpacaEval](https://github.com/tatsu-lab/alpaca_eval) for rigorous alignment benchmarks.\n",
|
|
"- **GRPO**: For tasks with a programmatic reward (math, code), see the GRPO trainer docs."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"accelerator": "GPU",
|
|
"colab": {
|
|
"gpuType": "A100",
|
|
"provenance": []
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"name": "python",
|
|
"version": "3.12.0"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
}
|