cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: clarify distribution pattern detection in test script

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 12:04:54 -04:00
parent 3049aa9b0a
commit e3ea60e6c6
1 changed files with 22 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
test_model_loading.py
View File

@@ -53,11 +53,29 @@ def test_model_loading():
# Check if model is distributed # Check if model is distributed
print("\n4. Distribution Check:") print("\n4. Distribution Check:")
print(" Model should be split across GPUs (not all on one GPU)")
# Count parameters per GPU # Get memory on each GPU
total_params = sum(p.numel() for p in model.parameters()) gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
print(f" Total parameters: {total_params / 1e9:.2f}B") 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("\n" + "=" * 80)
print("TEST PASSED: Model loaded and distributed across GPUs") print("TEST PASSED: Model loaded and distributed across GPUs")