From e3ea60e6c68f9601a89c88d82b5e8234e166fa6b Mon Sep 17 00:00:00 2001 From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com> Date: Thu, 2 Jul 2026 12:04:54 -0400 Subject: [PATCH] fix: clarify distribution pattern detection in test script --- test_model_loading.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/test_model_loading.py b/test_model_loading.py index a4290ac..122018c 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -53,11 +53,29 @@ def test_model_loading(): # 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") + # 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")