#!/usr/bin/env python3 """ Test multiple model loading strategies to find what works. Each strategy is tested independently. Model: deepreinforce-ai/Ornith-1.0-35B (Qwen3_5Moe architecture) """ import torch from transformers import AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig def get_layer_names(model_path): """Detect decoder layer class names from model config""" print(" Detecting layer names from config...") config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) # Common layer name patterns layer_names = [] # Check for decoder layer if hasattr(config, 'decoder_layer'): layer_names.append(config.decoder_layer) # Check for common patterns if hasattr(config, 'hidden_act'): # Some configs have layer info in different fields pass # If no standard field, try to infer from model type if not layer_names: model_type = config.model_type if 'moe' in model_type.lower(): layer_names.append(f"{model_type.title().replace('_', '')}DecoderLayer") layer_names.append(f"{model_type.title().replace('_', '')}SparseMoeBlock") elif 'qwen' in model_type.lower(): layer_names.append("Qwen2DecoderLayer") else: layer_names.append("DecoderLayer") print(f" Detected layers: {layer_names}") return layer_names 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 quantize_model_bnb(model, quant_type="4bit"): """Quantize model using BnB (BitsAndBytes)""" print(" Using BnB to quantize model...") from transformers import BitsAndBytesConfig from peft import prepare_model_for_kbit_training # Prepare model for k-bit training model = prepare_model_for_kbit_training( model, use_gradient_checkpointing=False, ) # Set quantization config if quant_type == "4bit": bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, ) else: bnb_config = BitsAndBytesConfig( load_in_8bit=True, ) # The actual quantization happens when we set device_map # For now, just return the prepared model print(f" ✓ Model prepared for {quant_type}-bit quantization") return model # 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_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_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 with BnB 4-bit, then move to GPU""" print("\n" + "=" * 80) print("TEST 6: bf16 to CPU → BnB 4-bit (device_map=cpu)") print("=" * 80) try: torch.cuda.empty_cache() print(" Step 1: Load bf16 model to CPU 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", quantization_config=bnb_config, device_map="cpu", # ← Quantize on CPU, not GPU! trust_remote_code=True, low_cpu_mem_usage=True, ) print(f" ✓ Model loaded: {type(model).__name__}") print(f" ✓ Model class: {model.__class__.__name__}") print(f" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") # Check CPU memory import psutil mem = psutil.virtual_memory() print(f" CPU RAM: {mem.used / 1e9:.2f}GB / {mem.total / 1e9:.2f}GB") print("\n Step 2: Move to GPU 0...") model = model.to("cuda:0") pattern = check_gpu_memory() print(f" Pattern after move to GPU 0: {pattern}") print("\n Step 3: Move to GPU 1...") model = model.to("cuda:1") pattern = check_gpu_memory() print(f" Pattern after move to GPU 1: {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_7(): """Test 7: bf16 to CPU with BnB 4-bit → accelerate distribute""" print("\n" + "=" * 80) print("TEST 7: bf16 to CPU → BnB 4-bit → accelerate") print("=" * 80) try: torch.cuda.empty_cache() # Detect layer names dynamically print(" Detecting layer names...") layer_names = get_layer_names("/data/models/Ornith-1.0-35B") print("\n Step 1: Load bf16 model to CPU 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", quantization_config=bnb_config, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") print("\n Step 2: Use accelerate to distribute across GPUs...") from accelerate import infer_auto_device_map # Create device map for quantized model device_map = infer_auto_device_map( model, max_memory={0: "15GB", 1: "15GB"}, no_split_module_classes=layer_names, ) print(f" Created device_map with {len(device_map)} entries") # Reload with device_map model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", quantization_config=bnb_config, device_map=device_map, 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 with BnB 4-bit → 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 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", quantization_config=bnb_config, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") print("\n Step 2: 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 with BnB 8-bit → 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 with BnB 8-bit...") bnb_config = BitsAndBytesConfig( load_in_8bit=True, ) model = AutoModelForCausalLM.from_pretrained( "/data/models/Ornith-1.0-35B", quantization_config=bnb_config, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print(" ✓ Model loaded to CPU with BnB 8-bit (~35GB)") print("\n Step 2: 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 with BnB 4-bit → FSDP""" print("\n" + "=" * 80) print("TEST 10: bf16 to CPU → BnB 4-bit → FSDP") print("=" * 80) try: torch.cuda.empty_cache() print(" Step 1: Load bf16 model to CPU 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", quantization_config=bnb_config, device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True, ) print(" ✓ Model loaded to CPU with BnB 4-bit (~17.5GB)") print("\n Step 2: 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") print("=" * 80) # Check GPU availability print(f"\n1. GPU Check:") print(f" CUDA available: {torch.cuda.is_available()}") print(f" GPU count: {torch.cuda.device_count()}") for i in range(torch.cuda.device_count()): props = torch.cuda.get_device_properties(i) print(f" GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)") # Run all tests 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 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: 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) for name, success, pattern in results: status = "✓ PASS" if success else "✗ FAIL" print(f"{status}: {name}") print(f" Pattern: {pattern}") # 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 strategies work!")