feat: test both strategies - device_map=auto and FSDP

test_model_loading.py

@@ -1,14 +1,145 @@
 #!/usr/bin/env python3
 """
 Test model loading and GPU distribution without training.
+Tests multiple strategies to find what works.
 """
 
 import torch
-from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM
 
-def test_model_loading():
+def test_strategy_1():
+    """Test 1: Load with device_map='auto' (no FSDP)"""
+    print("\n" + "=" * 80)
+    print("TEST 1: Load with device_map='auto' (no FSDP)")
     print("=" * 80)
-    print("Testing model loading and GPU distribution")
+    
+    try:
+        print("   Loading model...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            device_map="auto",
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded successfully")
+        
+        # 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
+        
+    except Exception as e:
+        print(f"\n   ✗ Test 1 FAILED: {e}")
+        return False
+
+def test_strategy_2():
+    """Test 2: Load with device_map='auto' then wrap with FSDP"""
+    print("\n" + "=" * 80)
+    print("TEST 2: Load with device_map='auto' then wrap with FSDP")
+    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...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B-4bit",
+            device_map="auto",
+            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
+        
+    except Exception as e:
+        print(f"\n   ✗ Test 2 FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+if __name__ == "__main__":
+    print("=" * 80)
+    print("Testing model loading strategies")
     print("=" * 80)
     
     # Check GPU availability
@@ -19,76 +150,29 @@ def test_model_loading():
         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,
-        )
+    # Test Strategy 1
+    strategy1_ok = test_strategy_1()
     
-        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")
+    # Test Strategy 2 (FSDP)
+    strategy2_ok = test_strategy_2()
     
+    # Summary
     print("\n" + "=" * 80)
-        print("TEST PASSED: Model loaded and distributed across GPUs")
+    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")
     
-    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
+    if strategy2_ok:
+        print("✓ Strategy 2 (FSDP): WORKS - model successfully sharded")
+    else:
+        print("✗ Strategy 2 (FSDP): FAILED - model not sharded")
     
-    return True
-
-if __name__ == "__main__":
-    success = test_model_loading()
-    exit(0 if success else 1)
+    if strategy1_ok or strategy2_ok:
+        print("\n✓ At least one strategy works!")
+        exit(0)
+    else:
+        print("\n✗ No strategy works - model cannot be distributed")
+        exit(1)