#!/usr/bin/env python3 """ Train Qwen2.5-27B on software development data via QLoRA (CPU-only). Features: - Checkpoint resumption (survives crashes/restarts) - Progress logging to file - CPU-only (no GPU needed) - Optimized for Ryzen 9 9900X + 60GB RAM Usage: python3 train_dev_model.py python3 train_dev_model.py --resume # Resume from last checkpoint """ import os import sys import json import argparse import time import logging from pathlib import Path from datetime import datetime # Force CPU-only PyTorch os.environ["CUDA_VISIBLE_DEVICES"] = "" os.environ["CUDA_DEVICE_ORDER"] = "" import torch from torch.utils.data import Dataset # Load model from local path MODEL_NAME = "/home/vincent/projects/bitcoin-ai/model" DATA_FILE = Path(__file__).parent.parent / "data" / "processed" / "dev_finetuning_massive.json" OUTPUT_DIR = Path(__file__).parent.parent / "output" / "dev_ai_model" LOG_FILE = Path(__file__).parent.parent / "output" / "training.log" # Training config MAX_STEPS = 3000 # More steps for thorough training BATCH_SIZE = 1 GRADIENT_ACCUM = 4 LEARNING_RATE = 2e-4 WARMUP_STEPS = 300 SAVE_INTERVAL = 200 # Save checkpoint every N steps # Setup logging logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.FileHandler(LOG_FILE), logging.StreamHandler(sys.stdout) ] ) logger = logging.getLogger(__name__) class DevDataset(Dataset): """Load instruction-tuning pairs for code training.""" def __init__(self, data_path, tokenizer, max_length=2048): with open(data_path, "r") as f: self.data = json.load(f) self.tokenizer = tokenizer self.max_length = max_length logger.info(f"Loaded {len(self.data)} training pairs") def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] messages = item["messages"] # Build chat format prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=False ) tokenized = self.tokenizer( prompt, truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt" ) input_ids = tokenized["input_ids"].squeeze() attention_mask = tokenized["attention_mask"].squeeze() # Labels = input_ids (for causal LM, mask padding with -100) labels = input_ids.clone() labels[attention_mask == 0] = -100 return { "input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, } def find_latest_checkpoint(output_dir): """Find the latest checkpoint directory for resumption.""" output_dir = Path(output_dir) checkpoints = [d for d in output_dir.iterdir() if d.is_dir() and d.name.startswith("checkpoint-")] if not checkpoints: return None # Extract step number and sort def step_num(d): return int(d.name.split("-")[1]) latest = sorted(checkpoints, key=step_num)[-1] logger.info(f"Found checkpoint: {latest.name}") return latest def main(resume=False): logger.info("=" * 60) logger.info("Dev-AI Training - Qwen2.5-27B QLoRA (CPU)") logger.info("=" * 60) # Check data exists if not DATA_FILE.exists(): logger.error(f"Dataset not found: {DATA_FILE}") logger.error("Run scripts/collect_code_dataset.py first") sys.exit(1) # Load model logger.info(f"Loading model from {MODEL_NAME}...") from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float32, bnb_4bit_use_double_quant=False, ) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, quantization_config=bnb_config, device_map="cpu", torch_dtype=torch.float32, low_cpu_mem_usage=True, ) logger.info(f"Model loaded: {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B parameters") # Setup LoRA from peft import LoraConfig, get_peft_model, TaskType lora_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_dropout=0.05, bias="none", task_type=TaskType.CAUSAL_LM, ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() # Load dataset dataset = DevDataset(DATA_FILE, tokenizer) # Training setup from transformers import ( TrainingArguments, Trainer, DataCollatorForSeq2Seq, ) # Resume checkpoint resume_from = None if resume: resume_from = find_latest_checkpoint(OUTPUT_DIR) if resume_from: logger.info(f"Resuming from {resume_from}") training_args = TrainingArguments( output_dir=str(OUTPUT_DIR), num_train_epochs=1, max_steps=MAX_STEPS, per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRADIENT_ACCUM, learning_rate=LEARNING_RATE, fp16=False, bf16=False, logging_steps=10, save_steps=SAVE_INTERVAL, save_total_limit=10, # Keep last 10 checkpoints optim="adamw_torch", warmup_steps=WARMUP_STEPS, report_to="none", dataloader_num_workers=0, ddp_find_unused_parameters=False, resume_from_checkpoint=resume_from, ) data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model) # Custom callback to log progress class ProgressCallback: def __init__(self): self.start_time = time.time() self.last_step_time = None self.last_step_num = None def on_step_end(self, args, state, control, **kwargs): elapsed = time.time() - self.start_time steps_done = state.global_step if self.last_step_time: step_time = time.time() - self.last_step_time else: step_time = 0 self.last_step_time = time.time() if steps_done % 50 == 0: eta_seconds = (MAX_STEPS - steps_done) * (elapsed / max(steps_done, 1)) eta_hours = eta_seconds / 3600 logger.info( f"Step {steps_done}/{MAX_STEPS} | " f"Loss: {state.log_history[-1].get('loss', 0):.4f} | " f"Step time: {step_time:.1f}s | " f"ETA: {eta_hours:.1f}h" ) trainer = Trainer( model=model, args=training_args, train_dataset=dataset, data_collator=data_collator, ) trainer.add_callback(ProgressCallback()) logger.info(f"Starting training: {MAX_STEPS} steps, batch={BATCH_SIZE}, accum={GRADIENT_ACCUM}") logger.info(f"Checkpoint interval: every {SAVE_INTERVAL} steps") logger.info(f"Log file: {LOG_FILE}") trainer.train(resume_from_checkpoint=resume_from) # Save final adapter final_path = OUTPUT_DIR / "final_adapter" model.save_pretrained(str(final_path)) tokenizer.save_pretrained(str(final_path)) logger.info("=" * 60) logger.info(f"TRAINING COMPLETE") logger.info(f"Final adapter saved to: {final_path}") logger.info(f"Total time: {(time.time() - ProgressCallback().start_time) / 3600:.1f} hours") logger.info("=" * 60) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--resume", action="store_true", help="Resume from latest checkpoint") args = parser.parse_args() main(resume=args.resume)