#!/usr/bin/env python3
"""
Test model loading and GPU distribution without training.
"""

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

def test_model_loading():
    print("=" * 80)
    print("Testing model loading and GPU distribution")
    print("=" * 80)
    
    # Check GPU availability
    print(f"\n1. GPU Check:")
    print(f"   CUDA available: {torch.cuda.is_available()}")
    print(f"   GPU count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        print(f"   GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
    
    # Test 1: Load with device_map="auto" (distributed)
    print("\n2. Test 1: Load with device_map='auto' (should distribute across GPUs)")
    try:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        
        print("   Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        # Check memory usage
        print("\n3. Memory Usage:")
        for i in range(torch.cuda.device_count()):
            mem_allocated = torch.cuda.memory_allocated(i) / 1e9
            mem_reserved = torch.cuda.memory_reserved(i) / 1e9
            total = torch.cuda.get_device_properties(i).total_memory / 1e9
            print(f"   GPU {i}:")
            print(f"      Allocated: {mem_allocated:.2f} GB")
            print(f"      Reserved: {mem_reserved:.2f} GB")
            print(f"      Total: {total:.2f} GB")
            print(f"      Free: {total - mem_allocated:.2f} GB")
        
        # Check if model is distributed
        print("\n4. Distribution Check:")
        print("   Model should be split across GPUs (not all on one GPU)")
        
        # Count parameters per GPU
        total_params = sum(p.numel() for p in model.parameters())
        print(f"   Total parameters: {total_params / 1e9:.2f}B")
        
        print("\n" + "=" * 80)
        print("TEST PASSED: Model loaded and distributed across GPUs")
        print("=" * 80)
        
    except Exception as e:
        print(f"\n✗ Test 1 FAILED: {e}")
        print("\nThis means the model is NOT being distributed properly!")
        print("It might be trying to fit the entire model on one GPU.")
        return False
    
    return True

if __name__ == "__main__":
    success = test_model_loading()
    exit(0 if success else 1)