561 lines
19 KiB
Python
561 lines
19 KiB
Python
#!/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)
|
|
Model class: Qwen3_5MoeForCausalLM
|
|
Layer classes: Qwen3_5MoeDecoderLayer, Qwen3_5MoeSparseMoeBlock
|
|
|
|
Note: Model does NOT have quantize_4bit() method - need manual quantization
|
|
"""
|
|
|
|
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 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, 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 using BnB
|
|
print(" Quantizing weights to 4-bit using BnB...")
|
|
from bitsandbytes import quantize_batch
|
|
# Note: This is a simplified approach - actual implementation may vary
|
|
print(" ⚠ Manual BnB quantization may need different approach")
|
|
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)
|
|
|
|
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")
|
|
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!")
|