#!/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']}") # Load model with multiple strategies print(f"\n[INFO] Loading {config['base_model']}...") errors = [] # ------------------------------------------------------------------ # Strategy 1: Load already-quantized 4-bit model (distributed across GPUs) # ------------------------------------------------------------------ print("\n[1/4] Trying: 4-bit model AS-IS (distributed across GPUs)...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], device_map="auto", # Distribute layers across GPUs automatically torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) print("✓ Success: 4-bit model distributed across GPUs (no FSDP)") except Exception as e: errors.append(("QLoRA 4-bit", e)) print(f"✗ Failed: {e}") # -------------------------------------------------------------- # Strategy 2: 4-bit model with FSDP (load to GPU, FSDP shards) # -------------------------------------------------------------- print("\n[2/4] Trying: 4-bit model with FSDP (load to GPU)...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], device_map="auto", # Load to GPU first torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True, ) print("✓ Success: 4-bit model loaded to GPU (FSDP will shard)") except Exception as e: errors.append(("QLoRA 4-bit FSDP", e)) print(f"✗ Failed: {e}") # -------------------------------------------------------------- # Strategy 3: BF16 CPU (model too large for single GPU) # -------------------------------------------------------------- print("\n[3/4] Trying: bf16 CPU...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], torch_dtype=torch.bfloat16, device_map="cpu", low_cpu_mem_usage=True, trust_remote_code=True, ) print("✓ Success: bf16 CPU") except Exception as e: errors.append(("bf16 CPU", e)) print(f"✗ Failed: {e}") # ---------------------------------------------------------- # Strategy 4: FP16 CPU (fallback) # ---------------------------------------------------------- print("\n[4/4] Trying: fp16 CPU...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print("✓ Success: fp16 CPU") except Exception as e: errors.append(("fp16 CPU", e)) print(f"✗ Failed: {e}") # ------------------------------------------------------ # Strategy 5: 4-bit AS-IS (already quantized) # ---------------------------------------------------------- print("\n[5/5] Trying: 4-bit AS-IS...") try: model = AutoModelForCausalLM.from_pretrained( config["base_model"], torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print("✓ Success: 4-bit AS-IS") except Exception as e: errors.append(("4-bit AS-IS", e)) print(f"✗ Failed: {e}") msg = "\n".join( f"{name}: {err}" for name, err in errors ) raise RuntimeError( f"All loading strategies failed:\n\n{msg}" ) # 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 paths - training/data/ is relative to project root repo_root = Path(__file__).parent train_path = str(repo_root / "training" / "data" / "train.jsonl") test_path = str(repo_root / "training" / "data" / "test.jsonl") print(f"Looking for dataset at: {train_path}") dataset = load_dataset( "json", data_files={ "train": train_path, "test": test_path, }, ) # Manually wrap model with FSDP before passing to trainer try: from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from functools import partial import torch.distributed as dist # Initialize distributed process group (required for FSDP) if not dist.is_initialized(): dist.init_process_group(backend="nccl") print("✓ Distributed process group initialized") def get_auto_wrap_policy(model): from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer return partial( transformer_auto_wrap_policy, transformer_layer_cls={Qwen3_5MoeDecoderLayer}, ) # Wrap model with FSDP on GPU print("Wrapping model with FSDP on GPU...") model = FSDP( model, auto_wrap_policy=get_auto_wrap_policy(model), device_id=torch.cuda.current_device(), # Keep on current GPU mixed_precision=None, sync_module_states=False, # Model is already on GPU use_orig_params=True, ) print("✓ Model wrapped with FSDP (will be sharded across GPUs during training)") # Training arguments (no FSDP config - we handle it manually) 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", "steps"), eval_steps=config.get("eval_steps", 100), bf16=True, gradient_checkpointing=False, # FSDP handles it # No fsdp config - we wrap manually above ) # SFT Trainer from trl import SFTTrainer trainer = SFTTrainer( model=model, processing_class=tokenizer, train_dataset=dataset["train"], eval_dataset=dataset["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']}") # Cleanup distributed process group if dist.is_initialized(): dist.destroy_process_group() # Return early - success return except Exception as e: errors.append(("QLoRA 4-bit FSDP training", e)) print(f"✗ Failed during FSDP training: {e}") # Cleanup distributed process group if initialized if dist.is_initialized(): dist.destroy_process_group() # Fall through to next strategy 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()