95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
#!/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-4bit", # Use already-quantized 4-bit model
|
|
device_map="auto",
|
|
torch_dtype=torch.float16,
|
|
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:")
|
|
|
|
# Get memory on each GPU
|
|
gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
|
|
gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
|
|
|
|
print(f" GPU 0 memory: {gpu0_mem:.2f} GB")
|
|
print(f" GPU 1 memory: {gpu1_mem:.2f} GB")
|
|
|
|
# Determine distribution pattern
|
|
print("\n5. Distribution Pattern:")
|
|
if abs(gpu0_mem - gpu1_mem) < 1.0: # Within 1GB
|
|
if gpu0_mem < 10.0: # Less than 10GB each
|
|
print(" ✓ DISTRIBUTED: Model split across both GPUs")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB of the model")
|
|
else:
|
|
print(" ⚠ DUPLICATE: Same model loaded on both GPUs")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB (wasteful but fits)")
|
|
else:
|
|
print(" ✗ NOT DISTRIBUTED: Model on one GPU only")
|
|
if gpu0_mem > gpu1_mem:
|
|
print(f" GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
|
|
else:
|
|
print(f" GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
|
|
|
|
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)
|