""" Llama 3.2-1B MMLU Evaluation Script (Laptop Optimized with Quantization) This script evaluates Llama 3.2-1B on the MMLU benchmark. Optimized for laptops with 4-bit or 8-bit quantization to reduce memory usage. Quantization options: - 4-bit: ~1.5 GB VRAM/RAM (default for laptop) - 8-bit: ~2.5 GB VRAM/RAM - No quantization: ~5 GB VRAM/RAM Usage: 1. Install: pip install transformers torch datasets accelerate tqdm bitsandbytes 2. Login: huggingface-cli login 3. Run: python llama_mmlu_eval_quantized.py Set QUANTIZATION_BITS below to choose quantization level. """ import torch from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from datasets import load_dataset import json from tqdm.auto import tqdm import os from datetime import datetime import sys import platform # ============================================================================ # CONFIGURATION - Modify these settings # ============================================================================ MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" # GPU settings # If True, will attempt to use the best available GPU (CUDA for NVIDIA, MPS for Apple Silicon) # If False, will always use CPU regardless of available hardware USE_GPU = True # Set to False to force CPU-only execution MAX_NEW_TOKENS = 1 # Quantization settings # Options: 4, 8, or None (default is None for full precision) # # To enable quantization, change QUANTIZATION_BITS to one of the following: # QUANTIZATION_BITS = 4 # 4-bit quantization: ~1.5 GB memory (most memory efficient) # QUANTIZATION_BITS = 8 # 8-bit quantization: ~2.5 GB memory (balanced quality/memory) # QUANTIZATION_BITS = None # No quantization: ~5 GB memory (full precision, best quality) # # Notes: # - Quantization requires the 'bitsandbytes' package: pip install bitsandbytes # - Quantization only works with CUDA (NVIDIA GPUs), not with Apple Metal (MPS) # - If using Apple Silicon, quantization will be automatically disabled QUANTIZATION_BITS = None # Change to 4 or 8 to enable quantization # For quick testing, you can reduce this list MMLU_SUBJECTS = [ # "abstract_algebra", "anatomy", "astronomy", "business_ethics", # "clinical_knowledge", "college_biology", "college_chemistry", # "college_computer_science", "college_mathematics", "college_medicine", # "college_physics", "computer_security", "conceptual_physics", # "econometrics", "electrical_engineering", "elementary_mathematics", # "formal_logic", "global_facts", "high_school_biology", # "high_school_chemistry", "high_school_computer_science", # "high_school_european_history", "high_school_geography", # "high_school_government_and_politics", "high_school_macroeconomics", # "high_school_mathematics", "high_school_microeconomics", # "high_school_physics", "high_school_psychology", "high_school_statistics", # "high_school_us_history", "high_school_world_history", "human_aging", # "human_sexuality", "international_law", "jurisprudence", # "logical_fallacies", "machine_learning", "management", "marketing", # "medical_genetics", "miscellaneous", "moral_disputes", "moral_scenarios", # "nutrition", "philosophy", "prehistory", "professional_accounting", # "professional_law", "professional_medicine", "professional_psychology", # "public_relations", "security_studies", "sociology", "us_foreign_policy", # "virology", "world_religions" ] def detect_device(): """Detect the best available device (CUDA, MPS, or CPU)""" # If GPU is disabled, always use CPU if not USE_GPU: return "cpu" # Check for CUDA if torch.cuda.is_available(): return "cuda" # Check for Apple Silicon with Metal if torch.backends.mps.is_available(): # Check if we're actually on Apple ARM is_apple_arm = platform.system() == "Darwin" and platform.processor() == "arm" if is_apple_arm: # Metal is available but incompatible with quantization if QUANTIZATION_BITS is not None: print("\n" + "="*70) print("ERROR: Metal and Quantization Conflict") print("="*70) print("Metal Performance Shaders (MPS) is incompatible with quantization.") print(f"You have USE_GPU = True and QUANTIZATION_BITS = {QUANTIZATION_BITS}") print("") print("Please choose one of the following options:") print(" 1. Set USE_GPU = False to use CPU with quantization") print(" 2. Set QUANTIZATION_BITS = None to use Metal without quantization") print("="*70 + "\n") sys.exit(1) return "mps" # Default to CPU return "cpu" def check_environment(): global QUANTIZATION_BITS """Check environment and dependencies""" print("="*70) print("Environment Check") print("="*70) # Check if in Colab try: import google.colab print("✓ Running in Google Colab") in_colab = True except: print("✓ Running locally (not in Colab)") in_colab = False # Check system info print(f"✓ Platform: {platform.system()} ({platform.machine()})") if platform.system() == "Darwin": print(f"✓ Processor: {platform.processor()}") # Detect and set device device = detect_device() # Check device if device == "cuda": gpu_name = torch.cuda.get_device_name(0) gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9 print(f"✓ GPU Available: {gpu_name}") print(f"✓ GPU Memory: {gpu_memory:.2f} GB") elif device == "mps": print("✓ Apple Metal (MPS) Available") print("✓ Using Metal Performance Shaders for GPU acceleration") else: print("⚠️ No GPU detected - running on CPU") # Check quantization support if QUANTIZATION_BITS is not None: try: import bitsandbytes print(f"✓ bitsandbytes installed - {QUANTIZATION_BITS}-bit quantization available") except ImportError: print(f"❌ bitsandbytes NOT installed - cannot use quantization") sys.exit(1) if device == 'mps': print(f"❌ Apple METAL is incompatible with quantization") print("✓ Quantization disabled - loading full precision model") QUANTIZATION_BITS = None sys.exit(1) else: print("✓ Quantization disabled - loading full precision model") # Check HF authentication try: from huggingface_hub import HfFolder token = HfFolder.get_token() if token: print("✓ Hugging Face authenticated") else: print("⚠️ No Hugging Face token found") print("Run: huggingface-cli login") except: print("⚠️ Could not check Hugging Face authentication") # Print configuration print("\n" + "="*70) print("Configuration") print("="*70) print(f"Model: {MODEL_NAME}") print(f"Device: {device}") if QUANTIZATION_BITS is not None: print(f"Quantization: {QUANTIZATION_BITS}-bit") if QUANTIZATION_BITS == 4: print(f"Expected memory: ~1.5 GB") elif QUANTIZATION_BITS == 8: print(f"Expected memory: ~2.5 GB") else: print(f"Quantization: None (full precision)") if device == "cuda": print(f"Expected memory: ~2.5 GB (FP16)") elif device == "mps": print(f"Expected memory: ~2.5 GB (FP16)") else: print(f"Expected memory: ~5 GB (FP32)") print(f"Number of subjects: {len(MMLU_SUBJECTS)}") print("="*70 + "\n") return in_colab, device def get_quantization_config(): """Create quantization config based on settings""" if QUANTIZATION_BITS is None: return None if QUANTIZATION_BITS == 4: # 4-bit quantization (most memory efficient) config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, # Double quantization for extra compression bnb_4bit_quant_type="nf4" # NormalFloat4 - better for LLMs ) print("Using 4-bit quantization (NF4 + double quant)") print("Memory usage: ~1.5 GB") elif QUANTIZATION_BITS == 8: # 8-bit quantization (balanced) config = BitsAndBytesConfig( load_in_8bit=True, llm_int8_threshold=6.0, llm_int8_has_fp16_weight=False ) print("Using 8-bit quantization") print("Memory usage: ~2.5 GB") else: raise ValueError(f"Invalid QUANTIZATION_BITS: {QUANTIZATION_BITS}. Use 4, 8, or None") return config def load_model_and_tokenizer(device): """Load Llama model with optional quantization""" print(f"\nLoading model {MODEL_NAME}...") print(f"Device: {device}") try: # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) print("✓ Tokenizer loaded") # Get quantization config quant_config = get_quantization_config() # Load model print("Loading model (this may take 2-3 minutes)...") if quant_config is not None: # Quantized model loading (only works with CUDA) model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, quantization_config=quant_config, device_map="auto", low_cpu_mem_usage=True ) else: # Non-quantized model loading if device == "cuda": model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True ) elif device == "mps": model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, dtype=torch.float16, low_cpu_mem_usage=True ) model = model.to(device) else: # CPU model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, dtype=torch.float32, low_cpu_mem_usage=True ) model = model.to(device) model.eval() # Print model info print("✓ Model loaded successfully!") print(f" Model device: {next(model.parameters()).device}") print(f" Model dtype: {next(model.parameters()).dtype}") # Check memory usage if torch.cuda.is_available(): memory_allocated = torch.cuda.memory_allocated(0) / 1e9 memory_reserved = torch.cuda.memory_reserved(0) / 1e9 print(f" GPU Memory: {memory_allocated:.2f} GB allocated, {memory_reserved:.2f} GB reserved") # Check if using quantization if quant_config is not None: print(f" Quantization: {QUANTIZATION_BITS}-bit active") elif device == "mps": print(f" Running on Apple Metal (MPS)") return model, tokenizer except Exception as e: print(f"\n❌ Error loading model: {e}") print("\nPossible causes:") print("1. No Hugging Face token - Run: huggingface-cli login") print("2. Llama license not accepted - Visit: https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct") print("3. bitsandbytes not installed - Run: pip install bitsandbytes") print("4. Out of memory - Try 4-bit quantization or smaller model") raise def format_mmlu_prompt(question, choices): """Format MMLU question as multiple choice""" choice_labels = ["A", "B", "C", "D"] prompt = f"{question}\n\n" for label, choice in zip(choice_labels, choices): prompt += f"{label}. {choice}\n" prompt += "\nAnswer:" return prompt def get_model_prediction(model, tokenizer, prompt): """Get model's prediction for multiple-choice question""" inputs = tokenizer(prompt, return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=MAX_NEW_TOKENS, pad_token_id=tokenizer.eos_token_id, do_sample=False, temperature=1.0 ) generated_text = tokenizer.decode( outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True ) answer = generated_text.strip()[:1].upper() if answer not in ["A", "B", "C", "D"]: for char in generated_text.upper(): if char in ["A", "B", "C", "D"]: answer = char break else: answer = "A" return answer def evaluate_subject(model, tokenizer, subject): """Evaluate model on a specific MMLU subject""" print(f"\n{'='*70}") print(f"Evaluating subject: {subject}") print(f"{'='*70}") try: dataset = load_dataset("cais/mmlu", subject, split="test") except Exception as e: print(f"❌ Error loading subject {subject}: {e}") return None correct = 0 total = 0 for example in tqdm(dataset, desc=f"Testing {subject}", leave=True): question = example["question"] choices = example["choices"] correct_answer_idx = example["answer"] correct_answer = ["A", "B", "C", "D"][correct_answer_idx] prompt = format_mmlu_prompt(question, choices) predicted_answer = get_model_prediction(model, tokenizer, prompt) if predicted_answer == correct_answer: correct += 1 total += 1 accuracy = (correct / total * 100) if total > 0 else 0 print(f"✓ Result: {correct}/{total} correct = {accuracy:.2f}%") return { "subject": subject, "correct": correct, "total": total, "accuracy": accuracy } def main(): """Main evaluation function""" print("\n" + "="*70) print("Llama 3.2-1B MMLU Evaluation (Quantized)") print("="*70 + "\n") # Check environment in_colab, device = check_environment() # Load model model, tokenizer = load_model_and_tokenizer(device) # Evaluate results = [] total_correct = 0 total_questions = 0 print(f"\n{'='*70}") print(f"Starting evaluation on {len(MMLU_SUBJECTS)} subjects") print(f"{'='*70}\n") start_time = datetime.now() for i, subject in enumerate(MMLU_SUBJECTS, 1): print(f"\nProgress: {i}/{len(MMLU_SUBJECTS)} subjects") result = evaluate_subject(model, tokenizer, subject) if result: results.append(result) total_correct += result["correct"] total_questions += result["total"] end_time = datetime.now() duration = (end_time - start_time).total_seconds() # Calculate overall accuracy overall_accuracy = (total_correct / total_questions * 100) if total_questions > 0 else 0 # Print summary print("\n" + "="*70) print("EVALUATION SUMMARY") print("="*70) print(f"Model: {MODEL_NAME}") print(f"Quantization: {QUANTIZATION_BITS}-bit" if QUANTIZATION_BITS else "None (full precision)") print(f"Total Subjects: {len(results)}") print(f"Total Questions: {total_questions}") print(f"Total Correct: {total_correct}") print(f"Overall Accuracy: {overall_accuracy:.2f}%") print(f"Duration: {duration/60:.1f} minutes") print("="*70) # Save results timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") quant_suffix = f"_{QUANTIZATION_BITS}bit" if QUANTIZATION_BITS else "_full" output_file = f"llama_3.2_1b_mmlu_results{quant_suffix}_{timestamp}.json" output_data = { "model": MODEL_NAME, "quantization_bits": QUANTIZATION_BITS, "timestamp": timestamp, "device": str(device), "duration_seconds": duration, "overall_accuracy": overall_accuracy, "total_correct": total_correct, "total_questions": total_questions, "subject_results": results } with open(output_file, "w") as f: json.dump(output_data, f, indent=2) print(f"\n✓ Results saved to: {output_file}") # Print top/bottom subjects if len(results) > 0: sorted_results = sorted(results, key=lambda x: x["accuracy"], reverse=True) print("\n📊 Top 5 Subjects:") for i, result in enumerate(sorted_results[:5], 1): print(f" {i}. {result['subject']}: {result['accuracy']:.2f}%") print("\n📉 Bottom 5 Subjects:") for i, result in enumerate(sorted_results[-5:], 1): print(f" {i}. {result['subject']}: {result['accuracy']:.2f}%") # Colab-specific instructions if in_colab: print("\n" + "="*70) print("💾 To download results in Colab:") print("="*70) print(f"from google.colab import files") print(f"files.download('{output_file}')") print("\n✅ Evaluation complete!") return output_file if __name__ == "__main__": try: output_file = main() except KeyboardInterrupt: print("\n\n⚠️ Evaluation interrupted by user") except Exception as e: print(f"\n❌ Error during evaluation: {e}") import traceback traceback.print_exc()