diff --git a/test_model_loading.py b/test_model_loading.py new file mode 100644 index 0000000..a4290ac --- /dev/null +++ b/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)