#!/usr/bin/env python3
"""
Test model loading and GPU distribution without training.
Tests multiple strategies to find what works.
"""

import torch
from transformers import AutoModelForCausalLM

def test_strategy_1():
    """Test 1: Load with device_map='auto' (no FSDP)"""
    print("\n" + "=" * 80)
    print("TEST 1: Load with device_map='auto' (no FSDP)")
    print("=" * 80)
    
    try:
        print("   Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B-4bit",
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        # Check memory usage
        print("\n   Memory Usage:")
        for i in range(torch.cuda.device_count()):
            mem_allocated = torch.cuda.memory_allocated(i) / 1e9
            total = torch.cuda.get_device_properties(i).total_memory / 1e9
            print(f"   GPU {i}: {mem_allocated:.2f} GB / {total:.2f} GB")
        
        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
        
        print("\n   Distribution Pattern:")
        if abs(gpu0_mem - gpu1_mem) < 1.0:
            if gpu0_mem < 10.0:
                print("   ✓ DISTRIBUTED: Model split across both GPUs")
                return True
            else:
                print("   ⚠ DUPLICATE: Same model on both GPUs")
                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
                return False
        else:
            print("   ✗ NOT DISTRIBUTED: Model on one GPU only")
            return False
        
    except Exception as e:
        print(f"\n   ✗ Test 1 FAILED: {e}")
        return False

def test_strategy_2():
    """Test 2: Load with device_map='auto' then wrap with FSDP"""
    print("\n" + "=" * 80)
    print("TEST 2: Load with device_map='auto' then wrap with FSDP")
    print("=" * 80)
    
    try:
        import torch.distributed as dist
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
        from functools import partial
        
        # Initialize distributed process group
        if not dist.is_initialized():
            dist.init_process_group(backend="nccl")
            print("   ✓ Distributed process group initialized")
        
        # Clear GPU memory
        torch.cuda.empty_cache()
        
        print("   Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B-4bit",
            device_map="auto",
            torch_dtype=torch.float16,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded to GPU")
        
        # Check memory before FSDP
        print("\n   Memory BEFORE FSDP:")
        for i in range(torch.cuda.device_count()):
            mem = torch.cuda.memory_allocated(i) / 1e9
            print(f"   GPU {i}: {mem:.2f} GB")
        
        # Define auto wrap policy
        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 with FSDP
        print("\n   Wrapping with FSDP...")
        model = FSDP(
            model,
            auto_wrap_policy=get_auto_wrap_policy(model),
            device_id=torch.cuda.current_device(),
            mixed_precision=None,
            sync_module_states=False,
            use_orig_params=True,
        )
        print("   ✓ Model wrapped with FSDP")
        
        # Check memory after FSDP
        print("\n   Memory AFTER FSDP (should be sharded):")
        for i in range(torch.cuda.device_count()):
            mem = torch.cuda.memory_allocated(i) / 1e9
            total = torch.cuda.get_device_properties(i).total_memory / 1e9
            print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
        
        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
        
        print("\n   Distribution Pattern:")
        if abs(gpu0_mem - gpu1_mem) < 2.0:  # Within 2GB (more lenient for FSDP)
            if gpu0_mem < 20.0:  # Less than 20GB each (sharded)
                print("   ✓ SUCCESSFULLY SHARDED: Model split across GPUs")
                print(f"      Each GPU has ~{gpu0_mem:.2f}GB (down from ~31GB)")
                return True
            else:
                print("   ⚠ NOT SHARDED: Still duplicate loading")
                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
                return False
        else:
            print("   ✗ FAILED: Uneven distribution")
            return False
        
    except Exception as e:
        print(f"\n   ✗ Test 2 FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("=" * 80)
    print("Testing model loading strategies")
    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 Strategy 1
    strategy1_ok = test_strategy_1()
    
    # Test Strategy 2 (FSDP)
    strategy2_ok = test_strategy_2()
    
    # Summary
    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)
    if strategy1_ok:
        print("✓ Strategy 1 (device_map=auto): WORKS - model distributed")
    else:
        print("✗ Strategy 1 (device_map=auto): FAILED - model not distributed")
    
    if strategy2_ok:
        print("✓ Strategy 2 (FSDP): WORKS - model successfully sharded")
    else:
        print("✗ Strategy 2 (FSDP): FAILED - model not sharded")
    
    if strategy1_ok or strategy2_ok:
        print("\n✓ At least one strategy works!")
        exit(0)
    else:
        print("\n✗ No strategy works - model cannot be distributed")
        exit(1)