\n",
@@ -733,13 +740,17 @@
" \n",
" \n",
"
"
+ ],
+ "text/plain": [
+ ""
]
},
- "metadata": {}
+ "metadata": {},
+ "output_type": "display_data"
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π SFT training completed!\n",
"πΎ Saving to: lfm2-vl-med\n"
@@ -798,7 +809,6 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "id": "QHB1ACVMzFZ5",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 84,
@@ -816,33 +826,34 @@
"ba5474abed8646f0997c998eb60c6584"
]
},
+ "id": "QHB1ACVMzFZ5",
"outputId": "6a91c28c-eac7-4ed9-ba89-d9ffb78595d6"
},
"outputs": [
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π Merging LoRA weights...\n"
]
},
{
- "output_type": "display_data",
"data": {
- "text/plain": [
- "Writing model shards: 0%| | 0/1 [00:00=0.18.2 peft>=0.15.2",
+ "execution_count": null,
"metadata": {
"id": "3FIcp_wo9nsR"
},
- "execution_count": null,
- "outputs": []
+ "outputs": [],
+ "source": [
+ "!uv pip install transformers==4.54.0 trl>=0.18.2 peft>=0.15.2"
+ ]
},
{
"cell_type": "markdown",
- "source": [
- "Let's now verify the packages are installed correctly"
- ],
"metadata": {
"id": "41UEf1uxCd6m"
- }
+ },
+ "source": [
+ "Let's now verify the packages are installed correctly"
+ ]
},
{
"cell_type": "code",
- "source": [
- "import torch\n",
- "import transformers\n",
- "import trl\n",
- "import os\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__}\")"
- ],
+ "execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
@@ -86,71 +62,60 @@
"id": "bSJgYtHT_Os4",
"outputId": "d483d722-d85a-4de3-d266-347c75abbfce"
},
- "execution_count": null,
"outputs": [
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π¦ PyTorch version: 2.6.0+cu124\n",
"π€ Transformers version: 4.54.0\n",
"π TRL version: 0.19.1\n"
]
}
+ ],
+ "source": [
+ "import torch\n",
+ "import transformers\n",
+ "import trl\n",
+ "import os\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__}\")"
]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "v_uXLzxQ_rnK"
+ },
"source": [
"# Loading the model from Transformers π€\n",
"\n"
- ],
- "metadata": {
- "id": "v_uXLzxQ_rnK"
- }
+ ]
},
{
"cell_type": "code",
- "source": [
- "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
- "import torch\n",
- "\n",
- "model_name = \"LiquidAI/LFM2.5-1.2B-Instruct\"\n",
- "\n",
- "print(\"π Loading tokenizer...\")\n",
- "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
- "\n",
- "print(\"π§ Loading model...\")\n",
- "model = AutoModelForCausalLM.from_pretrained(\n",
- " model_name,\n",
- " device_map=\"auto\",\n",
- " torch_dtype=\"auto\",\n",
- ")\n",
- "\n",
- "print(\"β
Local model loaded successfully!\")\n",
- "print(f\"π’ Parameters: {model.num_parameters():,}\")\n",
- "print(f\"π Vocab size: {len(tokenizer)}\")\n",
- "print(f\"πΎ Model size: ~{model.num_parameters() * 2 / 1e9:.1f} GB (bfloat16)\")"
- ],
+ "execution_count": null,
"metadata": {
- "id": "iA3erKM4-HhS",
"colab": {
"base_uri": "https://localhost:8080/"
},
+ "id": "iA3erKM4-HhS",
"outputId": "e19261f7-5e96-4756-e001-80af1d9d37b5"
},
- "execution_count": null,
"outputs": [
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π Loading tokenizer...\n"
]
},
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
"The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
@@ -161,8 +126,8 @@
]
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π§ Loading model...\n",
"β
Local model loaded successfully!\n",
@@ -171,53 +136,56 @@
"πΎ Model size: ~2.3 GB (bfloat16)\n"
]
}
+ ],
+ "source": [
+ "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+ "import torch\n",
+ "\n",
+ "model_name = \"LiquidAI/LFM2.5-1.2B-Instruct\"\n",
+ "\n",
+ "print(\"π Loading tokenizer...\")\n",
+ "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+ "\n",
+ "print(\"π§ Loading model...\")\n",
+ "model = AutoModelForCausalLM.from_pretrained(\n",
+ " model_name,\n",
+ " device_map=\"auto\",\n",
+ " torch_dtype=\"auto\",\n",
+ ")\n",
+ "\n",
+ "print(\"β
Local model loaded successfully!\")\n",
+ "print(f\"π’ Parameters: {model.num_parameters():,}\")\n",
+ "print(f\"π Vocab size: {len(tokenizer)}\")\n",
+ "print(f\"πΎ Model size: ~{model.num_parameters() * 2 / 1e9:.1f} GB (bfloat16)\")"
]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "o83NqFPNA_nk"
+ },
"source": [
"# π― Direct Preference Optimization (DPO + LoRA)\n",
"\n",
"DPO aligns the model with human preferences by learning from preference pairs (chosen vs rejected responses). This typically follows SFT training.\n",
"\n",
"DPO might be too compute heavy if you're running on one of the free-tier colab GPUs. Hence we use LoRA (Low-Rank Adaptation) to finetune the model by only training a small number of additional parameters. Perfect for limited compute resources!"
- ],
- "metadata": {
- "id": "o83NqFPNA_nk"
- }
+ ]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "0w-40n_XJL9H"
+ },
"source": [
"## Load a DPO Dataset\n",
"\n",
"We will use [mlabonne/orpo-dpo-mix-40k](https://huggingface.co/datasets/mlabonne/orpo-dpo-mix-40k), limiting ourselves to the first 2k samples for brevity. Feel free to change the limit by changing the slicing index in the parameter `split`. The size of the validation data can be adjusted by changing `test_size`."
- ],
- "metadata": {
- "id": "0w-40n_XJL9H"
- }
+ ]
},
{
"cell_type": "code",
- "source": [
- "from datasets import load_dataset\n",
- "\n",
- "print(\"π₯ Loading DPO dataset...\")\n",
- "\n",
- "dataset_dpo = load_dataset(\"mlabonne/orpo-dpo-mix-40k\", split=\"train[:2000]\")\n",
- "dataset_dpo = dataset_dpo.train_test_split(test_size=0.1, seed=42)\n",
- "train_dataset_dpo, eval_dataset_dpo = dataset_dpo['train'], dataset_dpo['test']\n",
- "\n",
- "print(\"β
DPO Dataset loaded:\")\n",
- "print(f\" π Train samples: {len(train_dataset_dpo)}\")\n",
- "print(f\" π§ͺ Eval samples: {len(eval_dataset_dpo)}\")\n",
- "\n",
- "sample = train_dataset_dpo[0]\n",
- "print(\"\\nπ Single Sample:\")\n",
- "print(f\" Prompt: {sample['prompt'][:100]}...\")\n",
- "print(f\" β
Chosen: {sample['chosen'][:100]}...\")\n",
- "print(f\" β Rejected: {sample['rejected'][:100]}...\")"
- ],
+ "execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
@@ -225,11 +193,10 @@
"id": "Ajq6EABUBAv_",
"outputId": "1b6e9f8f-e7a7-472d-e50b-b40b450d336e"
},
- "execution_count": null,
"outputs": [
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π₯ Loading DPO dataset...\n",
"β
DPO Dataset loaded:\n",
@@ -242,10 +209,32 @@
" β Rejected: [{'content': 'Classify the following instruments into their respective families (brass, strings, woodwinds, or percussion): Didgeridoo, Cuica, Euphonium, Guzheng', 'role': 'user'}, {'content': 'The classification of the mentioned instruments is as follows:\\n\\n1. Didgeridoo: Brass\\n2. Cuica: Percussion\\n3. Euphonium: Brass\\n4. Guzheng: Strings', 'role': 'assistant'}, {'content': 'Explain the construction and sound production process of each instrument mentioned, and how these processes relate to their classification in their respective families.', 'role': 'user'}, {'content': \"1. Didgeridoo: The didgeridoo is traditionally made from eucalyptus trees which have been naturally hollowed out by termites. The mouthpiece can be made of beeswax or shaped directly from the wood. The player blows into the instrument without interruption while using a special breathing technique called circular breathing. The vibration of the lips and the resonances in the tube produce the sound. Although the didgeridoo is not a typical brass instrument, it is often classified as such due to the lip-vibrated method of sound production, similar to a brass instrument.\\n\\n2. Cuica: The cuica is a Brazilian friction drum with a large pitch range, produced by changing tension on the head of the drum. It has a stick fixed in the center of the drum head, and the player rubs this stick with a wet cloth and simultaneously modulates the pitch by pressing the thumb on the drum head. This results in a sound that can be used to mimic a variety of effects, often the call of wild animals. It is classified as a percussion instrument due to the way the sound is produced by striking or rubbing.\\n\\n3. Euphonium: The euphonium is a conical-bore, tenor-voiced brass instrument. The mouthpiece of the euphonium is larger than that of a trumpet, and this, along with the conical bore, contributes to its rich, dark tone. The player buzzes their lips into the mouthpiece, creating vibrations that travel through the instrument and are amplified by the bell. The pitch can be altered by changing the lip tension and by using the valves to change the length of the tubing. Its classification as a brass instrument is due to the material it's made from and the method of sound production.\\n\\n4. Guzheng: The guzheng, also known as the Chinese zither, is a string instrument with movable bridges and 21 strings. It is made of a large, resonant wooden body, over which the strings are stretched. The player plucks the strings with their fingers, often with the aid of plectra attached to four of the fingers on the right or both hands. The left hand can also press the strings to the left of the bridges to change the pitch or produce vibrato. It is classified as a string instrument because the sound is produced by vibrating strings.\", 'role': 'assistant'}, {'content': 'Discuss the technique of circular breathing used in playing the Didgeridoo and its impact on the sound produced.', 'role': 'user'}, {'content': \"Circular breathing is a technique used by players of some wind instruments to produce a continuous tone without interruption. This technique is used extensively in playing the didgeridoo. It involves inhaling through the nose while maintaining airflow through the instrument, using air stored in the cheeks.\\n\\nHere's a basic breakdown of how it works:\\n\\n1. Fill your cheeks with air and begin to play the instrument.\\n2. While you're blowing this air out through your mouth into the instrument, start to breathe in through your nose.\\n3. As your lungs fill with air, continue to push the air out of your mouth using your cheeks and tongue.\\n4. Once your cheeks are empty, switch back to blowing air from your lungs into the instrument, and refill your cheeks.\\n\\nThis technique allows the player to sustain a drone for a long period of time, which is characteristic of didgeridoo music. It creates a continuous, rhythmic sound and allows for the layering of complex rhythms and sound effects. The circular breathing technique contributes to the mesmerizing, trance-like quality of the didgeridoo's sound. It also allows for the performance of more complex pieces that require extended periods of sound without breaks for breath.\", 'role': 'assistant'}, {'content': 'Can you elaborate on the physiological and mental challenges a musician might face while learning and mastering the technique of circular breathing?', 'role': 'user'}, {'content': \"Learning and mastering the technique of circular breathing can pose numerous physiological and mental challenges for instrumentalists. Here are some of the challenges a musician may face:\\n\\n1. Coordination: Circular breathing requires a high degree of muscular and mental coordination between respiratory, oral, and facial muscles. This includes skills like inhalation through the nose, simultaneous sound production, and control over airflow from the lungs and cheeks. The learning curve can be steep, and several hours of practice may be required to achieve the basic level of proficiency.\\n\\n2. Breath control: Attaining consistent breath control is crucial in circular breathing. A musician must be able to pulse and maintain a steady airflow through the nose while precisely regulating the rate and amount of air entering the lungs. This requires a high level of breath control and lung capacity.\\n\\n3. Cheek strength and control: Circular breathing demands significant strength and control of the cheek muscles. The cheeks must act as a second lung, storing air while the player inhales through the nose and releasing it back into the instrument when blowing from the lungs. This requires practice and muscle conditioning to build the endurance necessary for prolonged circular breathing sessions.\\n\\n4. Embouchure: Keeping a stable embouchure β the position and tension of the lips and mouth β is important while incorporating the circular breathing technique. It may take time to adjust to this combination and prevent air loss or a change in the quality of the sound while switching from blowing from the lungs to inhaling through the nose.\\n\\n5. Mental focus and concentration: Circular breathing involves sustaining multiple tasks simultaneously and requires a high concentration level. Musicians must maintain focus on the sound production while switching between inhaling and exhaling, ensuring an uninterrupted airflow and sound throughout the performance.\\n\\nIn addition to these physiological challenges, mastering circular breathing can also have mental benefits. It can enhance an individual's mindfulness, concentration, and breath awareness, leading to improved performance in various aspects of life. However, the learning process may be frustrating at times due to the high level of coordination and practice involved. Overcoming these challenges can be rewarding and offer a sense of accomplishment as the musician achieves greater proficiency and control in the circular breathing technique.\", 'role': 'assistant'}]...\n"
]
}
+ ],
+ "source": [
+ "from datasets import load_dataset\n",
+ "\n",
+ "print(\"π₯ Loading DPO dataset...\")\n",
+ "\n",
+ "dataset_dpo = load_dataset(\"mlabonne/orpo-dpo-mix-40k\", split=\"train[:2000]\")\n",
+ "dataset_dpo = dataset_dpo.train_test_split(test_size=0.1, seed=42)\n",
+ "train_dataset_dpo, eval_dataset_dpo = dataset_dpo['train'], dataset_dpo['test']\n",
+ "\n",
+ "print(\"β
DPO Dataset loaded:\")\n",
+ "print(f\" π Train samples: {len(train_dataset_dpo)}\")\n",
+ "print(f\" π§ͺ Eval samples: {len(eval_dataset_dpo)}\")\n",
+ "\n",
+ "sample = train_dataset_dpo[0]\n",
+ "print(\"\\nπ Single Sample:\")\n",
+ "print(f\" Prompt: {sample['prompt'][:100]}...\")\n",
+ "print(f\" β
Chosen: {sample['chosen'][:100]}...\")\n",
+ "print(f\" β Rejected: {sample['rejected'][:100]}...\")"
]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "v8or_m-UdKGc"
+ },
"source": [
"## Wrap the model with PEFT\n",
"\n",
@@ -254,13 +243,31 @@
"- lower -> needs even less compute resources\n",
"\n",
"You can skip this part if you have a premium GPU and want to go for a full finetune."
- ],
- "metadata": {
- "id": "v8or_m-UdKGc"
- }
+ ]
},
{
"cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "z7x5RhC5eLmc",
+ "outputId": "f6553df2-ba07-4839-9c73-2c55dec64495"
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "trainable params: 5,554,176 || all params: 1,175,894,784 || trainable%: 0.4723\n",
+ "β
LoRA configuration applied!\n",
+ "ποΈ LoRA rank: 8\n",
+ "π LoRA alpha: 16\n",
+ "π― Target modules: {'w2', 'out_proj', 'in_proj', 'q_proj', 'k_proj', 'w3', 'v_proj', 'w1'}\n"
+ ]
+ }
+ ],
"source": [
"from peft import LoraConfig, get_peft_model, TaskType\n",
"\n",
@@ -286,126 +293,65 @@
"print(f\"ποΈ LoRA rank: {lora_config.r}\")\n",
"print(f\"π LoRA alpha: {lora_config.lora_alpha}\")\n",
"print(f\"π― Target modules: {lora_config.target_modules}\")"
- ],
- "metadata": {
- "colab": {
- "base_uri": "https://localhost:8080/"
- },
- "id": "z7x5RhC5eLmc",
- "outputId": "f6553df2-ba07-4839-9c73-2c55dec64495"
- },
- "execution_count": null,
- "outputs": [
- {
- "output_type": "stream",
- "name": "stdout",
- "text": [
- "trainable params: 5,554,176 || all params: 1,175,894,784 || trainable%: 0.4723\n",
- "β
LoRA configuration applied!\n",
- "ποΈ LoRA rank: 8\n",
- "π LoRA alpha: 16\n",
- "π― Target modules: {'w2', 'out_proj', 'in_proj', 'q_proj', 'k_proj', 'w3', 'v_proj', 'w1'}\n"
- ]
- }
]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "yxuxfHu7Jopx"
+ },
"source": [
"## Launch Training\n",
"\n",
"We are now ready to launch a DPO run with `DPOTrainer`, feel free to modify `DPOConfig` to play around with different configurations.\n",
"\n",
"Replace `lora_model` with `model` if you have a premium-tier in colab and want to run a full finetune."
- ],
- "metadata": {
- "id": "yxuxfHu7Jopx"
- }
+ ]
},
{
"cell_type": "code",
- "source": [
- "from trl import DPOConfig, DPOTrainer\n",
- "\n",
- "# DPO Training configuration\n",
- "dpo_config = DPOConfig(\n",
- " output_dir=\"./lfm2-dpo\",\n",
- " num_train_epochs=1,\n",
- " per_device_train_batch_size=1,\n",
- " learning_rate=1e-6,\n",
- " lr_scheduler_type=\"linear\",\n",
- " logging_steps=10,\n",
- " save_strategy=\"epoch\",\n",
- " eval_strategy=\"epoch\",\n",
- " bf16=False # <- not all colab GPUs support bf16\n",
- ")\n",
- "\n",
- "# Create DPO trainer\n",
- "print(\"ποΈ Creating DPO trainer...\")\n",
- "dpo_trainer = DPOTrainer(\n",
- " model=lora_model,\n",
- " args=dpo_config,\n",
- " train_dataset=train_dataset_dpo,\n",
- " eval_dataset=eval_dataset_dpo,\n",
- " processing_class=tokenizer,\n",
- ")\n",
- "\n",
- "# Start DPO training\n",
- "print(\"\\nπ Starting DPO training...\")\n",
- "dpo_trainer.train()\n",
- "\n",
- "print(\"π DPO training completed!\")\n",
- "\n",
- "# Save the DPO model\n",
- "dpo_trainer.save_model()\n",
- "print(f\"πΎ DPO model saved to: {dpo_config.output_dir}\")"
- ],
+ "execution_count": null,
"metadata": {
- "id": "nGBgkf9LBHp6",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 263
},
+ "id": "nGBgkf9LBHp6",
"outputId": "5242d3ac-16e2-4ff3-c4b6-fdfb32c6cf23"
},
- "execution_count": null,
"outputs": [
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n",
"Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
]
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"ποΈ Creating DPO trainer...\n"
]
},
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"No label_names provided for model class `PeftModelForCausalLM`. Since `PeftModel` hides base models input arguments, if label_names is not given, label_names can't be set automatically within `Trainer`. Note that empty label_names list will be used instead.\n"
]
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"\n",
"π Starting DPO training...\n"
]
},
{
- "output_type": "display_data",
"data": {
- "text/plain": [
- ""
- ],
"text/html": [
"\n",
" \n",
@@ -445,66 +391,122 @@
" \n",
" \n",
"
"
+ ],
+ "text/plain": [
+ ""
]
},
- "metadata": {}
+ "metadata": {},
+ "output_type": "display_data"
},
{
- "output_type": "stream",
"name": "stderr",
+ "output_type": "stream",
"text": [
"Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
]
},
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"π DPO training completed!\n",
"πΎ DPO model saved to: ./lfm2-dpo\n"
]
}
+ ],
+ "source": [
+ "from trl import DPOConfig, DPOTrainer\n",
+ "\n",
+ "# DPO Training configuration\n",
+ "dpo_config = DPOConfig(\n",
+ " output_dir=\"./lfm2-dpo\",\n",
+ " num_train_epochs=1,\n",
+ " per_device_train_batch_size=1,\n",
+ " learning_rate=1e-6,\n",
+ " lr_scheduler_type=\"linear\",\n",
+ " logging_steps=10,\n",
+ " save_strategy=\"epoch\",\n",
+ " eval_strategy=\"epoch\",\n",
+ " bf16=False # <- not all colab GPUs support bf16\n",
+ ")\n",
+ "\n",
+ "# Create DPO trainer\n",
+ "print(\"ποΈ Creating DPO trainer...\")\n",
+ "dpo_trainer = DPOTrainer(\n",
+ " model=lora_model,\n",
+ " args=dpo_config,\n",
+ " train_dataset=train_dataset_dpo,\n",
+ " eval_dataset=eval_dataset_dpo,\n",
+ " processing_class=tokenizer,\n",
+ ")\n",
+ "\n",
+ "# Start DPO training\n",
+ "print(\"\\nπ Starting DPO training...\")\n",
+ "dpo_trainer.train()\n",
+ "\n",
+ "print(\"π DPO training completed!\")\n",
+ "\n",
+ "# Save the DPO model\n",
+ "dpo_trainer.save_model()\n",
+ "print(f\"πΎ DPO model saved to: {dpo_config.output_dir}\")"
]
},
{
"cell_type": "markdown",
+ "metadata": {
+ "id": "WIrCmjUsfNIg"
+ },
"source": [
"## Save merged model\n",
"\n",
"If you have used LoRA. Merge the extra weights learned with LoRA back into the model to obtain a \"normal\" model checkpoint."
- ],
- "metadata": {
- "id": "WIrCmjUsfNIg"
- }
+ ]
},
{
"cell_type": "code",
- "source": [
- "print(\"\\nπ Merging LoRA weights...\")\n",
- "merged_model = lora_model.merge_and_unload()\n",
- "merged_model.save_pretrained(\"./lfm2-lora-merged\")\n",
- "tokenizer.save_pretrained(\"./lfm2-lora-merged\")\n",
- "print(\"πΎ Merged model saved to: ./lfm2-lora-merged\")"
- ],
+ "execution_count": null,
"metadata": {
- "id": "RUJoU3BwfPTg",
"colab": {
"base_uri": "https://localhost:8080/"
},
+ "id": "RUJoU3BwfPTg",
"outputId": "45e09b8f-9089-4237-ffff-c4d59db6ff89"
},
- "execution_count": null,
"outputs": [
{
- "output_type": "stream",
"name": "stdout",
+ "output_type": "stream",
"text": [
"\n",
"π Merging LoRA weights...\n",
"πΎ Merged model saved to: ./lfm2-lora-merged\n"
]
}
+ ],
+ "source": [
+ "print(\"\\nπ Merging LoRA weights...\")\n",
+ "merged_model = lora_model.merge_and_unload()\n",
+ "merged_model.save_pretrained(\"./lfm2-lora-merged\")\n",
+ "tokenizer.save_pretrained(\"./lfm2-lora-merged\")\n",
+ "print(\"πΎ Merged model saved to: ./lfm2-lora-merged\")"
]
}
- ]
-}
\ No newline at end of file
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "gpuType": "T4",
+ "provenance": []
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/util/README.md b/util/README.md
new file mode 100644
index 0000000..8924740
--- /dev/null
+++ b/util/README.md
@@ -0,0 +1,83 @@
+# Notebook Testing Utilities
+
+This directory contains tools for running Jupyter notebooks as smoke tests on [Modal](https://modal.com/) GPUs. The system extracts code cells from `.ipynb` files, collects their dependencies, and executes them remotely to verify they work end-to-end.
+
+## Overview
+
+| File | Purpose |
+|---|---|
+| `modal_runner.py` | Modal app that runs Python code on a remote GPU container |
+| `run_notebook_test.py` | CLI that parses a notebook, extracts dependencies, and submits the combined script to Modal |
+
+## How it works
+
+1. **`run_notebook_test.py`** reads a notebook and extracts all code cells.
+2. Pip install lines (`!pip install ...`, `!uv pip install ...`) are extracted from **every** cell, collected as image-level dependencies, and stripped from the script so they don't run at execution time.
+3. Shell commands (`!` lines) are converted to `subprocess.run()` calls; `%%bash` cells become setup commands.
+4. The remaining cells are concatenated into a single Python script, preserving shared state (variables, imports) across cells just like a real notebook.
+5. The script, dependency list, and setup commands are sent to **`modal_runner.py`**'s `run_code` function via Modal's remote invocation.
+
+**`modal_runner.py`** runs on an A10G GPU (default) with a 10-minute timeout. It:
+- Installs pip packages in order (each `pip install` line becomes a separate install group to preserve ordering)
+- Runs any setup commands (shell commands, `%%bash` cells)
+- Executes the combined Python script as a subprocess and captures stdout/stderr
+
+## Usage
+
+### Prerequisites
+
+- A Modal account with a deployed `ci-runner` app
+- Modal token configured locally
+
+### Deploy the Modal app (one-time)
+
+```bash
+modal deploy util/modal_runner.py
+```
+
+### Run a notebook
+
+```bash
+python util/run_notebook_test.py --notebook notebooks/MyNotebook.ipynb
+```
+
+### Options
+
+| Flag | Description | Default |
+|---|---|---|
+| `--notebook` | Path to `.ipynb` file (required) | β |
+| `--gpu` | GPU type | `A10G` |
+| `--dry-run` | Print the combined script and packages without running | off |
+| `--skip-packages PKG ...` | Package names to exclude from installation | none |
+
+### Examples
+
+```bash
+# Dry run to inspect what would be executed
+python util/run_notebook_test.py --notebook notebooks/LFM2_Inference.ipynb --dry-run
+
+# Skip a package that doesn't build in CI
+python util/run_notebook_test.py --notebook notebooks/LFM2_Inference.ipynb --skip-packages flash-attn
+
+# Use a different GPU
+python util/run_notebook_test.py --notebook notebooks/LFM2_Inference.ipynb --gpu A100
+```
+
+## Notebook directives
+
+Control which cells and lines are included in the test run:
+
+| Directive | Scope | Effect |
+|---|---|---|
+| `# test:skip` | Cell | Skip the entire cell |
+| `# !modal_skip` | Line or cell | Skip that line (in dependency parsing) or cell (in cell filtering) |
+| `# !modal_skip_rest` | Line or cell | Stop processing β ignore this and all subsequent lines/cells |
+
+These can be placed in comments within any code cell.
+
+## CI workflow
+
+The GitHub Actions workflow (`.github/workflows/run-notebooks.yaml`) automatically:
+1. Discovers all `.ipynb` files in the `notebooks/` directory
+2. Runs each notebook in parallel as a separate matrix job
+3. Triggers on pushes/PRs that change `notebooks/**`, or manually via `workflow_dispatch`
diff --git a/util/modal_runner.py b/util/modal_runner.py
new file mode 100644
index 0000000..b75bada
--- /dev/null
+++ b/util/modal_runner.py
@@ -0,0 +1,62 @@
+import modal
+import subprocess, sys, os, tempfile
+
+app = modal.App("ci-runner")
+
+@app.function(
+ gpu="A10G",
+ timeout=600,
+ image=modal.Image.debian_slim(python_version="3.12")
+ .apt_install("curl", "wget", "zstd", "git")
+ .pip_install("uv", "typing_extensions>=4.14.0"),
+)
+def run_code(code: str, pip_packages: list[list[str]] = [], setup_commands: list[str] = []) -> dict:
+
+ # Disable vLLM v1 multiprocessing (incompatible with containers)
+ os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
+
+ # Restore real stdout/stderr β Modal wraps them with objects that
+ # lack fileno(), which breaks subprocesses and libraries like vLLM.
+ real_stdout = open("/dev/stdout", "w")
+ real_stderr = open("/dev/stderr", "w")
+ sys.stdout = real_stdout
+ sys.stderr = real_stderr
+
+ for group in pip_packages:
+ print(f"[ci-runner] Installing: {group}")
+ subprocess.check_call(
+ ["uv", "pip", "install", "--system", *group],
+ stdout=real_stdout, stderr=real_stderr,
+ )
+
+ for cmd in setup_commands:
+ print(f"[ci-runner] Setup: {cmd}")
+ subprocess.run(cmd, shell=True, check=True, stdout=real_stdout, stderr=real_stderr)
+
+ # Remove Modal's bundled deps from sys.path to prevent shadowing
+ clean_path = [p for p in sys.path if "/__modal/deps" not in p]
+
+ # Write code to a temp file and run as a subprocess so that
+ # stdout/stderr are real file descriptors (required by vLLM, etc.)
+ with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+ f.write(code)
+ script_path = f.name
+
+ env = os.environ.copy()
+ env["PYTHONPATH"] = ":".join(clean_path)
+
+ proc = subprocess.run(
+ [sys.executable, script_path],
+ capture_output=True,
+ text=True,
+ timeout=540,
+ )
+
+ os.unlink(script_path)
+
+ return {
+ "success": proc.returncode == 0,
+ "stdout": proc.stdout,
+ "stderr": proc.stderr,
+ "error": None if proc.returncode == 0 else f"Exit code {proc.returncode}\n{proc.stderr[-2000:] if proc.stderr else ''}",
+ }
\ No newline at end of file
diff --git a/util/run_notebook_test.py b/util/run_notebook_test.py
new file mode 100644
index 0000000..3775de8
--- /dev/null
+++ b/util/run_notebook_test.py
@@ -0,0 +1,324 @@
+#!/usr/bin/env python3
+"""
+Run code cells from a notebook on a Modal GPU as a single combined script.
+
+Pip install lines are extracted from ALL cells, collected as image dependencies,
+and stripped from the script so they don't run at execution time.
+Non-skipped cells are concatenated into one script and executed together,
+preserving shared state (variables, imports) across cells just like a real notebook.
+
+Usage:
+ python run_notebook_test.py --notebook notebooks/LFM2_Inference_with_Transformers.ipynb
+ python run_notebook_test.py --notebook notebooks/LFM2_Inference_with_Transformers.ipynb --gpu A10G
+ python run_notebook_test.py --notebook notebooks/LFM2_Inference_with_Transformers.ipynb --dry-run
+"""
+
+import argparse
+import re
+import sys
+import time
+from pathlib import Path
+import json
+import modal
+
+sys.path.insert(0, str(Path(__file__).resolve().parent / "tests"))
+
+
+def extract_code_cells(notebook_path: Path) -> list[dict]:
+ """Return a list of code cell dicts with 'source', 'index', and 'skipped' keys."""
+ with open(notebook_path) as f:
+ nb = json.load(f)
+
+ cells = []
+ code_index = 0
+ for cell in nb.get("cells", []):
+ if cell["cell_type"] != "code":
+ continue
+
+ source = "".join(cell["source"])
+ skip = (
+ "# test:skip" in source
+ or cell.get("metadata", {}).get("test_skip", False)
+ )
+
+ cells.append({
+ "index": code_index,
+ "source": source,
+ "skipped": skip,
+ })
+ code_index += 1
+
+ return cells
+
+def parse_packages_from_cell(source: str) -> tuple[list[list[str]], list[str]]:
+ """Parse dependency cell into pip package groups and shell setup commands.
+
+ Returns (pip_package_groups, setup_commands). Each ``pip install`` line
+ becomes its own group so that install ordering is preserved (e.g. packages
+ that need torch at build time can be installed in a later group).
+ Non-pip ``!`` lines become setup commands that run before the Python script.
+ Lines after ``# !modal_skip_rest`` are ignored; lines with
+ ``# !modal_skip`` are skipped individually.
+ """
+ package_groups: list[list[str]] = []
+ setup_commands = []
+ for line in source.splitlines():
+ line = line.strip()
+ # Stop processing if we hit a modal_skip_rest directive
+ if "!modal_skip_rest" in line:
+ break
+ # Skip individual lines marked with modal_skip
+ if "!modal_skip" in line:
+ continue
+ # Match: !pip install ..., !uv pip install ..., pip install ...
+ match = re.match(r"^!?\s*(?:uv\s+)?pip\s+install\s+(.+)", line)
+ if match:
+ group = []
+ for token in re.split(r"[,\s]+", match.group(1)):
+ token = token.strip('"').strip("'").strip(",")
+ if token and not token.startswith("-"):
+ group.append(token)
+ if group:
+ package_groups.append(group)
+ elif line.startswith("!"):
+ setup_commands.append(line[1:].strip())
+ return package_groups, setup_commands
+
+
+def strip_pip_lines(source: str) -> str:
+ """Remove pip install lines (and their continuations) from cell source.
+
+ This is used after packages have been extracted so that pip installs
+ run as image-level dependencies rather than at script execution time.
+ """
+ lines = source.splitlines()
+ cleaned: list[str] = []
+ i = 0
+ while i < len(lines):
+ stripped = lines[i].strip()
+ if re.match(r"^!?\s*(?:uv\s+)?pip\s+install\s+", stripped):
+ # Skip this line and any backslash-continuation lines
+ while stripped.endswith("\\") and i + 1 < len(lines):
+ i += 1
+ stripped = lines[i].strip()
+ i += 1
+ continue
+ cleaned.append(lines[i])
+ i += 1
+ return "\n".join(cleaned)
+
+
+TRAINER_PATTERNS = re.compile(r"\b(?:SFTTrainer|GRPOTrainer|DPOTrainer|Trainer)\s*\(")
+
+
+def is_training_notebook(code_cells: list[dict]) -> bool:
+ """Return True if any cell instantiates a HF Trainer."""
+ return any(TRAINER_PATTERNS.search(c["source"]) for c in code_cells)
+
+# Default to lower epoch to speed up testing pipeline
+def patch_training_epochs(source: str, epochs: float = 0.01) -> str:
+ """Replace num_train_epochs= with a minimal value for smoke testing."""
+ return re.sub(
+ r"num_train_epochs\s*=\s*[0-9.]+",
+ f"num_train_epochs={epochs}",
+ source,
+ )
+
+
+def filter_cells(code_cells: list[dict]) -> list[dict]:
+ """Apply !modal_skip and !modal_skip_rest directives.
+
+ - !modal_skip: skip only that cell
+ - !modal_skip_rest: skip that cell and all remaining cells
+ """
+ filtered = []
+ for cell in code_cells:
+ source = cell["source"]
+ if "!modal_skip_rest" in source:
+ break
+ if "!modal_skip" in source and "!modal_skip_rest" not in source:
+ continue
+ filtered.append(cell)
+ return filtered
+
+
+def preprocess_cell(source: str) -> tuple[str, list[str]]:
+ """Transform cell source so shell lines become valid Python.
+
+ Returns (transformed_source, setup_commands).
+ - ``%%bash`` cells are collected as setup commands and excluded from the script.
+ - ``!`` lines are converted to ``subprocess.run()`` calls.
+ - Pure Python lines are unchanged.
+ """
+ lines = source.splitlines()
+ setup_commands: list[str] = []
+
+ # Detect %%bash magic (may follow comment lines)
+ non_comment = [l for l in lines if l.strip() and not l.strip().startswith("#")]
+ if non_comment and non_comment[0].strip() == "%%bash":
+ bash_lines: list[str] = []
+ found_magic = False
+ for line in lines:
+ if not found_magic:
+ if line.strip() == "%%bash":
+ found_magic = True
+ continue
+ bash_lines.append(line)
+ if bash_lines:
+ setup_commands.append("\n".join(bash_lines))
+ return "", setup_commands
+
+ # Convert ! lines to subprocess.run(), joining backslash continuations
+ transformed: list[str] = []
+ i = 0
+ while i < len(lines):
+ stripped = lines[i].lstrip()
+ if stripped.startswith("!"):
+ indent = lines[i][: len(lines[i]) - len(stripped)]
+ cmd_parts = [stripped[1:]]
+ # Collect continuation lines ending with backslash
+ while cmd_parts[-1].rstrip().endswith("\\") and i + 1 < len(lines):
+ i += 1
+ cmd_parts.append(lines[i])
+ full_cmd = "\n".join(cmd_parts)
+ # Ensure uv/pip installs use --system (no venv in containers)
+ if re.match(r"(?:uv\s+)?pip\s+install", full_cmd) and "--system" not in full_cmd:
+ full_cmd = full_cmd.replace("pip install", "pip install --system", 1)
+ transformed.append(f"{indent}subprocess.run({full_cmd!r}, shell=True, check=True)")
+ else:
+ transformed.append(lines[i])
+ i += 1
+ return "\n".join(transformed), setup_commands
+
+
+def combine_cells(code_cells: list[dict]) -> tuple[str, list[str]]:
+ """Concatenate code cells into a single script with cell markers.
+
+ Returns (combined_code, setup_commands) after preprocessing each cell.
+ """
+ parts = []
+ all_setup: list[str] = []
+ has_subprocess_import = False
+
+ for cell in code_cells:
+ source, setup_cmds = preprocess_cell(cell["source"])
+ all_setup.extend(setup_cmds)
+ if source.strip():
+ if not has_subprocess_import and "subprocess.run(" in source:
+ has_subprocess_import = True
+ parts.append(f"# --- cell {cell['index']} ---")
+ parts.append(source)
+
+ code = "\n\n".join(parts)
+ # Prepend subprocess import if we generated subprocess.run() calls
+ if has_subprocess_import:
+ code = "import subprocess\n\n" + code
+ return code, all_setup
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Run notebook code cells on Modal as a single script")
+ parser.add_argument("--notebook", required=True, help="Path to .ipynb file")
+ parser.add_argument("--gpu", default="A10G", help="GPU type (default: A10G)")
+ parser.add_argument("--dry-run", action="store_true", help="Print combined script and packages without running")
+ parser.add_argument("--skip-packages", nargs="+", default=[], metavar="PKG",
+ help="Package names to exclude from installation (e.g. --skip-packages flash-attn)")
+ args = parser.parse_args()
+
+ notebook_path = Path(args.notebook)
+ if not notebook_path.exists():
+ print(f"Error: notebook not found: {notebook_path}")
+ sys.exit(1)
+
+ cells = extract_code_cells(notebook_path)
+ if not cells:
+ print("No code cells found.")
+ sys.exit(1)
+
+ # Filter out skipped / modal_skip cells
+ code_cells = [c for c in cells if not c["skipped"]]
+ code_cells = filter_cells(code_cells)
+
+ if not code_cells:
+ print(f"Notebook {notebook_path.name}: no runnable code cells (all skipped) β passing.")
+ sys.exit(0)
+
+ # Extract pip packages from ALL cells, then strip those lines from source.
+ # Only collect packages here β shell commands (! lines, %%bash) in regular
+ # cells are handled by preprocess_cell inside combine_cells.
+ pip_packages: list[list[str]] = []
+ for cell in code_cells:
+ pkgs, _setup = parse_packages_from_cell(cell["source"])
+ pip_packages.extend(pkgs)
+ cell["source"] = strip_pip_lines(cell["source"])
+
+ # Filter out skipped packages
+ if args.skip_packages:
+ skip_set = set(args.skip_packages)
+ pip_packages = [
+ [pkg for pkg in group if pkg not in skip_set]
+ for group in pip_packages
+ ]
+ pip_packages = [g for g in pip_packages if g] # drop empty groups
+
+ combined, setup_commands = combine_cells(code_cells)
+ skipped = len(cells) - len(code_cells)
+
+ # Auto-detect training notebooks: upgrade GPU and patch epochs for smoke testing
+ is_training = is_training_notebook(code_cells)
+ if is_training:
+ if args.gpu == "A10G": # only override if user didn't specify
+ args.gpu = "H100"
+ combined = patch_training_epochs(combined, epochs=0.01)
+
+ print(f"{'=' * 50}")
+ print(f"Notebook: {notebook_path.name}")
+ if is_training:
+ print(f"Type: training (auto: GPUβ{args.gpu}, epochsβ0.01)")
+ all_packages = [pkg for group in pip_packages for pkg in group]
+ print(f"Packages: {all_packages or '(none)'} ({len(pip_packages)} install group(s))")
+ if setup_commands:
+ print(f"Setup: {len(setup_commands)} command(s)")
+ print(f"GPU: {args.gpu}")
+ print(f"Cells: {len(code_cells)} to run, {skipped} skipped")
+ print(f"{'=' * 50}\n")
+
+ if args.dry_run:
+ if setup_commands:
+ print("=== SETUP COMMANDS ===")
+ for i, cmd in enumerate(setup_commands):
+ print(f"\n--- setup {i} ---")
+ print(cmd)
+ print(f"\n=== PYTHON SCRIPT ===\n")
+ print(combined)
+ return
+
+ try:
+ run_code = modal.Function.from_name("ci-runner", "run_code")
+ except modal.exception.NotFoundError:
+ print("Error: Modal app 'ci-runner' not deployed.")
+ print("Run this first: modal deploy modal_runner.py")
+ sys.exit(1)
+
+ print("Submitting to Modal...")
+ start_time = time.time()
+ result = run_code.remote(code=combined, pip_packages=pip_packages, setup_commands=setup_commands)
+
+ print("\n--- STDOUT ---")
+ print(result["stdout"] or "(empty)")
+
+ if result["stderr"]:
+ print("\n--- STDERR ---")
+ print(result["stderr"])
+
+ if result["error"]:
+ print(f"\nβ FAILED: {result['error']}")
+ sys.exit(1)
+
+ elapsed = time.time() - start_time
+ print(f"\nRuntime: {elapsed:.2f}s")
+ print(f"β All {len(code_cells)} cells passed")
+
+
+if __name__ == "__main__":
+ main()