fix: clarify distribution pattern detection in test script

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")