330 lines
11 KiB
Python
330 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test multiple model loading strategies to find what works.
|
|
Each strategy is tested independently.
|
|
"""
|
|
|
|
import torch
|
|
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: 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, quantize with BnB, then move to GPU"""
|
|
print("\n" + "=" * 80)
|
|
print("TEST 6: bf16 to CPU → BnB 4-bit quantize → 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)")
|
|
|
|
# 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: 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")
|
|
|
|
# Actually quantize the model
|
|
from bitsandbytes.nn.modules import Params4bit
|
|
print(" Quantizing weights to 4-bit...")
|
|
model.quantize_4bit()
|
|
print(" ✓ Model quantized to 4-bit (~17.5GB)")
|
|
|
|
print("\n Step 3: 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 4: 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)
|
|
|
|
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),
|
|
]
|
|
|
|
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!")
|