feat: comprehensive test of 5 loading strategies

test_model_loading.py

@@ -1,16 +1,39 @@
 #!/usr/bin/env python3
 """
-Test model loading and GPU distribution without training.
-Tests multiple strategies to find what works.
+Test multiple model loading strategies to find what works.
+Each strategy is tested independently.
 """
 
 import torch
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+def check_gpu_memory():
+    """Check memory usage on all GPUs."""
+    print("   Memory Usage:")
+    for i in range(torch.cuda.device_count()):
+        mem = torch.cuda.memory_allocated(i) / 1e9
+        total = torch.cuda.get_device_properties(i).total_memory / 1e9
+        print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
+    
+    gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
+    gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
+    
+    # Determine pattern
+    if abs(gpu0_mem - gpu1_mem) < 2.0:  # Within 2GB
+        if gpu0_mem < 15.0:
+            return "DISTRIBUTED"
+        else:
+            return "DUPLICATE"
+    else:
+        if gpu0_mem > gpu1_mem:
+            return f"GPU0_ONLY ({gpu0_mem:.1f}GB)"
+        else:
+            return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
 
 def test_strategy_1():
-    """Test 1: Load with device_map='auto' (no FSDP)"""
+    """Test 1: device_map='auto' (no quantization config)"""
     print("\n" + "=" * 80)
-    print("TEST 1: Load with device_map='auto' (no FSDP)")
+    print("TEST 1: device_map='auto' (no BnB)")
     print("=" * 80)
     
     try:
@@ -24,163 +47,175 @@ def test_strategy_1():
         )
         print("   ✓ Model loaded successfully")
         
-        # Test 1b: Try with load_in_4bit=True (force quantization)
-        print("\n   Testing with load_in_4bit=True (force quantization)...")
-        torch.cuda.empty_cache()
+        pattern = check_gpu_memory()
+        print(f"\n   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        return False, str(e)
 
-        from transformers import BitsAndBytesConfig
+def test_strategy_2():
+    """Test 2: device_map='auto' with BnB 4-bit"""
+    print("\n" + "=" * 80)
+    print("TEST 2: device_map='auto' + BnB 4-bit")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Loading model with BnB 4-bit...")
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=torch.bfloat16,
         )
-        
-        model2 = AutoModelForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             "/data/models/Ornith-1.0-35B-4bit",
             quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
-        print("   ✓ Model loaded with BnB 4-bit")
-        
-        print("\n   Memory with BnB 4-bit:")
-        for i in range(torch.cuda.device_count()):
-            mem = torch.cuda.memory_allocated(i) / 1e9
-            print(f"   GPU {i}: {mem:.2f} GB")
-        
-        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
-        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
-        
-        print("\n   Distribution Pattern:")
-        if abs(gpu0_mem - gpu1_mem) < 1.0:
-            if gpu0_mem < 10.0:
-                print("   ✓ DISTRIBUTED: Model split across both GPUs")
-                return True
-            else:
-                print("   ⚠ DUPLICATE: Same model on both GPUs")
-                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
-                return False
-        else:
-            print("   ✗ NOT DISTRIBUTED: Model on one GPU only")
-            return False
-        
-        # Check memory usage
-        print("\n   Memory Usage:")
-        for i in range(torch.cuda.device_count()):
-            mem_allocated = torch.cuda.memory_allocated(i) / 1e9
-            total = torch.cuda.get_device_properties(i).total_memory / 1e9
-            print(f"   GPU {i}: {mem_allocated:.2f} GB / {total:.2f} GB")
-        
-        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
-        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
-        
-        print("\n   Distribution Pattern:")
-        if abs(gpu0_mem - gpu1_mem) < 1.0:
-            if gpu0_mem < 10.0:
-                print("   ✓ DISTRIBUTED: Model split across both GPUs")
-                return True
-            else:
-                print("   ⚠ DUPLICATE: Same model on both GPUs")
-                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
-                return False
-        else:
-            print("   ✗ NOT DISTRIBUTED: Model on one GPU only")
-            return False
+        print("   ✓ Model loaded successfully")
         
+        pattern = check_gpu_memory()
+        print(f"\n   Pattern: {pattern}")
+        return True, pattern
     except Exception as e:
-        print(f"\n   ✗ Test 1 FAILED: {e}")
-        return False
+        print(f"\n   ✗ FAILED: {e}")
+        return False, str(e)
 
-def test_strategy_2():
-    """Test 2: Load with device_map='auto' then wrap with FSDP"""
+def test_strategy_3():
+    """Test 3: device_map with explicit GPU assignment"""
     print("\n" + "=" * 80)
-    print("TEST 2: Load with device_map='auto' then wrap with FSDP")
+    print("TEST 3: device_map with explicit GPU assignment")
     print("=" * 80)
     
     try:
-        import torch.distributed as dist
-        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
-        from functools import partial
-        
-        # Initialize distributed process group
-        if not dist.is_initialized():
-            dist.init_process_group(backend="nccl")
-            print("   ✓ Distributed process group initialized")
-        
-        # Clear GPU memory
         torch.cuda.empty_cache()
+        print("   Loading model with explicit device_map...")
         
-        print("   Loading model...")
+        # Get model config to determine layers
+        from transformers import AutoConfig
+        config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B-4bit", trust_remote_code=True)
+        num_layers = config.num_hidden_layers
+        
+        # Split layers: first half on GPU 0, second half on GPU 1
+        device_map = {}
+        for i in range(num_layers):
+            if i < num_layers // 2:
+                device_map[f"model.layers.{i}"] = 0
+            else:
+                device_map[f"model.layers.{i}"] = 1
+        
+        # Embeddings and norm on GPU 0
+        device_map["model.embed_tokens"] = 0
+        device_map["model.norm"] = 0
+        device_map["lm_head"] = 0
+        
+        print(f"   Created device_map with {len(device_map)} entries")
         model = AutoModelForCausalLM.from_pretrained(
             "/data/models/Ornith-1.0-35B-4bit",
-            device_map="auto",
+            device_map=device_map,
             torch_dtype=torch.float16,
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
-        print("   ✓ Model loaded to GPU")
-        
-        # Check memory before FSDP
-        print("\n   Memory BEFORE FSDP:")
-        for i in range(torch.cuda.device_count()):
-            mem = torch.cuda.memory_allocated(i) / 1e9
-            print(f"   GPU {i}: {mem:.2f} GB")
-        
-        # Define auto wrap policy
-        def get_auto_wrap_policy(model):
-            from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
-            return partial(
-                transformer_auto_wrap_policy,
-                transformer_layer_cls={Qwen3_5MoeDecoderLayer},
-            )
-        
-        # Wrap with FSDP
-        print("\n   Wrapping with FSDP...")
-        model = FSDP(
-            model,
-            auto_wrap_policy=get_auto_wrap_policy(model),
-            device_id=torch.cuda.current_device(),
-            mixed_precision=None,
-            sync_module_states=False,
-            use_orig_params=True,
-        )
-        print("   ✓ Model wrapped with FSDP")
-        
-        # Check memory after FSDP
-        print("\n   Memory AFTER FSDP (should be sharded):")
-        for i in range(torch.cuda.device_count()):
-            mem = torch.cuda.memory_allocated(i) / 1e9
-            total = torch.cuda.get_device_properties(i).total_memory / 1e9
-            print(f"   GPU {i}: {mem:.2f} GB / {total:.2f} GB")
-        
-        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
-        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
-        
-        print("\n   Distribution Pattern:")
-        if abs(gpu0_mem - gpu1_mem) < 2.0:  # Within 2GB (more lenient for FSDP)
-            if gpu0_mem < 20.0:  # Less than 20GB each (sharded)
-                print("   ✓ SUCCESSFULLY SHARDED: Model split across GPUs")
-                print(f"      Each GPU has ~{gpu0_mem:.2f}GB (down from ~31GB)")
-                return True
-            else:
-                print("   ⚠ NOT SHARDED: Still duplicate loading")
-                print(f"      Each GPU has ~{gpu0_mem:.2f}GB")
-                return False
-        else:
-            print("   ✗ FAILED: Uneven distribution")
-            return False
+        print("   ✓ Model loaded successfully")
         
+        pattern = check_gpu_memory()
+        print(f"\n   Pattern: {pattern}")
+        return True, pattern
     except Exception as e:
-        print(f"\n   ✗ Test 2 FAILED: {e}")
-        import traceback
-        traceback.print_exc()
-        return False
+        print(f"\n   ✗ FAILED: {e}")
+        return False, str(e)
+
+def test_strategy_4():
+    """Test 4: Load to CPU, then move to GPU manually"""
+    print("\n" + "=" * 80)
+    print("TEST 4: Load to CPU, then move to GPU")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Loading model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            device_map="cpu",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to CPU")
+        
+        # Count params on CPU
+        cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == 'cpu')
+        print(f"   CPU parameters: {cpu_params / 1e9:.2f}B")
+        
+        # Move to GPU 0
+        print("\n   Moving to GPU 0...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 0: {pattern}")
+        
+        # Move to GPU 1
+        print("\n   Moving to GPU 1...")
+        model = model.to("cuda:1")
+        pattern = check_gpu_memory()
+        print(f"   Pattern after move to GPU 1: {pattern}")
+        
+        return True, "LOADED_TO_CPU_THEN_GPU"
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        return False, str(e)
+
+def test_strategy_5():
+    """Test 5: Sequential layer loading (manual distribution)"""
+    print("\n" + "=" * 80)
+    print("TEST 5: Sequential layer loading (manual distribution)")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Loading model layer by layer...")
+        
+        # This is a simplified version - in reality would need more complex logic
+        # For now, just test if we can load to one GPU
+        print("   Loading to GPU 0 only...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            device_map={"": 0},
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to GPU 0")
+        
+        pattern = check_gpu_memory()
+        print(f"\n   Pattern: {pattern}")
+        
+        # Now try GPU 1
+        torch.cuda.empty_cache()
+        print("\n   Loading to GPU 1 only...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            device_map={"": 1},
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to GPU 1")
+        
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        
+        return True, "SEQUENTIAL_LOAD"
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        return False, str(e)
 
 if __name__ == "__main__":
     print("=" * 80)
-    print("Testing model loading strategies")
+    print("Testing multiple model loading strategies")
     print("=" * 80)
     
     # Check GPU availability
@@ -191,29 +226,43 @@ if __name__ == "__main__":
         props = torch.cuda.get_device_properties(i)
         print(f"   GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
     
-    # Test Strategy 1
-    strategy1_ok = test_strategy_1()
+    # Run all tests
+    results = []
     
-    # Test Strategy 2 (FSDP)
-    strategy2_ok = test_strategy_2()
+    tests = [
+        ("Test 1: device_map=auto (no BnB)", test_strategy_1),
+        ("Test 2: device_map=auto + BnB 4-bit", test_strategy_2),
+        ("Test 3: Explicit device_map", test_strategy_3),
+        ("Test 4: Load to CPU then GPU", test_strategy_4),
+        ("Test 5: Sequential layer loading", test_strategy_5),
+    ]
+    
+    for name, test_func in tests:
+        try:
+            success, pattern = test_func()
+            results.append((name, success, pattern))
+        except Exception as e:
+            print(f"\n   ✗ Test crashed: {e}")
+            results.append((name, False, str(e)))
+        
+        # Clear GPU memory between tests
+        torch.cuda.empty_cache()
     
     # Summary
     print("\n" + "=" * 80)
     print("SUMMARY")
     print("=" * 80)
-    if strategy1_ok:
-        print("✓ Strategy 1 (device_map=auto): WORKS - model distributed")
-    else:
-        print("✗ Strategy 1 (device_map=auto): FAILED - model not distributed")
     
-    if strategy2_ok:
-        print("✓ Strategy 2 (FSDP): WORKS - model successfully sharded")
-    else:
-        print("✗ Strategy 2 (FSDP): FAILED - model not sharded")
+    for name, success, pattern in results:
+        status = "✓ PASS" if success else "✗ FAIL"
+        print(f"{status}: {name}")
+        print(f"         Pattern: {pattern}")
     
-    if strategy1_ok or strategy2_ok:
-        print("\n✓ At least one strategy works!")
-        exit(0)
+    # Find working strategies
+    working = [name for name, success, _ in results if success]
+    if working:
+        print(f"\n✓ {len(working)} strategy/strategies work:")
+        for w in working:
+            print(f"  - {w}")
     else:
-        print("\n✗ No strategy works - model cannot be distributed")
-        exit(1)
+        print("\n✗ No strategies work!")