vlm support (#5)

Author: EdoardoMosca

README.md

@@ -58,8 +58,64 @@ Unless you overwrote `output_dir`, results will be stored in `outputs/training_t
 
 When training is done, you can bundle your output checkpoint with `leap-bundle` to use it directly within LEAP. Checkout our [Quick Start guide](https://leap.liquid.ai/docs/leap-bundle/quick-start?utm_source=github&utm_medium=link&utm_campaign=LEAP&utm_content=general).
 
+## 📊 Expected Dataset Formats
+
+### SFT (Supervised Fine-Tuning)
+
+```json
+{
+  "messages": [
+    { "role": "user", "content": "What is the capital of France?" },
+    { "role": "assistant", "content": "The capital of France is Paris." }
+  ]
+}
+```
+
+### DPO (Direct Preference Optimization)
+
+```json
+{
+  "prompt": "What is the capital of France?",
+  "chosen": "The capital of France is Paris.",
+  "rejected": "The capital of France is London."
+}
+```
+
+### VLM SFT (Vision-Language Model)
+
+```json
+{
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "You are an image-based assistant. Answer questions based on the provided image."
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": [
+        { "type": "image", "image": "/path/to/image.jpg" },
+        { "type": "text", "text": "What do you see in this image?" }
+      ]
+    },
+    {
+      "role": "assistant",
+      "content": [{ "type": "text", "text": "I see a car in the image." }]
+    }
+  ]
+}
+```
+
+> **Note**: VLM datasets commonly have images in a separate row and are referenced in the messages column. If your image URLs or Paths are in a separate column from your messages, you'll need to merge the images into the 'messages' section like above.
+
 ## 🧪 Advanced Configuration
 
+### Default Configs Location and Adding New Configs
+
 The default configurations are located in:
 
 - **SFT Training**: [`src/leap_finetune/configs/sft_configs.py`](./src/leap_finetune/configs/sft_configs.py)

config.py

@@ -27,6 +27,10 @@
     "mlabonne/orpo-dpo-mix-40k", "dpo", limit=1000, test_size=0.2, subset="default"
 )
 
+example_vlm_sft_dataset = DatasetLoader(
+    "alay2shah/example-vlm-sft-dataset", "vlm_sft", limit=None, test_size=0.2
+)
+
 
 #################################
 #       Training Config         #
@@ -37,10 +41,12 @@
 
 Default SFT: TrainingConfig.DEFAULT_SFT
 Default DPO: TrainingConfig.DEFAULT_DPO
+Default VLM SFT: TrainingConfig.DEFAULT_VLM_SFT
 
-Default LORA: PeftConfig.DEFAULT_LORA
-High R LORA: PeftConfig.HIGH_R_LORA
 No LORA (full finetuning): PeftConfig.NO_LORA
+Default LORA: PeftConfig.DEFAULT_LORA
+High R LoRA: PeftConfig.HIGH_R_LORA
+VLM LoRA: PeftConfig.DEFAULT_VLM_LORA
 
 Args:
     output_dir: Output directory for training artifacts
@@ -76,7 +82,7 @@
 """
 
 JOB_CONFIG = JobConfig(
-    job_name="my_job",
+    job_name="my_job_name",
     model_name="LFM2-1.2B",
     training_type="sft",
     dataset=example_sft_dataset,

pyproject.toml

@@ -16,7 +16,9 @@ dependencies = [
     "transformers>=4.55.0",
     "trl>=0.18.2",
     "rich>=14.1.0",
-    "liger-kernel>=0.6.1",
+    "pillow>=11.3.0",
+    "mpi4py>=4.1.0",
+    "liger-kernel>=0.6.2",
 ]
 
 [project.scripts]

src/leap_finetune/__init__.py

@@ -1,6 +1,6 @@
 import sys
 
-from leap_finetune.utils.logging import setup_training_environment
+from leap_finetune.utils.logging_utils import setup_training_environment
 
 setup_training_environment()
 

src/leap_finetune/configs/__init__.py

@@ -5,15 +5,20 @@
 from .sft_configs import (
     DEFAULT_SFT_CONFIG,
 )
+from .vlm_sft_config import (
+    DEFAULT_VLM_SFT_CONFIG,
+)
 from .peft_configs import (
     LFM2_LORA_DEFAULT_CONFIG,
     LFM2_LORA_HIGH_R_CONFIG,
+    LFM2_VLM_LORA_CONFIG,
 )
 
 
 class TrainingConfig(Enum):
     DEFAULT_SFT = DEFAULT_SFT_CONFIG
     DEFAULT_DPO = DEFAULT_DPO_CONFIG
+    DEFAULT_VLM_SFT = DEFAULT_VLM_SFT_CONFIG
 
     def override(self, **overrides):
         """Create a custom TrainingConfig with overrides"""
@@ -29,6 +34,7 @@ class _CustomTrainingConfig(Enum):
 
 
 class PeftConfig(Enum):
+    NO_LORA = None
     DEFAULT_LORA = LFM2_LORA_DEFAULT_CONFIG
     HIGH_R_LORA = LFM2_LORA_HIGH_R_CONFIG
-    NO_LORA = None
+    DEFAULT_VLM_LORA = LFM2_VLM_LORA_CONFIG

src/leap_finetune/configs/job_config.py

@@ -1,16 +1,21 @@
 from dataclasses import dataclass
 from typing import Any, Literal
+from pathlib import Path
 
 from datasets import Dataset
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
 
 from leap_finetune.configs import PeftConfig, TrainingConfig
-from leap_finetune.utils.output_paths import (
-    is_job_name_unique,
-    resolve_model_output_path,
-)
+from leap_finetune.data_loaders.dataset_loader import DatasetLoader
+
+
+def is_job_name_unique(output_dir: str) -> bool:
+    """
+    Checks if a job with the given name exists for a given training type.
+    """
+    return not Path(output_dir).exists()
 
 
 @dataclass
@@ -19,13 +24,17 @@ class JobConfig:
 
     job_name: str
     model_name: str = "LFM2-1.2B"
-    training_type: Literal["sft", "dpo"] = "sft"
-    dataset: Dataset | None = None
+    training_type: Literal["sft", "dpo", "vlm_sft"] = "sft"
+    dataset: DatasetLoader | tuple[Dataset, Dataset] | None = None
     training_config: TrainingConfig = TrainingConfig.DEFAULT_SFT
     peft_config: PeftConfig | None = PeftConfig.DEFAULT_LORA
 
     def __post_init__(self):
-        self.dataset = self.dataset.load()  # Load dataset after init
+        if isinstance(self.dataset, DatasetLoader):
+            self.dataset = self.dataset.load()
+        self.training_config.value["output_dir"] = str(
+            Path(self.training_config.value.get("output_dir")) / self.job_name
+        )
         self._validate_job_name()
         self._validate_training_config()
 
@@ -38,10 +47,10 @@ def _validate_job_name(self):
 
         # Check if job dir already exists - warn but don't fail
         # (Ray workers might import config after directory is created)
-        if not is_job_name_unique(self.training_type, self.job_name):
+        if not is_job_name_unique(self.training_config.value.get("output_dir")):
             raise ValueError(
-                f"Job directory already exists for job '{self.job_name}' with training type '{self.training_type}'. "
-                f"This might be from a previous run or concurrent Ray worker initialization."
+                "Job output directory already exists\n"
+                "This might be from a previous run or concurrent Ray worker initialization."
             )
 
     def _validate_training_config(self):
@@ -75,7 +84,7 @@ def print_config_summary(self):
         console = Console()
 
         # Calculate output directory
-        output_dir = resolve_model_output_path(self.training_type, self.job_name)
+        output_dir = self.training_config.value.get("output_dir")
 
         # Create a table for the configuration
         table = Table(show_header=False, box=None, padding=(0, 2))

src/leap_finetune/configs/peft_configs.py

@@ -22,3 +22,18 @@
     lora_dropout=0.1,
     target_modules=GLU_MODULES + MHA_MODULES + CONV_MODULES,
 )
+
+
+LFM_MODULES = ["q_proj", "k_proj", "v_proj", "out_proj", "in_proj"]
+VISION_TOWER_MODULES = ["fc1", "fc2"]
+MULTI_MODAL_PROJECTOR_MODULES = ["linear_1", "linear_2"]
+
+LFM2_VLM_LORA_CONFIG = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    inference_mode=False,
+    r=8,
+    lora_alpha=16,
+    lora_dropout=0.1,
+    bias="none",
+    target_modules=LFM_MODULES + VISION_TOWER_MODULES + MULTI_MODAL_PROJECTOR_MODULES,
+)

src/leap_finetune/configs/vlm_sft_config.py

@@ -0,0 +1,59 @@
+from leap_finetune.utils.constants import SFT_OUTPUT_PATH
+
+########################
+#   DEEPSEED CONFIGS   #
+########################
+
+
+DEEPSEED_CONFIG = {
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": True,
+    },
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "gradient_clipping": "auto",
+    "gradient_accumulation_steps": "auto",
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",  # Uses learning_rate from training config
+            "betas": "auto",  # DEFAULT: (0.9, 0.999)
+            "eps": "auto",  # DEFAULT: 1e-8
+            "weight_decay": "auto",  # DEFAULT: 0.01
+        },
+    },
+    "bf16": {"enabled": "auto"},
+    "activation_checkpointing": {
+        "partition_activations": False,
+        "cpu_checkpointing": False,
+        "contiguous_memory_optimization": False,
+        "number_checkpoints": None,
+        "synchronize_checkpoint_boundary": False,
+        "profile": False,
+    },
+}
+
+
+########################
+#     SFT CONFIGS      #
+########################
+
+
+DEFAULT_VLM_SFT_CONFIG = {
+    "training_type": "vlm_sft",
+    "output_dir": SFT_OUTPUT_PATH,
+    "num_train_epochs": 3,  # 1 to 5 generally (post-training goes for 2-3)
+    "per_device_train_batch_size": 4,  # adjust based on context length (post-training goes for 1-2 at 32k context length)
+    "learning_rate": 5e-5,  # anything from 1e-5 to 5e-5 seems ok. "end_learning_rate" would be 1e-7, not easy to set up with out-of-the-box SFTConfig
+    "lr_scheduler_type": "linear",
+    "warmup_steps": 100,
+    "warmup_ratio": 0.2,
+    "logging_steps": 10,
+    "save_strategy": "epoch",
+    "eval_strategy": "epoch",
+    "load_best_model_at_end": True,
+    "gradient_checkpointing": True,
+    "dataset_kwargs": {"skip_prepare_dataset": True},
+    "deepspeed": DEEPSEED_CONFIG,
+}

src/leap_finetune/data_loaders/dataset_loader.py

@@ -12,7 +12,7 @@ class DatasetLoader:
     """Dataset loader for training and testing datasets"""
 
     dataset_path: str
-    dataset_type: Literal["sft", "dpo"]
+    dataset_type: Literal["sft", "dpo", "vlm_sft"]
     limit: Optional[int] = None  # Default: all samples
     split: str = "train"  # Default: "train"
     test_size: float = 0.2  # Default: 80/20 split

src/leap_finetune/data_loaders/image_loader.py

@@ -0,0 +1,15 @@
+from PIL import Image
+import io
+import requests
+
+
+def load_image(src):
+    """Load image from various sources and return PIL Image."""
+    # bytes -> PIL
+    if isinstance(src, (bytes, bytearray)):
+        return Image.open(io.BytesIO(src))
+    # URL -> PIL
+    if isinstance(src, str) and src.startswith(("http://", "https://")):
+        return Image.open(requests.get(src, stream=True).raw)
+    # file path -> PIL
+    return Image.open(src)