From 048e68f91ab524c58508f7f4ebd282c73237fd7d Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Thu, 2 Jul 2026 13:45:17 -0400
Subject: [PATCH] refactor: comment out failing tests, add Test 7-10 variations

---
 test_model_loading.py | 572 ++++++++++++++++++++++++++++--------------
 1 file changed, 383 insertions(+), 189 deletions(-)

diff --git a/test_model_loading.py b/test_model_loading.py
index 328f7dc..0c19aff 100644
--- a/test_model_loading.py
+++ b/test_model_loading.py
@@ -30,194 +30,194 @@ def check_gpu_memory():
         else:
             return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
 
-def test_strategy_1():
-    """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
-    print("\n" + "=" * 80)
-    print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
-    print("=" * 80)
-    
-    try:
-        print("   Loading bf16 model with BnB 4-bit...")
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-        )
-        model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B",  # ← bf16 model
-            quantization_config=bnb_config,
-            device_map="auto",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("   ✓ Model loaded successfully")
-        
-        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)
+# def test_strategy_1():
+#     """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
+#     print("\n" + "=" * 80)
+#     print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
+#     print("=" * 80)
+#     
+#     try:
+#         print("   Loading bf16 model with BnB 4-bit...")
+#         bnb_config = BitsAndBytesConfig(
+#             load_in_4bit=True,
+#             bnb_4bit_quant_type="nf4",
+#             bnb_4bit_compute_dtype=torch.bfloat16,
+#         )
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",  # ← bf16 model
+#             quantization_config=bnb_config,
+#             device_map="auto",
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         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)
 
-def test_strategy_2():
-    """Test 2: bf16 model + BnB 4-bit (alternative config)"""
-    print("\n" + "=" * 80)
-    print("TEST 2: bf16 model + BnB 4-bit (alt config)")
-    print("=" * 80)
-    
-    try:
-        torch.cuda.empty_cache()
-        print("   Loading bf16 model with BnB 4-bit...")
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True,
-        )
-        model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B",  # ← bf16 model
-            quantization_config=bnb_config,
-            device_map="auto",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("   ✓ Model loaded successfully")
-        
-        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)
+# def test_strategy_2():
+#     """Test 2: bf16 model + BnB 4-bit (alternative config)"""
+#     print("\n" + "=" * 80)
+#     print("TEST 2: bf16 model + BnB 4-bit (alt config)")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading bf16 model with BnB 4-bit...")
+#         bnb_config = BitsAndBytesConfig(
+#             load_in_4bit=True,
+#             bnb_4bit_quant_type="nf4",
+#             bnb_4bit_compute_dtype=torch.bfloat16,
+#             bnb_4bit_use_double_quant=True,
+#         )
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",  # ← bf16 model
+#             quantization_config=bnb_config,
+#             device_map="auto",
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         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)
 
-def test_strategy_3():
-    """Test 3: device_map with explicit GPU assignment"""
-    print("\n" + "=" * 80)
-    print("TEST 3: device_map with explicit GPU assignment")
-    print("=" * 80)
-    
-    try:
-        torch.cuda.empty_cache()
-        print("   Loading model with explicit device_map...")
-        
-        # Get model config to determine layers
-        from transformers import AutoConfig
-        config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", 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",
-            device_map=device_map,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
-        )
-        print("   ✓ Model loaded successfully")
-        
-        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)
+# def test_strategy_3():
+#     """Test 3: device_map with explicit GPU assignment"""
+#     print("\n" + "=" * 80)
+#     print("TEST 3: device_map with explicit GPU assignment")
+#     print("=" * 80)
+#     
+#     try:
+#         torch.cuda.empty_cache()
+#         print("   Loading model with explicit device_map...")
+#         
+#         # Get model config to determine layers
+#         from transformers import AutoConfig
+#         config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", 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",
+#             device_map=device_map,
+#             torch_dtype=torch.bfloat16,
+#             trust_remote_code=True,
+#             low_cpu_mem_usage=True,
+#         )
+#         print("   ✓ Model loaded successfully")
+#         
+#         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)
 
-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 bf16 model to CPU...")
-        model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B",
-            device_map="cpu",
-            torch_dtype=torch.bfloat16,
-            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_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 bf16 model to CPU...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map="cpu",
+#             torch_dtype=torch.bfloat16,
+#             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 bf16 to GPU 0 only...")
-        model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B",
-            device_map={"": 0},
-            torch_dtype=torch.bfloat16,
-            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 bf16 to GPU 1 only...")
-        model = AutoModelForCausalLM.from_pretrained(
-            "/data/models/Ornith-1.0-35B",
-            device_map={"": 1},
-            torch_dtype=torch.bfloat16,
-            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)
+# 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 bf16 to GPU 0 only...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map={"": 0},
+#             torch_dtype=torch.bfloat16,
+#             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 bf16 to GPU 1 only...")
+#         model = AutoModelForCausalLM.from_pretrained(
+#             "/data/models/Ornith-1.0-35B",
+#             device_map={"": 1},
+#             torch_dtype=torch.bfloat16,
+#             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)
 
 def test_strategy_6():
     """Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU"""
@@ -273,6 +273,196 @@ def test_strategy_6():
         traceback.print_exc()
         return False, str(e)
 
+def test_strategy_7():
+    """Test 7: bf16 to CPU → BnB 4-bit → Use accelerate to distribute"""
+    print("\n" + "=" * 80)
+    print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate device_map")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to CPU (~70GB)")
+        
+        print("\n   Step 2: Apply BnB 4-bit quantization...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model prepared for k-bit training")
+        
+        print("   Step 3: Quantize weights to 4-bit...")
+        model.quantize_4bit()
+        print("   ✓ Model quantized to 4-bit (~17.5GB)")
+        
+        print("\n   Step 4: Use accelerate to distribute across GPUs...")
+        from accelerate import infer_auto_device_map, init_empty_weights
+        from accelerate.utils import get_balanced_memory
+        
+        # Create device map
+        device_map = infer_auto_device_map(
+            model,
+            max_memory={0: "15GB", 1: "15GB"},
+            no_split_module_classes=["Qwen3_5MoeDecoderLayer"],
+        )
+        print(f"   Created device_map with {len(device_map)} entries")
+        
+        # Load model with device_map
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map=device_map,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded with device_map")
+        
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_8():
+    """Test 8: bf16 to CPU → BnB 4-bit → Load to GPU 0 only"""
+    print("\n" + "=" * 80)
+    print("TEST 8: bf16 to CPU → BnB 4-bit → GPU 0 only")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to CPU (~70GB)")
+        
+        print("\n   Step 2: Apply BnB 4-bit quantization...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model prepared for k-bit training")
+        
+        print("   Step 3: Quantize weights to 4-bit...")
+        model.quantize_4bit()
+        print("   ✓ Model quantized to 4-bit (~17.5GB)")
+        
+        print("\n   Step 4: Move to GPU 0 only...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_9():
+    """Test 9: bf16 to CPU → BnB 4-bit (int8) → GPU"""
+    print("\n" + "=" * 80)
+    print("TEST 9: bf16 to CPU → BnB 8-bit → GPU")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to CPU (~70GB)")
+        
+        print("\n   Step 2: Apply BnB 8-bit quantization...")
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(
+            model,
+            use_gradient_checkpointing=False,
+        )
+        print("   ✓ Model prepared for k-bit training")
+        
+        print("   Step 3: Quantize weights to 8-bit...")
+        # Use int8 quantization instead of 4-bit
+        from bitsandbytes.nn.modules import Params8bit
+        # Note: This is a simplified version - actual int8 quantization may need different approach
+        print("   ⚠ int8 quantization may not be fully implemented")
+        
+        print("\n   Step 4: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
+def test_strategy_10():
+    """Test 10: bf16 to CPU → FSDP → GPU (single process)"""
+    print("\n" + "=" * 80)
+    print("TEST 10: bf16 to CPU → FSDP (single process)")
+    print("=" * 80)
+    
+    try:
+        torch.cuda.empty_cache()
+        print("   Step 1: Load bf16 model to CPU...")
+        model = AutoModelForCausalLM.from_pretrained(
+            "/data/models/Ornith-1.0-35B",
+            device_map="cpu",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+        )
+        print("   ✓ Model loaded to CPU (~70GB)")
+        
+        print("\n   Step 2: Apply FSDP wrapping...")
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        from torch.distributed.fsdp import MixedPrecision
+        
+        # Wrap model with FSDP
+        model = FSDP(
+            model,
+            mixed_precision=MixedPrecision(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                buffer_dtype=torch.float32,
+            ),
+            auto_wrap_policy=None,
+        )
+        print("   ✓ Model wrapped with FSDP")
+        
+        print("\n   Step 3: Move to GPU...")
+        model = model.to("cuda:0")
+        pattern = check_gpu_memory()
+        print(f"   Pattern: {pattern}")
+        return True, pattern
+    except Exception as e:
+        print(f"\n   ✗ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        return False, str(e)
+
 if __name__ == "__main__":
     print("=" * 80)
     print("Testing multiple model loading strategies")
@@ -290,12 +480,16 @@ if __name__ == "__main__":
     results = []
     
     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),
+        # ("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),
         ("Test 6: bf16 to CPU → BnB 4-bit → GPU", test_strategy_6),
+        ("Test 7: bf16 to CPU → BnB 4-bit → accelerate", test_strategy_7),
+        ("Test 8: bf16 to CPU → BnB 4-bit → GPU 0 only", test_strategy_8),
+        ("Test 9: bf16 to CPU → BnB 8-bit → GPU", test_strategy_9),
+        ("Test 10: bf16 to CPU → FSDP → GPU", test_strategy_10),
     ]
     
     for name, test_func in tests: