#!/usr/bin/env python3 """ JEDI_LFM2.5 LoRA Fine-Tuning Script Trains the model on our cross-domain + Veritas dataset using QLoRA. Targets: ~2.67M tokens of connective/Machiavelli/Veritas training data. Usage: python3 tune_jedi.py # Full training run python3 tune_jedi.py --quick # Quick test run (100 examples) python3 tune_jedi.py --resume # Resume from checkpoint """ import json, os, sys, gc, math, random from pathlib import Path import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from transformers import ( AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling, BitsAndBytesConfig, set_seed, ) from peft import ( LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel, ) import bitsandbytes as bnb # ─── CONFIG ───────────────────────────────────────────────────── MODEL_ID = "LiquidAI/LFM2.5-1.2B-Instruct" MASTER_DATA = "/root/JEDI/training_data_master.jsonl" OUTPUT_DIR = "/root/JEDI/lora_checkpoints" FINAL_ADAPTER = "/root/JEDI/jedi_lora_adapter" SEED = 42 set_seed(SEED) # Training params (CPU-friendly) LORA_R = 16 LORA_ALPHA = 32 LORA_DROPOUT = 0.05 TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] MAX_LENGTH = 1024 BATCH_SIZE = 2 # tiny for CPU GRAD_ACCUM = 4 # effective batch = 8 LEARNING_RATE = 3e-4 NUM_EPOCHS = 1 WARMUP_STEPS = 50 LOGGING_STEPS = 10 SAVE_STEPS = 200 MAX_STEPS = 1000 # ~2.67M tokens / (1024*8) ≈ 325 steps per epoch, over-estimate QUICK_MODE = "--quick" in sys.argv RESUME = "--resume" in sys.argv if QUICK_MODE: MAX_STEPS = 20 print("[QUICK MODE] Training on 100 examples, 20 steps") # ─── DATASET ──────────────────────────────────────────────────── class ShareGPTDataset(Dataset): """Load ShareGPT-format JSONL and format for causal LM training.""" def __init__(self, path, tokenizer, max_length=1024, max_examples=None): self.tokenizer = tokenizer self.max_length = max_length self.examples = [] with open(path) as f: for line in f: line = line.strip() if line: try: item = json.loads(line) self.examples.append(item) except json.JSONDecodeError: continue if max_examples and len(self.examples) > max_examples: random.shuffle(self.examples) self.examples = self.examples[:max_examples] print(f"Loaded {len(self.examples)} training examples") def __len__(self): return len(self.examples) def __getitem__(self, idx): item = self.examples[idx] messages = item.get("messages", []) # Format as ChatML (LFM2.5 template) formatted = "" for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") if role == "system": formatted += f"<|im_start|>system\n{content}<|im_end|>\n" elif role == "user": formatted += f"<|im_start|>user\n{content}<|im_end|>\n" elif role == "assistant": formatted += f"<|im_start|>assistant\n{content}<|im_end|>\n" formatted += "<|im_start|>assistant\n" # prompt for generation encoded = self.tokenizer( formatted, truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt", ) labels = encoded["input_ids"].clone() # Mask user input (don't compute loss on it) user_tokens = self.tokenizer( "<|im_start|>user\n", add_special_tokens=False )["input_ids"] assistant_tokens = self.tokenizer( "<|im_start|>assistant\n", add_special_tokens=False )["input_ids"] # Find all user sections and mask their labels input_ids = encoded["input_ids"][0] labels_seq = labels[0].clone() # Simple approach: mask everything before the last assistant token # Find positions of assistant tokens assistant_len = len(assistant_tokens) input_len = len(input_ids) # Find last occurrence of assistant header last_asst_pos = -1 for i in range(input_len - assistant_len): if (input_ids[i:i+assistant_len] == torch.tensor(assistant_tokens)).all(): last_asst_pos = i if last_asst_pos > 0: # Mask everything before the last assistant turn labels_seq[:last_asst_pos] = -100 else: # If no assistant found, mask everything (safety) labels_seq = torch.full_like(labels_seq, -100) return { "input_ids": encoded["input_ids"][0], "attention_mask": encoded["attention_mask"][0], "labels": labels_seq, } # ─── MODEL SETUP ──────────────────────────────────────────────── def setup_model(): """Load model with QLoRA 4-bit quantization.""" print(f"Loading model: {MODEL_ID}") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float32, # CPU doesn't support bfloat16 ) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, quantization_config=bnb_config, device_map="auto", trust_remote_code=True, torch_dtype=torch.float32, ) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" # Prepare for k-bit training model = prepare_model_for_kbit_training(model) # Configure LoRA lora_config = LoraConfig( r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=TARGET_MODULES, lora_dropout=LORA_DROPOUT, bias="none", task_type="CAUSAL_LM", ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() return model, tokenizer # ─── MAIN ─────────────────────────────────────────────────────── def train(): model, tokenizer = setup_model() max_examples = 100 if QUICK_MODE else None dataset = ShareGPTDataset(MASTER_DATA, tokenizer, MAX_LENGTH, max_examples) # Data collator collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=False, ) # Training args args = TrainingArguments( output_dir=OUTPUT_DIR, overwrite_output_dir=True, num_train_epochs=NUM_EPOCHS, max_steps=MAX_STEPS, per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM, warmup_steps=WARMUP_STEPS, logging_steps=LOGGING_STEPS, save_steps=SAVE_STEPS, learning_rate=LEARNING_RATE, lr_scheduler_type="cosine", optim="adamw_8bit" if torch.cuda.is_available() else "adamw_torch", report_to="none", ddp_find_unused_parameters=False, gradient_checkpointing=True, fp16=False, bf16=False, dataloader_pin_memory=False, max_grad_norm=0.3, remove_unused_columns=False, ) trainer = Trainer( model=model, args=args, train_dataset=dataset, data_collator=collator, tokenizer=tokenizer, ) # Disable caching model.config.use_cache = False print(f"\n{'='*50}") print(f"Starting training:") print(f" Model: {MODEL_ID}") print(f" Data: {len(dataset)} examples") print(f" Steps: {MAX_STEPS}") print(f" Batch: {BATCH_SIZE} (eff: {BATCH_SIZE * GRAD_ACCUM})") print(f" LoRA r={LORA_R}, α={LORA_ALPHA}") print(f" Device: {'CPU' if not torch.cuda.is_available() else 'GPU'}") print(f"{'='*50}\n") # Count parameters trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) print(f"Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)") trainer.train(resume_from_checkpoint=RESUME) # Save adapter model.save_pretrained(FINAL_ADAPTER) tokenizer.save_pretrained(FINAL_ADAPTER) print(f"\nLoRA adapter saved to: {FINAL_ADAPTER}") # Save merged model path for reference with open(os.path.join(FINAL_ADAPTER, "base_model.txt"), "w") as f: f.write(MODEL_ID) print("Training complete!") def apply_adapter(): """Apply the trained LoRA adapter back to the base model for testing.""" print(f"Loading base model + LoRA adapter from {FINAL_ADAPTER}...") model = PeftModel.from_pretrained( AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True, device_map="auto"), FINAL_ADAPTER, ) tokenizer = AutoTokenizer.from_pretrained(FINAL_ADAPTER) # Test prompt = "<|im_start|>system\nYou are JEDI — forensic analytical engine. Connect everything to psychology and Machiavelli.<|im_end|>\n<|im_start|>user\nWhat is the connection between Machiavelli's 'trust is safer to fear than love' and zero-trust architecture?<|im_end|>\n<|im_start|>assistant\n" inputs = tokenizer(prompt, return_tensors="pt") outputs = model.generate( **inputs, max_new_tokens=200, temperature=0.3, do_sample=True, ) response = tokenizer.decode(outputs[0], skip_special_tokens=False) print("\n" + "="*50) print("TEST GENERATION:") print(response[len(prompt):]) print("="*50) if __name__ == "__main__": if "--apply" in sys.argv: apply_adapter() elif "--help" in sys.argv or "-h" in sys.argv: print("Usage: python3 tune_jedi.py [--quick|--resume|--apply]") print(" (no flag) Full training run") print(" --quick Quick test (100 examples, 20 steps)") print(" --resume Resume from checkpoint") print(" --apply Test the trained adapter") else: train()