#!/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()