""" NOVA-EVO: Genetic algorithm for hyperparameter and architecture search """ import random import json from pathlib import Path from typing import List, Tuple, Optional import time from tqdm import tqdm import copy from .config import EvolutionConfig, Individual from .fitness import FitnessEvaluator class EvolutionEngine: """ Genetic algorithm engine for evolving NOVA configurations Features: - Multi-objective fitness (loss, latency, memory, quality) - Elitism with Pareto selection - Mutation and crossover - Hall of Fame for best individuals - Rollback on regression """ def __init__( self, config: EvolutionConfig, fitness_evaluator: FitnessEvaluator, ): """ Args: config: Evolution configuration fitness_evaluator: Fitness evaluation engine """ self.config = config self.evaluator = fitness_evaluator # Population self.population: List[Individual] = [] self.generation = 0 # Hall of Fame - best individuals self.hall_of_fame: List[Individual] = [] self.max_hof_size = 10 # Tracking self.evolution_history = [] self.start_time = None # Setup Path(config.save_dir).mkdir(parents=True, exist_ok=True) random.seed(config.seed) def initialize_population(self) -> List[Individual]: """Create initial random population""" print(f"Initializing population of {self.config.population_size}...") population = [] for i in range(self.config.population_size): individual = Individual( learning_rate=random.uniform(self.config.lr_min, self.config.lr_max) if self.config.search_learning_rate else 3e-4, batch_size=random.choice(self.config.batch_size_options) if self.config.search_batch_size else 8, warmup_steps=random.randint(self.config.warmup_min, self.config.warmup_max) if self.config.search_warmup_steps else 1000, weight_decay=random.uniform(self.config.wd_min, self.config.wd_max) if self.config.search_weight_decay else 0.1, rope_theta=random.choice(self.config.rope_theta_options) if self.config.search_rope_theta else 10000.0, hidden_act=random.choice(self.config.activation_options) if self.config.search_activation else "swiglu", norm_type=random.choice(self.config.norm_options) if self.config.search_norm else "rmsnorm", generation=0, ) population.append(individual) return population def evaluate_population(self, population: List[Individual]) -> List[Individual]: """Evaluate fitness for all individuals in population""" print(f"\nEvaluating {len(population)} individuals...") for idx, individual in enumerate(tqdm(population, desc="Evaluating")): # Skip if already evaluated if individual.fitness is not None: continue # Evaluate metrics = self.evaluator.evaluate(individual) # Store metrics individual.loss = metrics['loss'] individual.perplexity = metrics.get('perplexity') individual.latency_ms = metrics.get('latency_ms') individual.memory_mb = metrics.get('memory_mb') individual.quality_score = metrics.get('quality_score', 0.0) # Calculate multi-objective fitness individual.fitness = self._calculate_fitness(individual) return population def _calculate_fitness(self, individual: Individual) -> float: """ Calculate multi-objective fitness score Lower is better (we're minimizing) """ fitness = 0.0 # Loss component (lower is better) if individual.loss is not None: fitness += individual.loss * self.config.loss_weight # Latency component (lower is better, normalized) if individual.latency_ms is not None: normalized_latency = individual.latency_ms / 1000.0 # Normalize to seconds fitness += normalized_latency * self.config.latency_weight # Memory component (lower is better, normalized) if individual.memory_mb is not None: normalized_memory = individual.memory_mb / 1000.0 # Normalize to GB fitness += normalized_memory * self.config.memory_weight # Quality component (higher is better, so negate) if individual.quality_score is not None: fitness -= individual.quality_score * self.config.quality_weight return fitness def select_parents(self, population: List[Individual]) -> List[Individual]: """ Select parents for next generation using elitism Args: population: Current population (should be evaluated) Returns: Elite individuals to keep """ # Sort by fitness (lower is better) sorted_pop = sorted(population, key=lambda x: x.fitness if x.fitness is not None else float('inf')) # Select top performers num_elite = max(1, int(len(population) * self.config.elite_ratio)) elite = sorted_pop[:num_elite] return elite def crossover(self, parent1: Individual, parent2: Individual) -> Individual: """ Create offspring by combining two parents Uses uniform crossover - randomly picks from each parent """ child = Individual( learning_rate=random.choice([parent1.learning_rate, parent2.learning_rate]), batch_size=random.choice([parent1.batch_size, parent2.batch_size]), warmup_steps=random.choice([parent1.warmup_steps, parent2.warmup_steps]), weight_decay=random.choice([parent1.weight_decay, parent2.weight_decay]), rope_theta=random.choice([parent1.rope_theta, parent2.rope_theta]), hidden_act=random.choice([parent1.hidden_act, parent2.hidden_act]), norm_type=random.choice([parent1.norm_type, parent2.norm_type]), generation=self.generation + 1, parent_ids=[id(parent1), id(parent2)], ) return child def mutate(self, individual: Individual) -> Individual: """ Mutate an individual with random changes Args: individual: Individual to mutate Returns: Mutated copy """ mutated = copy.deepcopy(individual) mutated.generation = self.generation + 1 # Mutate each gene with some probability if random.random() < self.config.mutation_rate: mutated.learning_rate = random.uniform(self.config.lr_min, self.config.lr_max) if random.random() < self.config.mutation_rate: mutated.batch_size = random.choice(self.config.batch_size_options) if random.random() < self.config.mutation_rate: mutated.warmup_steps = random.randint(self.config.warmup_min, self.config.warmup_max) if random.random() < self.config.mutation_rate: mutated.weight_decay = random.uniform(self.config.wd_min, self.config.wd_max) if random.random() < self.config.mutation_rate: mutated.rope_theta = random.choice(self.config.rope_theta_options) if random.random() < self.config.mutation_rate: mutated.hidden_act = random.choice(self.config.activation_options) if random.random() < self.config.mutation_rate: mutated.norm_type = random.choice(self.config.norm_options) # Reset fitness (needs re-evaluation) mutated.fitness = None mutated.loss = None return mutated def create_next_generation(self, parents: List[Individual]) -> List[Individual]: """Create next generation from parents""" next_gen = [] # Keep elite unchanged next_gen.extend(copy.deepcopy(parents)) # Fill rest with offspring while len(next_gen) < self.config.population_size: # Select two random parents parent1, parent2 = random.sample(parents, 2) # Crossover child = self.crossover(parent1, parent2) # Mutate child = self.mutate(child) next_gen.append(child) return next_gen def update_hall_of_fame(self, population: List[Individual]): """Update hall of fame with best individuals""" # Add current best to hall of fame for ind in population: if ind.fitness is not None: self.hall_of_fame.append(copy.deepcopy(ind)) # Sort by fitness self.hall_of_fame.sort(key=lambda x: x.fitness if x.fitness is not None else float('inf')) # Keep only top N self.hall_of_fame = self.hall_of_fame[:self.max_hof_size] def save_checkpoint(self): """Save evolution state""" checkpoint_path = Path(self.config.save_dir) / f"generation_{self.generation}.json" checkpoint = { 'generation': self.generation, 'population': [ind.to_dict() for ind in self.population], 'hall_of_fame': [ind.to_dict() for ind in self.hall_of_fame], 'config': self.config.__dict__, } with open(checkpoint_path, 'w') as f: json.dump(checkpoint, f, indent=2) print(f" Checkpoint saved: {checkpoint_path}") def run(self): """Run the evolution process""" print("=" * 60) print("NOVA-EVO: Genetic Algorithm Evolution") print("=" * 60) self.start_time = time.time() # Initialize population self.population = self.initialize_population() # Evolution loop for gen in range(self.config.num_generations): self.generation = gen print(f"\n{'='*60}") print(f"Generation {gen + 1}/{self.config.num_generations}") print(f"{'='*60}") # Evaluate self.population = self.evaluate_population(self.population) # Select parents parents = self.select_parents(self.population) # Update hall of fame self.update_hall_of_fame(self.population) # Report best individual best = self.hall_of_fame[0] if self.hall_of_fame else None if best: print(f"\nšŸ† Best individual so far:") print(f" Fitness: {best.fitness:.4f}") print(f" Loss: {best.loss:.4f}") print(f" LR: {best.learning_rate:.2e}, BS: {best.batch_size}") print(f" Activation: {best.hidden_act}, Norm: {best.norm_type}") # Checkpoint if (gen + 1) % self.config.checkpoint_every_n_generations == 0: self.save_checkpoint() # Create next generation if gen < self.config.num_generations - 1: self.population = self.create_next_generation(parents) # Final checkpoint self.save_checkpoint() print("\n" + "=" * 60) print("Evolution Complete!") print("=" * 60) print(f"Total time: {(time.time() - self.start_time) / 3600:.2f} hours") print(f"\nTop 3 individuals:") for i, ind in enumerate(self.hall_of_fame[:3]): print(f"\n{i+1}. Fitness: {ind.fitness:.4f}") print(f" Loss: {ind.loss:.4f}, LR: {ind.learning_rate:.2e}") print(f" Batch size: {ind.batch_size}, Warmup: {ind.warmup_steps}") print(f" Activation: {ind.hidden_act}, Norm: {ind.norm_type}")