freight/scripts/train_qwen_lora.py

#!/usr/bin/env python3
"""
scripts/train_qwen_lora.py — QLoRA / SFT Training Recipe for Qwen 3.8 9B on Freight Negotiation.

Trains Qwen 3.8 9B (or Qwen 2.5 7B/14B) on tool-calling freight dialogues using 4-bit QLoRA.
Supports execution on:
- Local GPU / Homelab (8GB VRAM with paged_adamw_8bit)
- Ephemeral HF Space (A10G ~$1.05/hr)
- Google Colab / Lambda Labs (A100/T4)

Usage:
    python3 scripts/train_qwen_lora.py --dataset_path data/freight_negotiation_sample.jsonl --epochs 3
"""

import os
import sys
import json
import argparse
from typing import Dict, Any

def main():
    parser = argparse.ArgumentParser(description="QLoRA Fine-Tuning for Freight LLM")
    parser.add_argument("--model_id", default="Qwen/Qwen2.5-7B-Instruct", help="Base model ID on Hugging Face")
    parser.add_argument("--dataset_path", default="freight/data/freight_negotiation_sample.jsonl", help="JSONL dataset path")
    parser.add_argument("--output_dir", default="models/loadeta-qwen3.8-9b-freight", help="Output directory for adapters")
    parser.add_argument("--lora_r", type=int, default=16, help="LoRA rank")
    parser.add_argument("--lora_alpha", type=int, default=32, help="LoRA alpha")
    parser.add_argument("--batch_size", type=int, default=1, help="Per device train batch size")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=8, help="Gradient accumulation steps")
    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")
    parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs")
    parser.add_argument("--max_seq_length", type=int, default=2048, help="Max sequence length")
    parser.add_argument("--push_to_hub", action="store_true", help="Push trained adapter to Hugging Face Hub")
    parser.add_argument("--hub_model_id", default="abalanescu/loadeta-qwen3.8-9b-freight", help="HF Hub repo ID")
    parser.add_argument("--dry_run", action="store_true", help="Print config and validate dependencies without training")
    args = parser.parse_args()

    print("=== LoadETA Qwen 3.8 9B Freight QLoRA Trainer ===")
    print(f"Base Model: {args.model_id}")
    print(f"Dataset:    {args.dataset_path}")
    print(f"Output:     {args.output_dir}")
    print(f"LoRA Config: r={args.lora_r}, alpha={args.lora_alpha}, target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj']")
    print(f"Hyperparams: lr={args.learning_rate}, batch_size={args.batch_size}x{args.gradient_accumulation_steps} (effective {args.batch_size*args.gradient_accumulation_steps}), epochs={args.epochs}")

    if not os.path.exists(args.dataset_path):
        print(f"Error: Dataset not found at {args.dataset_path}")
        sys.exit(1)

    # Count dataset samples
    with open(args.dataset_path, "r", encoding="utf-8") as f:
        count = sum(1 for line in f if line.strip())
    print(f"Found {count} conversations in dataset.")

    if args.dry_run:
        print("Dry run complete. Ready for GPU training execution.")
        return

    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
        from trl import SFTTrainer, SFTConfig
        from datasets import load_dataset
    except ImportError as e:
        print(f"\n[Notice] Missing ML training dependencies: {e}")
        print("To run actual GPU training, install requirements:")
        print("pip install torch transformers peft bitsandbytes trl datasets accelerate")
        print("\nOr fine-tune locally on Apple Silicon using MLX:")
        print(f"mlx_lm.lora --model {args.model_id} --train --data {args.dataset_path} --batch-size 2 --iters 600")
        return

    # 4-bit Quantization Config (QLoRA)
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
        bnb_4bit_use_double_quant=True,
    )

    print("\nLoading tokenizer and quantized base model...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )
    model = prepare_model_for_kbit_training(model)

    lora_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    dataset = load_dataset("json", data_files=args.dataset_path, split="train")

    training_args = SFTConfig(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        num_train_epochs=args.epochs,
        logging_steps=5,
        save_strategy="epoch",
        optim="paged_adamw_8bit",
        fp16=not torch.cuda.is_bf16_supported(),
        bf16=torch.cuda.is_bf16_supported(),
        max_grad_norm=0.3,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        report_to="none",
        max_seq_length=args.max_seq_length,
    )

    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset,
        peft_config=lora_config,
        args=training_args,
    )

    print("\nStarting training loop...")
    trainer.train()

    print(f"\nSaving fine-tuned LoRA adapters to {args.output_dir}...")
    trainer.model.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    if args.push_to_hub:
        print(f"Pushing to Hugging Face Hub: {args.hub_model_id}...")
        trainer.model.push_to_hub(args.hub_model_id)
        tokenizer.push_to_hub(args.hub_model_id)

    print("\nTraining completed successfully!")
    print("Next step: Merge adapters and convert to GGUF using llama.cpp:")
    print(f"python3 llama.cpp/convert_hf_to_gguf.py {args.output_dir} --outtype q8_0")


if __name__ == "__main__":
    main()