#!/usr/bin/env python3 """ Train LoRA adapter on Cyron summary dataset. Uses Hugging Face TRL for SFT training. """ import argparse import os import yaml from pathlib import Path import torch def train(config_path): """Train LoRA adapter using TRL.""" with open(config_path) as f: config = yaml.safe_load(f) from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, ) from peft import ( LoraConfig, get_peft_model, ) from trl import SFTTrainer print(f"Loading model: {config['base_model']}") # Create index file for sharded model import glob as glob_mod import json as json_mod safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors") shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name]) index_file = Path(config["base_model"]) / "model.safetensors.index.json" if not index_file.exists() and shards: print(f"\n[INFO] Creating index file for {len(shards)} shards...") weight_map = {} for shard_name in shards: shard_path = Path(config["base_model"]) / shard_name ckpt = torch.load(str(shard_path), map_location="cpu", weights_only=False) for key in ckpt.keys(): if isinstance(ckpt[key], torch.Tensor): weight_map[key] = shard_name index = { "metadata": {"total_size": sum((Path(config["base_model"]) / s).stat().st_size for s in shards)}, "weight_map": weight_map } with open(index_file, 'w') as f: json_mod.dump(index, f) print(f"✓ Created index ({len(weight_map)} weights)") # Remove quantization_config to prevent re-quantization config_json_path = Path(config["base_model"]) / "config.json" if config_json_path.exists(): with open(config_json_path, 'r') as f: config_data = json_mod.load(f) if 'quantization_config' in config_data: del config_data['quantization_config'] with open(config_json_path, 'w') as f: json_mod.dump(config_data, f) print(f"\n[INFO] Loading pre-quantized BnB 4-bit model...") model = AutoModelForCausalLM.from_pretrained( config["base_model"], device_map="cpu", torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) print("✓ Model loaded to CPU (BnB 4-bit)") # Move to GPU print(" Moving to GPU 0...") model = model.to("cuda:0") print("✓ Success: Model loaded to GPU 0 (4-bit)") print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB") print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB") # Add LoRA lora_config = LoraConfig( r=config["lora_r"], lora_alpha=config["lora_alpha"], lora_dropout=config["lora_dropout"], target_modules=config["target_modules"], task_type=config.get("lora_task_type", "CAUSAL_LM"), ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(config["base_model"]) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" # Load dataset from datasets import load_dataset import os # Get dataset path from config dataset_path = config["dataset"][0]["path"] print(f"Loading dataset from: {dataset_path}") dataset = load_dataset( "json", data_files={ "train": dataset_path, }, ) # Model is on single GPU print("✓ Model loaded to single GPU") print(f" GPU 0: {torch.cuda.memory_allocated(0) / 1e9:.2f} GB") print(f" Free VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9 - torch.cuda.memory_allocated(0) / 1e9:.2f} GB") # Training arguments training_args = TrainingArguments( output_dir=config["train_params"]["output_dir"], num_train_epochs=config["train_params"]["num_train_epochs"], per_device_train_batch_size=config["train_params"]["per_device_train_batch_size"], gradient_accumulation_steps=config["train_params"]["gradient_accumulation_steps"], learning_rate=float(config["train_params"]["learning_rate"]), lr_scheduler_type=config["train_params"]["lr_scheduler_type"], weight_decay=config["train_params"]["weight_decay"], warmup_steps=int(config["train_params"]["warmup_ratio"] * config["train_params"]["num_train_epochs"] * len(dataset["train"]) // config["train_params"]["per_device_train_batch_size"]), logging_steps=config["train_params"]["logging_steps"], save_steps=config["train_params"]["save_steps"], save_total_limit=config["train_params"]["save_total_limit"], eval_strategy=config.get("eval_strategy", "no"), bf16=True, gradient_checkpointing=config.get("gradient_checkpointing", False), optim=config["train_params"].get("optim", "adamw_torch"), optim_args=config["train_params"].get("optim_args"), ) print(f"Using optimizer: {training_args.optim}") # SFT Trainer from trl import SFTTrainer # Get text column from config (default to 'text' if not specified) text_column = config["dataset"][0].get("text_column", "text") print(f"Using text column: {text_column}") # Rename column to 'text' if needed (SFTTrainer expects 'text') if text_column != "text": print(f" Renaming '{text_column}' column to 'text'...") dataset["train"] = dataset["train"].rename_column(text_column, "text") if "test" in dataset and dataset["test"] is not None: dataset["test"] = dataset["test"].rename_column(text_column, "text") trainer = SFTTrainer( model=model, processing_class=tokenizer, train_dataset=dataset["train"], eval_dataset=dataset.get("test"), args=training_args, ) # Train print("Starting training...") trainer.train() # Save trainer.save_model(config["train_params"]["output_dir"]) tokenizer.save_pretrained(config["train_params"]["output_dir"]) print(f"Training complete! Model saved to {config['train_params']['output_dir']}") def main(): parser = argparse.ArgumentParser(description="Train LoRA adapter") parser.add_argument("--config", type=str, default="training/configs/ornith-35b-lora.yaml", help="Training configuration file") parser.add_argument("--check-only", action="store_true", help="Validate config and dependencies without training") args = parser.parse_args() if args.check_only: check_setup(args.config) else: train(args.config) def check_setup(config_path): """Validate config and dependencies without loading model.""" print("=== Checking Setup ===") # Check config file print(f"\n1. Config file: {config_path}") if not Path(config_path).exists(): print(f" ERROR: Config file not found: {config_path}") return False print(" ✓ Config file exists") # Load and validate config with open(config_path) as f: config = yaml.safe_load(f) print(" ✓ Config file is valid YAML") # Check model path print(f"\n2. Model path: {config['base_model']}") if Path(config['base_model']).exists(): print(f" ✓ Model path exists: {config['base_model']}") else: print(f" ⚠ Model path not found: {config['base_model']}") print(" (Will download from HuggingFace during training)") # Check dataset files print("\n3. Dataset files:") repo_root = Path(__file__).parent train_path = repo_root / "training" / "data" / "train.jsonl" test_path = repo_root / "training" / "data" / "test.jsonl" if train_path.exists(): print(f" ✓ Train data: {train_path}") else: print(f" ✗ Train data missing: {train_path}") if test_path.exists(): print(f" ✓ Test data: {test_path}") else: print(f" ✗ Test data missing: {test_path}") # Check GPU print("\n4. GPU:") import torch if torch.cuda.is_available(): print(f" ✓ GPU available: {torch.cuda.get_device_name(0)}") print(f" ✓ VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") print(f" ✓ GPU count: {torch.cuda.device_count()}") else: print(" ✗ No GPU detected!") return False # Check required packages print("\n5. Required packages:") packages = ['transformers', 'datasets', 'trl', 'peft', 'accelerate', 'deepspeed'] for pkg in packages: try: __import__(pkg) print(f" ✓ {pkg}") except ImportError: print(f" ✗ {pkg} not installed") # Check DeepSpeed config print("\n6. DeepSpeed config:") if 'deepspeed_config' in config: print(" ✓ DeepSpeed config present") ds_config = config['deepspeed_config'] if 'zero_optimization' in ds_config: stage = ds_config['zero_optimization'].get('stage', 'N/A') print(f" ✓ ZeRO stage: {stage}") else: print(" ✗ No DeepSpeed config found") print("\n=== Check Complete ===") print("If all checks pass, you can run training with:") print(f" bash train-on-this-server.sh") if __name__ == "__main__": main()