#!/usr/bin/env python3 """Fine-tune Qwen2.5-Coder-7B with Zenith framework""" import torch import argparse from pathlib import Path import sys sys.path.append(str(Path(__file__).parent)) from configs.zenith_config import get_7b_config from models.zenith_model import ZenithForCausalLM from data.openthoughts_processor import OpenThoughtsProcessor, OpenThoughtsConfig from training.trainer import Trainer, TrainerConfig from data.advanced_tokenizer import AdvancedTokenizer def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser(description="Fine-tune Qwen2.5-Coder-7B") # Model and data parser.add_argument( "--base_model", type=str, default="Qwen/Qwen2.5-Coder-7B", help="Base model from Hugging Face" ) parser.add_argument( "--train_data", type=str, required=True, help="Path to training data (JSON or HuggingFace dataset)" ) parser.add_argument( "--eval_data", type=str, default=None, help="Path to evaluation data" ) # Training hyperparameters parser.add_argument( "--epochs", type=int, default=3, help="Number of training epochs" ) parser.add_argument( "--batch_size", type=int, default=4, help="Per-device batch size" ) parser.add_argument( "--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps" ) parser.add_argument( "--learning_rate", type=float, default=2e-5, help="Learning rate" ) parser.add_argument( "--weight_decay", type=float, default=0.01, help="Weight decay" ) parser.add_argument( "--warmup_steps", type=int, default=100, help="Warmup steps" ) parser.add_argument( "--max_grad_norm", type=float, default=1.0, help="Gradient clipping norm" ) # LoRA/QLoRA settings parser.add_argument( "--use_lora", action="store_true", help="Use LoRA for efficient fine-tuning" ) 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( "--lora_dropout", type=float, default=0.1, help="LoRA dropout" ) parser.add_argument( "--use_qlora", action="store_true", help="Use 4-bit quantization with LoRA" ) # Zenith features parser.add_argument( "--use_eq_adapter", action="store_true", help="Enable EQ adapter during fine-tuning" ) parser.add_argument( "--eq_loss_weight", type=float, default=0.1, help="Weight for EQ auxiliary loss" ) parser.add_argument( "--use_moe", action="store_true", help="Enable MoE layers" ) parser.add_argument( "--num_experts", type=int, default=8, help="Number of MoE experts" ) # Data processing parser.add_argument( "--max_seq_length", type=int, default=8192, help="Maximum sequence length" ) parser.add_argument( "--use_quality_filter", action="store_true", help="Apply quality filtering" ) parser.add_argument( "--use_curriculum", action="store_true", help="Use curriculum learning" ) parser.add_argument( "--use_augmentation", action="store_true", help="Use data augmentation" ) # Output and logging parser.add_argument( "--output_dir", type=str, default="./outputs/zenith-7b-finetuned", help="Output directory for checkpoints" ) parser.add_argument( "--logging_steps", type=int, default=10, help="Logging frequency" ) parser.add_argument( "--eval_steps", type=int, default=500, help="Evaluation frequency" ) parser.add_argument( "--save_steps", type=int, default=1000, help="Checkpoint save frequency" ) # Mixed precision parser.add_argument( "--mixed_precision", type=str, default="bf16", choices=["no", "fp16", "bf16"], help="Mixed precision training" ) # Resume training parser.add_argument( "--resume_from_checkpoint", type=str, default=None, help="Resume from checkpoint" ) return parser.parse_args() def main(): args = parse_args() # Create output directory output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) # Load base configuration config = get_7b_config() # Apply fine-tuning specific settings config.use_lora = args.use_lora config.lora_r = args.lora_r config.lora_alpha = args.lora_alpha config.lora_dropout = args.lora_dropout config.use_qlora = args.use_qlora config.use_eq_adapter = args.use_eq_adapter config.eq_loss_weight = args.eq_loss_weight config.use_moe = args.use_moe config.num_experts = args.num_experts config.max_seq_len = args.max_seq_length # Initialize tokenizer from base model print(f"Loading tokenizer from {args.base_model}...") tokenizer = AdvancedTokenizer.from_pretrained(args.base_model) config.vocab_size = tokenizer.get_vocab_size() # Load base model print(f"Loading base model: {args.base_model}") model = ZenithForCausalLM.from_pretrained( args.base_model, config=config, device_map="auto" ) # Prepare OpenThoughts data processor ot_config = OpenThoughtsConfig( dataset_name="open-thoughts/OpenThoughts3-1.2M", streaming=True, max_seq_length=args.max_seq_length, quality_filtering=args.use_quality_filter, curriculum_learning=args.use_curriculum, augmentation=args.use_augmentation, tokenizer=tokenizer ) data_processor = OpenThoughtsProcessor(ot_config) # Load datasets print("Loading and preprocessing datasets...") train_dataset = data_processor.load_dataset(args.train_data, split="train") eval_dataset = None if args.eval_data: eval_dataset = data_processor.load_dataset(args.eval_data, split="validation") # Create dataloaders train_dataloader = data_processor.create_dataloader( train_dataset, batch_size=args.batch_size, shuffle=True ) eval_dataloader = None if eval_dataset: eval_dataloader = data_processor.create_dataloader( eval_dataset, batch_size=args.batch_size, shuffle=False ) # Trainer configuration trainer_config = TrainerConfig( output_dir=args.output_dir, num_train_epochs=args.epochs, per_device_train_batch_size=args.batch_size, per_device_eval_batch_size=args.batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, learning_rate=args.learning_rate, weight_decay=args.weight_decay, warmup_steps=args.warmup_steps, max_grad_norm=args.max_grad_norm, mixed_precision=args.mixed_precision, logging_steps=args.logging_steps, eval_steps=args.eval_steps, save_steps=args.save_steps, save_total_limit=3, report_to="tensorboard", run_name=f"zenith-7b-finetune-{args.base_model.split('/')[-1]}" ) # Initialize trainer print("Initializing trainer...") trainer = Trainer( model=model, config=trainer_config, train_dataloader=train_dataloader, eval_dataloader=eval_dataloader, tokenizer=tokenizer ) # Resume from checkpoint if specified if args.resume_from_checkpoint: print(f"Resuming from checkpoint: {args.resume_from_checkpoint}") trainer.load_checkpoint(args.resume_from_checkpoint) # Start training print("Starting fine-tuning...") trainer.train() # Save final model print(f"Saving final model to {output_dir}") model.save_pretrained(output_dir) tokenizer.save_pretrained(output_dir) print("Fine-tuning complete!") if __name__ == "__main__": main()