feat: add test script to verify model loading and GPU distribution

test_model_loading.py

@@ -0,0 +1,76 @@
+#!/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)