openadapt-ml 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/training/trl_trainer.py

@@ -0,0 +1,403 @@
+"""Simplified training using TRL SFTTrainer + Unsloth.
+
+This module provides a minimal, efficient training path for VLMs:
+- Unsloth for 2x speed, 50% less VRAM
+- TRL SFTTrainer for production-grade training
+- Direct integration with openadapt-ml data format
+
+Usage:
+    from openadapt_ml.training.trl_trainer import train_with_trl
+
+    # Train on episodes
+    train_with_trl(
+        episodes=episodes,
+        model_name="unsloth/Qwen2.5-VL-7B-Instruct",
+        output_dir="checkpoints/my_model",
+    )
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from PIL import Image
+
+
+@dataclass
+class TRLTrainingConfig:
+    """Configuration for TRL-based training."""
+
+    # Model
+    model_name: str = "unsloth/Qwen2.5-VL-7B-Instruct"
+    load_in_4bit: bool = True
+    max_seq_length: int = 4096
+
+    # LoRA
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.0
+    finetune_vision_layers: bool = False  # Set True if grounding needs improvement
+
+    # Training
+    num_epochs: int = 3
+    batch_size: int = 1
+    gradient_accumulation_steps: int = 4
+    learning_rate: float = 2e-4
+    warmup_ratio: float = 0.03
+
+    # Output
+    output_dir: str = "checkpoints"
+    logging_steps: int = 10
+    save_strategy: str = "epoch"
+
+
+def _load_unsloth_model(config: TRLTrainingConfig):
+    """Load model with Unsloth optimizations.
+
+    Returns:
+        tuple: (model, tokenizer, is_unsloth) - is_unsloth indicates if Unsloth was used
+    """
+    # Check if Unsloth is explicitly disabled via environment variable
+    if os.environ.get("OPENADAPT_DISABLE_UNSLOTH", "").lower() in ("1", "true", "yes"):
+        print("Unsloth disabled via OPENADAPT_DISABLE_UNSLOTH environment variable")
+        return _load_standard_model(config)
+
+    try:
+        from unsloth import FastVisionModel
+
+        model, tokenizer = FastVisionModel.from_pretrained(
+            config.model_name,
+            load_in_4bit=config.load_in_4bit,
+            use_gradient_checkpointing="unsloth",
+            max_seq_length=config.max_seq_length,
+        )
+
+        # Apply LoRA
+        model = FastVisionModel.get_peft_model(
+            model,
+            finetune_vision_layers=config.finetune_vision_layers,
+            finetune_language_layers=True,
+            finetune_attention_modules=True,
+            finetune_mlp_modules=True,
+            r=config.lora_r,
+            lora_alpha=config.lora_alpha,
+            lora_dropout=config.lora_dropout,
+            random_state=42,
+        )
+
+        # Enable training mode
+        FastVisionModel.for_training(model)
+
+        print(
+            f"✓ Loaded {config.model_name} with Unsloth (4-bit: {config.load_in_4bit})"
+        )
+        return model, tokenizer, True
+
+    except ImportError:
+        print("⚠ Unsloth not installed, falling back to standard transformers")
+        return _load_standard_model(config)
+
+
+def _load_standard_model(config: TRLTrainingConfig):
+    """Fallback: Load model with standard transformers + peft.
+
+    Automatically detects vision-language models and uses the appropriate
+    model class (Qwen2VLForConditionalGeneration for VL models,
+    AutoModelForCausalLM for text-only models).
+    """
+    from transformers import AutoConfig, AutoProcessor
+    from peft import LoraConfig, get_peft_model
+    import torch
+
+    # Check if this is a vision-language model
+    model_config = AutoConfig.from_pretrained(
+        config.model_name, trust_remote_code=True
+    )
+    is_vl_model = (
+        "VL" in config.model_name.upper()
+        or "vision" in config.model_name.lower()
+        or hasattr(model_config, "vision_config")
+    )
+
+    if is_vl_model:
+        # Vision-language model - use Qwen2VLForConditionalGeneration or AutoModelForVision2Seq
+        try:
+            from transformers import Qwen2VLForConditionalGeneration
+
+            model = Qwen2VLForConditionalGeneration.from_pretrained(
+                config.model_name,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                trust_remote_code=True,
+            )
+            print("  Using Qwen2VLForConditionalGeneration for VL model")
+        except (ImportError, ValueError, RuntimeError, TypeError):
+            # Fallback to AutoModelForVision2Seq for other VL models
+            from transformers import AutoModelForVision2Seq
+
+            model = AutoModelForVision2Seq.from_pretrained(
+                config.model_name,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+                trust_remote_code=True,
+            )
+            print("  Using AutoModelForVision2Seq for VL model")
+    else:
+        # Text-only model
+        from transformers import AutoModelForCausalLM
+
+        model = AutoModelForCausalLM.from_pretrained(
+            config.model_name,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        print("  Using AutoModelForCausalLM for text-only model")
+
+    processor = AutoProcessor.from_pretrained(config.model_name, trust_remote_code=True)
+
+    # Apply LoRA - use SEQ_2_SEQ_LM for VL models, CAUSAL_LM for text-only
+    peft_config = LoraConfig(
+        r=config.lora_r,
+        lora_alpha=config.lora_alpha,
+        lora_dropout=config.lora_dropout,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        task_type="SEQ_2_SEQ_LM" if is_vl_model else "CAUSAL_LM",
+    )
+    model = get_peft_model(model, peft_config)
+
+    print(f"✓ Loaded {config.model_name} with standard transformers")
+    return model, processor, False
+
+
+def _convert_samples_to_trl_format(
+    samples: List[Dict[str, Any]],
+    base_path: Optional[Path] = None,
+) -> List[Dict[str, Any]]:
+    """Convert openadapt-ml samples to TRL format.
+
+    The only change is loading image paths as PIL Images.
+
+    Args:
+        samples: List of samples from build_next_action_sft_samples()
+        base_path: Optional base path to resolve relative image paths
+
+    Returns:
+        List of samples with PIL Images instead of paths
+    """
+    trl_samples = []
+
+    for sample in samples:
+        # Load images as PIL
+        pil_images = []
+        for img_path in sample["images"]:
+            path = Path(img_path)
+            if base_path and not path.is_absolute():
+                path = base_path / path
+
+            if path.exists():
+                pil_images.append(Image.open(path).convert("RGB"))
+            else:
+                print(f"⚠ Image not found: {path}")
+                continue
+
+        if not pil_images:
+            continue  # Skip samples with missing images
+
+        trl_samples.append(
+            {
+                "images": pil_images,
+                "messages": sample["messages"],
+            }
+        )
+
+    return trl_samples
+
+
+def train_with_trl(
+    episodes: List,
+    config: Optional[TRLTrainingConfig] = None,
+    use_som: bool = False,
+    base_path: Optional[Path] = None,
+) -> str:
+    """Train a VLM using TRL SFTTrainer + Unsloth.
+
+    This is the simplified training entry point that replaces the legacy
+    custom training loop. It:
+    1. Converts episodes to TRL format
+    2. Loads model with Unsloth (or fallback)
+    3. Trains with TRL's SFTTrainer
+    4. Saves LoRA adapter
+
+    Args:
+        episodes: List of Episode objects from openadapt-ml schema
+        config: Training configuration (uses defaults if None)
+        use_som: If True, use Set-of-Marks DSL instead of coordinates
+        base_path: Base path for resolving relative image paths
+
+    Returns:
+        Path to saved checkpoint
+    """
+    from datasets import Dataset
+    from openadapt_ml.datasets.next_action import build_next_action_sft_samples
+
+    config = config or TRLTrainingConfig()
+
+    # Step 1: Convert episodes to SFT samples
+    print(f"Converting {len(episodes)} episodes to training samples...")
+    raw_samples = build_next_action_sft_samples(episodes, use_som=use_som)
+    print(f"  Generated {len(raw_samples)} training samples")
+
+    # Step 2: Convert to TRL format (load images as PIL)
+    print("Loading images...")
+    trl_samples = _convert_samples_to_trl_format(raw_samples, base_path)
+    print(f"  Loaded {len(trl_samples)} samples with images")
+
+    if not trl_samples:
+        raise ValueError("No valid training samples after loading images")
+
+    # Step 3: Create HuggingFace Dataset
+    dataset = Dataset.from_list(trl_samples)
+
+    # Step 4: Load model with Unsloth (or fallback)
+    model, tokenizer, is_unsloth = _load_unsloth_model(config)
+
+    # Step 5: Configure and run training
+    try:
+        from trl import SFTTrainer, SFTConfig
+
+        if is_unsloth:
+            # Unsloth-specific configuration
+            from unsloth.trainer import UnslothVisionDataCollator
+
+            training_args = SFTConfig(
+                output_dir=config.output_dir,
+                per_device_train_batch_size=config.batch_size,
+                gradient_accumulation_steps=config.gradient_accumulation_steps,
+                learning_rate=config.learning_rate,
+                num_train_epochs=config.num_epochs,
+                warmup_ratio=config.warmup_ratio,
+                lr_scheduler_type="cosine",
+                logging_steps=config.logging_steps,
+                save_strategy=config.save_strategy,
+                # Unsloth-specific settings
+                remove_unused_columns=False,
+                dataset_text_field="",
+                dataset_kwargs={"skip_prepare_dataset": True},
+            )
+
+            trainer = SFTTrainer(
+                model=model,
+                tokenizer=tokenizer,
+                data_collator=UnslothVisionDataCollator(model, tokenizer),
+                train_dataset=dataset,
+                args=training_args,
+            )
+        else:
+            # Standard TRL configuration
+            training_args = SFTConfig(
+                output_dir=config.output_dir,
+                per_device_train_batch_size=config.batch_size,
+                gradient_accumulation_steps=config.gradient_accumulation_steps,
+                learning_rate=config.learning_rate,
+                num_train_epochs=config.num_epochs,
+                warmup_ratio=config.warmup_ratio,
+                lr_scheduler_type="cosine",
+                logging_steps=config.logging_steps,
+                save_strategy=config.save_strategy,
+                max_length=None,  # Critical for VLMs
+                assistant_only_loss=False,  # Not supported for VL models yet
+            )
+
+            trainer = SFTTrainer(
+                model=model,
+                train_dataset=dataset,
+                args=training_args,
+            )
+
+        print(f"\n{'=' * 50}")
+        print("Starting training:")
+        print(f"  Model: {config.model_name}")
+        print(f"  Samples: {len(trl_samples)}")
+        print(f"  Epochs: {config.num_epochs}")
+        print(f"  Batch size: {config.batch_size}")
+        print(f"  Unsloth: {is_unsloth}")
+        print(f"  Output: {config.output_dir}")
+        print(f"{'=' * 50}\n")
+
+        trainer.train()
+
+        # Save the LoRA adapter
+        checkpoint_path = Path(config.output_dir) / "final"
+        trainer.save_model(str(checkpoint_path))
+        print(f"\n✓ Saved checkpoint to {checkpoint_path}")
+
+        return str(checkpoint_path)
+
+    except ImportError as e:
+        raise ImportError(
+            f"TRL not installed. Install with: pip install trl\nOriginal error: {e}"
+        )
+
+
+def train_from_parquet(
+    parquet_path: str,
+    config: Optional[TRLTrainingConfig] = None,
+    use_som: bool = False,
+) -> str:
+    """Train from a parquet file exported by openadapt-ml.
+
+    Args:
+        parquet_path: Path to parquet file with episode data
+        config: Training configuration
+        use_som: Use Set-of-Marks DSL
+
+    Returns:
+        Path to saved checkpoint
+    """
+    from openadapt_ml.export import from_parquet
+
+    print(f"Loading episodes from {parquet_path}...")
+    episodes = from_parquet(parquet_path)
+
+    base_path = Path(parquet_path).parent
+
+    return train_with_trl(
+        episodes=episodes,
+        config=config,
+        use_som=use_som,
+        base_path=base_path,
+    )
+
+
+if __name__ == "__main__":
+    # Simple CLI for testing
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Train VLM with TRL + Unsloth")
+    parser.add_argument("--parquet", required=True, help="Path to parquet file")
+    parser.add_argument("--output", default="checkpoints", help="Output directory")
+    parser.add_argument(
+        "--model", default="unsloth/Qwen2.5-VL-7B-Instruct", help="Model name"
+    )
+    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs")
+    parser.add_argument("--use-som", action="store_true", help="Use Set-of-Marks DSL")
+
+    args = parser.parse_args()
+
+    config = TRLTrainingConfig(
+        model_name=args.model,
+        output_dir=args.output,
+        num_epochs=args.epochs,
+    )
+
+    checkpoint = train_from_parquet(
+        parquet_path=args.parquet,
+        config=config,
+        use_som=args.use_som,
+    )
+
+    print(f"\nTraining complete! Checkpoint: {checkpoint}")