{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "e3cb8e81",
      "metadata": {
        "id": "e3cb8e81"
      },
      "source": [
        "# Continued Pre-Training (CPT) of LFM3-1.7B with TRL\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Liquid4All/cookbook/blob/main/finetuning/notebooks/lfm3_cpt_with_trl.ipynb) [![GitHub](https://img.shields.io/badge/GitHub-Liquid4All%2Fcookbook-181717?logo=github)](https://github.com/Liquid4All/cookbook) [![Discord](https://img.shields.io/discord/1385439864920739850?color=7289da&label=Discord&logo=discord&logoColor=white)](https://discord.com/invite/liquid-ai)\n",
        "\n",
        "This notebook continues the pre-training of [LFM3-1.7B](https://huggingface.co/LiquidAI/LFM3-1.7B) on raw text from a new domain, using [TRL](https://github.com/huggingface/trl) and LoRA. It runs on a free Colab T4 GPU, or any GPU with at least 16 GB of memory.\n",
        "\n",
        "Continued pre-training (CPT) teaches a model the vocabulary, style, and knowledge of a domain from unlabeled text, with the same next-token prediction objective as pre-training. Our domain is **Turkish**: LFM3-1.7B officially supports 16 languages, and Turkish is not one of them. We train on Turkish Wikipedia articles and measure how well the model predicts held-out Turkish text before and after.\n",
        "\n",
        "## 🎯 What you'll find\n",
        "\n",
        "- **Raw text data**: stream a slice of a large corpus from the Hugging Face Hub\n",
        "- **LoRA + CPT**: train on plain text packed into full rows, as in pre-training\n",
        "- **Evaluation**: perplexity on held-out articles and on out-of-distribution web pages, and a Turkish grammar test, before and after training\n",
        "- **Export**: merge the LoRA weights into a standalone checkpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8688f341",
      "metadata": {
        "id": "8688f341"
      },
      "source": [
        "# 📦 Installation & Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "89714441",
      "metadata": {
        "tags": [
          "install"
        ],
        "id": "89714441"
      },
      "outputs": [],
      "source": [
        "!pip install -qqq \"torch==2.14.1\" \"torchvision==0.29.1\" \"torchao==0.18.0\" \"trl==1.16.0\" \"transformers==5.20.2\" \"peft==0.21.2\" \"accelerate==1.15.0\" \"datasets==5.1.0\" \"huggingface-hub==2.3.0\" \"tokenizers==0.23.3\" \"safetensors==0.8.0\" --progress-bar off"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4229d5cb",
      "metadata": {
        "id": "4229d5cb"
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import os\n",
        "import torch\n",
        "import transformers\n",
        "import trl\n",
        "import peft\n",
        "\n",
        "os.environ[\"WANDB_DISABLED\"] = \"true\"\n",
        "\n",
        "print(f\"📦 PyTorch version: {torch.__version__}\")\n",
        "print(f\"🤗 Transformers version: {transformers.__version__}\")\n",
        "print(f\"📊 TRL version: {trl.__version__}\")\n",
        "print(f\"🎛️ PEFT version: {peft.__version__}\")\n",
        "\n",
        "# LFM3-1.7B is released in bfloat16. GPUs without bfloat16 support, like the T4, load it in float16 instead:\n",
        "# the conversion is exact for almost every weight, and the LoRA adapters train in float32 either way.\n",
        "use_bf16 = torch.cuda.is_bf16_supported(including_emulation=False)\n",
        "dtype = torch.bfloat16 if use_bf16 else torch.float16\n",
        "print(f\"🔢 Precision: {dtype}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2c04dd0b",
      "metadata": {
        "id": "2c04dd0b"
      },
      "source": [
        "# Loading the model from Transformers 🤗\n",
        "\n",
        "We continue training the post-trained LFM3-1.7B, which keeps its chat abilities for the instruction tuning that usually follows CPT. For large CPT runs, consider starting from the base checkpoint instead."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e32ce88b",
      "metadata": {
        "id": "e32ce88b"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "model_id = \"LiquidAI/LFM3-1.7B\"\n",
        "\n",
        "print(\"📚 Loading tokenizer...\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)\n",
        "\n",
        "print(\"🧠 Loading model...\")\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_id,\n",
        "    dtype=dtype,\n",
        "    device_map=\"auto\",\n",
        "    trust_remote_code=True,\n",
        ")\n",
        "\n",
        "print(\"✅ Model loaded successfully!\")\n",
        "print(f\"🔢 Parameters: {model.num_parameters():,}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7901bf09",
      "metadata": {
        "id": "7901bf09"
      },
      "source": [
        "# 📚 Load a domain corpus\n",
        "\n",
        "[Wikipedia](https://huggingface.co/datasets/wikimedia/wikipedia) has one subset per language, and the Turkish one alone has more than 500,000 articles. We stream it, so that only the articles we use are downloaded: 3,000 for training and 100 for evaluation. To check that the model learns Turkish and not just the style of Wikipedia, we also evaluate on 100 Turkish web pages from [FineWeb-2](https://huggingface.co/datasets/HuggingFaceFW/fineweb-2) (news, forums, shops), which the model never sees during training.\n",
        "\n",
        "Each article becomes one `text` example, which is the format TRL expects for plain language modeling. We add the end-of-sequence token so that the model learns where documents end."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "398db5e2",
      "metadata": {
        "id": "398db5e2"
      },
      "outputs": [],
      "source": [
        "from datasets import Dataset, load_dataset\n",
        "\n",
        "stream = load_dataset(\"wikimedia/wikipedia\", \"20231101.tr\", split=\"train\", streaming=True).shuffle(seed=42, buffer_size=10_000)\n",
        "\n",
        "def to_text(article):\n",
        "    return {\"text\": f\"{article['title']}\\n\\n{article['text']}{tokenizer.eos_token}\"}\n",
        "\n",
        "# Skip stub articles, which are mostly lists and templates\n",
        "articles = (to_text(article) for article in stream if len(article[\"text\"]) > 2_000)\n",
        "eval_dataset = Dataset.from_list([next(articles) for _ in range(100)])\n",
        "train_dataset = Dataset.from_list([next(articles) for _ in range(3_000)])\n",
        "\n",
        "# Out-of-distribution check: Turkish web pages (news, forums, shops), a different register from Wikipedia\n",
        "web = load_dataset(\"HuggingFaceFW/fineweb-2\", name=\"tur_Latn\", split=\"test\", streaming=True)\n",
        "ood_dataset = Dataset.from_list([{\"text\": page[\"text\"] + tokenizer.eos_token} for page in web.take(100)])\n",
        "\n",
        "n_tokens = sum(len(ids) for ids in tokenizer(list(train_dataset[\"text\"]))[\"input_ids\"])\n",
        "print(f\"📚 Train articles: {len(train_dataset)} ({n_tokens / 1e6:.1f}M tokens)\")\n",
        "print(f\"🧪 Eval articles: {len(eval_dataset)}\")\n",
        "print(f\"🌐 Out-of-distribution web pages: {len(ood_dataset)}\")\n",
        "print(f\"\\n📝 Sample:\\n{train_dataset[0]['text'][:500]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "27878ee8",
      "metadata": {
        "id": "27878ee8"
      },
      "source": [
        "# 🔍 A grammar test: Turkish minimal pairs\n",
        "\n",
        "Perplexity tells us that the model predicts Turkish text better. To check that it learned Turkish grammar, and not just the vocabulary of Wikipedia, we use [TurBLiMP](https://huggingface.co/datasets/juletxara/turblimp), a benchmark of minimal pairs: two sentences that differ in one detail, where only one is grammatical. A model \"passes\" a pair when it gives a higher probability to the grammatical sentence. The pairs cover 16 phenomena (agreement, binding, word order, and more), and none of them come from Wikipedia.\n",
        "\n",
        "We use 100 pairs per phenomenon, and score the model before training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "981226fc",
      "metadata": {
        "id": "981226fc"
      },
      "outputs": [],
      "source": [
        "from datasets import concatenate_datasets\n",
        "\n",
        "phenomena = [\"anaphor_agreement\", \"argument_structure_ditransitive\", \"argument_structure_transitive\", \"binding\", \"determiners\", \"ellipsis\", \"irregular_forms\", \"island_effects\", \"nominalization\", \"npi_licensing\", \"passives\", \"quantifiers\", \"relative_clauses\", \"scrambling\", \"subject_agreement\", \"suspended_affixation\"]\n",
        "pairs = concatenate_datasets([load_dataset(\"juletxara/turblimp\", p, split=\"train\").shuffle(seed=0).select(range(100)) for p in phenomena])\n",
        "\n",
        "@torch.no_grad()\n",
        "def sentence_logprobs(model, sentences, batch_size=64):\n",
        "    \"\"\"Total log-probability of each sentence.\"\"\"\n",
        "    model.eval()\n",
        "    tokenizer.padding_side = \"right\"\n",
        "    scores = []\n",
        "    for i in range(0, len(sentences), batch_size):\n",
        "        inputs = tokenizer([tokenizer.bos_token + s for s in sentences[i:i + batch_size]], return_tensors=\"pt\", padding=True, add_special_tokens=False).to(model.device)\n",
        "        logits = model(**inputs, use_cache=False).logits.float()\n",
        "        token_logprobs = torch.log_softmax(logits[:, :-1], dim=-1).gather(-1, inputs[\"input_ids\"][:, 1:, None])[..., 0]\n",
        "        scores += (token_logprobs * inputs[\"attention_mask\"][:, 1:]).sum(-1).tolist()\n",
        "    return scores\n",
        "\n",
        "def grammar_test(model):\n",
        "    \"\"\"For each pair, whether the model prefers the grammatical sentence.\"\"\"\n",
        "    good = sentence_logprobs(model, list(pairs[\"sentence_good\"]))\n",
        "    bad = sentence_logprobs(model, list(pairs[\"sentence_bad\"]))\n",
        "    return [g > b for g, b in zip(good, bad)]\n",
        "\n",
        "def by_phenomenon(passed):\n",
        "    return {p: sum(ok for ok, q in zip(passed, pairs[\"linguistic_phenomenon\"]) if q == p) / sum(q == p for q in pairs[\"linguistic_phenomenon\"]) for p in sorted(set(pairs[\"linguistic_phenomenon\"]))}\n",
        "\n",
        "original_passed = grammar_test(model)\n",
        "print(f\"📏 TurBLiMP accuracy before training: {sum(original_passed) / len(pairs):.1%}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8953e4ea",
      "metadata": {
        "id": "8953e4ea"
      },
      "source": [
        "# 🎯 Continued Pre-Training (CPT + LoRA)\n",
        "\n",
        "## Wrap the model with PEFT\n",
        "\n",
        "For CPT we use a higher LoRA rank than for SFT, because the model has more to learn. We adapt the attention projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`), the gated short convolutions (`in_proj`, `out_proj`), and the MLP (`w1`, `w2`, `w3`). LFM3's attention layers also have a per-head output gate, `gate_proj`; adding it to the targets is a reasonable choice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a96dd03a",
      "metadata": {
        "id": "a96dd03a"
      },
      "outputs": [],
      "source": [
        "from peft import LoraConfig\n",
        "\n",
        "peft_config = LoraConfig(\n",
        "    r=64,\n",
        "    lora_alpha=64,\n",
        "    lora_dropout=0.0,\n",
        "    target_modules=[\n",
        "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"in_proj\", \"out_proj\", \"w1\", \"w2\", \"w3\",\n",
        "        # \"gate_proj\",  # optional: LFM3's per-head attention gate\n",
        "    ],\n",
        "    task_type=\"CAUSAL_LM\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7507ea1a",
      "metadata": {
        "id": "7507ea1a"
      },
      "source": [
        "## Configure training\n",
        "\n",
        "The dataset has a `text` column, so `SFTTrainer` trains on every token, like pre-training. We also enable packing, as in pre-training: `packing_strategy=\"wrapped\"` concatenates all the articles, separated by the end-of-sequence token, and cuts the stream into rows of exactly 2,048 tokens. No compute is wasted on padding, and long articles are split across rows instead of truncated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5a9725a7",
      "metadata": {
        "id": "5a9725a7"
      },
      "outputs": [],
      "source": [
        "from trl import SFTConfig, SFTTrainer\n",
        "\n",
        "cpt_config = SFTConfig(\n",
        "    output_dir=\"./LFM3-1.7B-CPT-tr\",\n",
        "    num_train_epochs=1,\n",
        "    per_device_train_batch_size=2,\n",
        "    gradient_accumulation_steps=4,  # <- 8 rows of 2,048 tokens per step\n",
        "    learning_rate=2e-4,\n",
        "    lr_scheduler_type=\"cosine\",\n",
        "    warmup_steps=0.05,  # <- a float is a fraction of the total steps\n",
        "    max_length=2048,\n",
        "    packing=True,\n",
        "    packing_strategy=\"wrapped\",\n",
        "    dataset_text_field=\"text\",\n",
        "    logging_steps=10,\n",
        "    eval_strategy=\"steps\",\n",
        "    eval_steps=100,\n",
        "    per_device_eval_batch_size=2,\n",
        "    save_strategy=\"no\",\n",
        "    bf16=use_bf16,\n",
        "    fp16=not use_bf16,  # <- float16 mixed precision with loss scaling on GPUs without bfloat16\n",
        "    report_to=\"none\",\n",
        ")\n",
        "\n",
        "print(\"🏗️  Creating CPT trainer...\")\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    args=cpt_config,\n",
        "    train_dataset=train_dataset,\n",
        "    eval_dataset=eval_dataset,\n",
        "    processing_class=tokenizer,\n",
        "    peft_config=peft_config,\n",
        ")\n",
        "trainer.model.print_trainable_parameters()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ef7fc19c",
      "metadata": {
        "id": "ef7fc19c"
      },
      "source": [
        "## Evaluate the model before training\n",
        "\n",
        "Perplexity is the exponential of the average loss per token: lower means the model predicts the held-out Turkish text better."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d197b0e3",
      "metadata": {
        "id": "d197b0e3"
      },
      "outputs": [],
      "source": [
        "baseline = trainer.evaluate()\n",
        "baseline_ood = trainer.evaluate(eval_dataset=ood_dataset, metric_key_prefix=\"ood\")\n",
        "print(f\"📏 Perplexity before training: {math.exp(baseline['eval_loss']):.2f} (Wikipedia), {math.exp(baseline_ood['ood_loss']):.2f} (web pages)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0a33b515",
      "metadata": {
        "id": "0a33b515"
      },
      "source": [
        "## Launch training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3b73adda",
      "metadata": {
        "id": "3b73adda"
      },
      "outputs": [],
      "source": [
        "print(\"🚀 Starting continued pre-training...\")\n",
        "trainer.train()\n",
        "print(\"🎉 Continued pre-training completed!\")\n",
        "print(f\"🧮 Peak GPU memory: {torch.cuda.max_memory_reserved() / 1e9:.1f} GB\")\n",
        "\n",
        "final = trainer.evaluate()\n",
        "final_ood = trainer.evaluate(eval_dataset=ood_dataset, metric_key_prefix=\"ood\")\n",
        "print(f\"📏 Perplexity before training: {math.exp(baseline['eval_loss']):.2f} (Wikipedia), {math.exp(baseline_ood['ood_loss']):.2f} (web pages)\")\n",
        "print(f\"📏 Perplexity after training:  {math.exp(final['eval_loss']):.2f} (Wikipedia), {math.exp(final_ood['ood_loss']):.2f} (web pages)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7a6da81a",
      "metadata": {
        "id": "7a6da81a"
      },
      "source": [
        "# 🔍 Original vs fine-tuned model\n",
        "\n",
        "We run the grammar test again with the trained model, and look at pairs where training changed the model's choice."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a026048d",
      "metadata": {
        "id": "a026048d"
      },
      "outputs": [],
      "source": [
        "finetuned_passed = grammar_test(trainer.model)\n",
        "original_scores, finetuned_scores = by_phenomenon(original_passed), by_phenomenon(finetuned_passed)\n",
        "\n",
        "# Pairs the original model got wrong and the fine-tuned model gets right, in the phenomena that improved most\n",
        "most_improved = sorted(phenomena, key=lambda p: finetuned_scores[p] - original_scores[p], reverse=True)[:4]\n",
        "for p in most_improved:\n",
        "    i = next(i for i, (o, f, q) in enumerate(zip(original_passed, finetuned_passed, pairs[\"linguistic_phenomenon\"])) if q == p and not o and f)\n",
        "    print(f\"\\n🧩 {p.replace('_', ' ')}: {original_scores[p]:.0%} -> {finetuned_scores[p]:.0%}\")\n",
        "    print(f\"   ✅ grammatical:   {pairs[i]['sentence_good']}\")\n",
        "    print(f\"   ❌ ungrammatical: {pairs[i]['sentence_bad']}\")\n",
        "    print(f\"   Original prefers the ungrammatical one, the fine-tuned model prefers the grammatical one.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "72fa4ea9",
      "metadata": {
        "id": "72fa4ea9"
      },
      "source": [
        "And the summary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4e741848",
      "metadata": {
        "id": "4e741848"
      },
      "outputs": [],
      "source": [
        "print(f\"{'':<34}{'original':>10}{'fine-tuned':>12}\")\n",
        "print(f\"{'perplexity, Wikipedia (held out)':<34}{math.exp(baseline['eval_loss']):>10.2f}{math.exp(final['eval_loss']):>12.2f}\")\n",
        "print(f\"{'perplexity, web pages (OOD)':<34}{math.exp(baseline_ood['ood_loss']):>10.2f}{math.exp(final_ood['ood_loss']):>12.2f}\")\n",
        "print(f\"{'TurBLiMP grammar accuracy':<34}{sum(original_passed) / len(pairs):>10.1%}{sum(finetuned_passed) / len(pairs):>12.1%}\")\n",
        "for p in phenomena:\n",
        "    print(f\"{'  ' + p.replace('_', ' '):<34}{original_scores[p]:>10.0%}{finetuned_scores[p]:>12.0%}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bc99e328",
      "metadata": {
        "id": "bc99e328"
      },
      "source": [
        "# 💾 Save and merge the model\n",
        "\n",
        "We save the LoRA adapter, then merge it into the base weights. `save_pretrained` also copies LFM3's modeling code next to the weights, so the merged model loads with `trust_remote_code=True` like the original one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "364f76d8",
      "metadata": {
        "id": "364f76d8"
      },
      "outputs": [],
      "source": [
        "trainer.save_model(\"./LFM3-1.7B-CPT-tr-lora\")\n",
        "print(\"💾 LoRA adapter saved to: ./LFM3-1.7B-CPT-tr-lora\")\n",
        "\n",
        "output_dir = \"./LFM3-1.7B-CPT-tr-merged\"\n",
        "print(\"🔄 Merging LoRA weights...\")\n",
        "merged_model = trainer.model.merge_and_unload()\n",
        "merged_model.save_pretrained(output_dir)\n",
        "tokenizer.save_pretrained(output_dir)\n",
        "print(f\"💾 Merged model saved to: {output_dir}\")\n",
        "\n",
        "# To push the merged model to the Hugging Face Hub, uncomment:\n",
        "# merged_model.push_to_hub(\"your-username/LFM3-1.7B-CPT-tr\")\n",
        "# tokenizer.push_to_hub(\"your-username/LFM3-1.7B-CPT-tr\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8bef0d87",
      "metadata": {
        "id": "8bef0d87"
      },
      "source": [
        "# 🚀 Next steps\n",
        "\n",
        "- **Use more data**: CPT benefits from far more text than this demo. Increase the number of articles, or stream your own corpus with `load_dataset(\"json\", data_files=..., streaming=True)`.\n",
        "- **Teach the model to follow instructions in the new domain**: after CPT, run a short SFT on conversations from that domain (see [`lfm3_sft_with_trl.ipynb`](./lfm3_sft_with_trl.ipynb)), starting from the merged checkpoint.\n",
        "- **Train the embeddings too**: for a new language with many unseen words, also training the input embeddings helps."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}