cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: comment out failing tests, add Test 7-10 variations

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 13:45:17 -04:00
parent e2b351e04c
commit 048e68f91a
1 changed files with 383 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

572
test_model_loading.py
View File

@@ -30,194 +30,194 @@ def check_gpu_memory():
else: else:
return f"GPU1_ONLY ({gpu1_mem:.1f}GB)" return f"GPU1_ONLY ({gpu1_mem:.1f}GB)"
def test_strategy_1(): # def test_strategy_1():
"""Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)""" # """Test 1: bf16 model + BnB 4-bit (ON-THE-FLY quantization)"""
print("\n" + "=" * 80) # print("\n" + "=" * 80)
print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)") # print("TEST 1: bf16 model + BnB 4-bit (ON-THE-FLY)")
print("=" * 80) # print("=" * 80)
#
try: # try:
print(" Loading bf16 model with BnB 4-bit...") # print(" Loading bf16 model with BnB 4-bit...")
bnb_config = BitsAndBytesConfig( # bnb_config = BitsAndBytesConfig(
load_in_4bit=True, # load_in_4bit=True,
bnb_4bit_quant_type="nf4", # bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, # bnb_4bit_compute_dtype=torch.bfloat16,
) # )
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # ← bf16 model # "/data/models/Ornith-1.0-35B", # ← bf16 model
quantization_config=bnb_config, # quantization_config=bnb_config,
device_map="auto", # device_map="auto",
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded successfully") # print(" ✓ Model loaded successfully")
#
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f"\n Pattern: {pattern}") # print(f"\n Pattern: {pattern}")
return True, pattern # return True, pattern
except Exception as e: # except Exception as e:
print(f"\n ✗ FAILED: {e}") # print(f"\n ✗ FAILED: {e}")
return False, str(e) # return False, str(e)
def test_strategy_2(): # def test_strategy_2():
"""Test 2: bf16 model + BnB 4-bit (alternative config)""" # """Test 2: bf16 model + BnB 4-bit (alternative config)"""
print("\n" + "=" * 80) # print("\n" + "=" * 80)
print("TEST 2: bf16 model + BnB 4-bit (alt config)") # print("TEST 2: bf16 model + BnB 4-bit (alt config)")
print("=" * 80) # print("=" * 80)
#
try: # try:
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print(" Loading bf16 model with BnB 4-bit...") # print(" Loading bf16 model with BnB 4-bit...")
bnb_config = BitsAndBytesConfig( # bnb_config = BitsAndBytesConfig(
load_in_4bit=True, # load_in_4bit=True,
bnb_4bit_quant_type="nf4", # bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, # bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True, # bnb_4bit_use_double_quant=True,
) # )
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # ← bf16 model # "/data/models/Ornith-1.0-35B", # ← bf16 model
quantization_config=bnb_config, # quantization_config=bnb_config,
device_map="auto", # device_map="auto",
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded successfully") # print(" ✓ Model loaded successfully")
#
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f"\n Pattern: {pattern}") # print(f"\n Pattern: {pattern}")
return True, pattern # return True, pattern
except Exception as e: # except Exception as e:
print(f"\n ✗ FAILED: {e}") # print(f"\n ✗ FAILED: {e}")
return False, str(e) # return False, str(e)
def test_strategy_3(): # def test_strategy_3():
"""Test 3: device_map with explicit GPU assignment""" # """Test 3: device_map with explicit GPU assignment"""
print("\n" + "=" * 80) # print("\n" + "=" * 80)
print("TEST 3: device_map with explicit GPU assignment") # print("TEST 3: device_map with explicit GPU assignment")
print("=" * 80) # print("=" * 80)
#
try: # try:
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print(" Loading model with explicit device_map...") # print(" Loading model with explicit device_map...")
#
# Get model config to determine layers # # Get model config to determine layers
from transformers import AutoConfig # from transformers import AutoConfig
config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True) # config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B", trust_remote_code=True)
num_layers = config.num_hidden_layers # num_layers = config.num_hidden_layers
#
# Split layers: first half on GPU 0, second half on GPU 1 # # Split layers: first half on GPU 0, second half on GPU 1
device_map = {} # device_map = {}
for i in range(num_layers): # for i in range(num_layers):
if i < num_layers // 2: # if i < num_layers // 2:
device_map[f"model.layers.{i}"] = 0 # device_map[f"model.layers.{i}"] = 0
else: # else:
device_map[f"model.layers.{i}"] = 1 # device_map[f"model.layers.{i}"] = 1
#
# Embeddings and norm on GPU 0 # # Embeddings and norm on GPU 0
device_map["model.embed_tokens"] = 0 # device_map["model.embed_tokens"] = 0
device_map["model.norm"] = 0 # device_map["model.norm"] = 0
device_map["lm_head"] = 0 # device_map["lm_head"] = 0
#
print(f" Created device_map with {len(device_map)} entries") # print(f" Created device_map with {len(device_map)} entries")
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # "/data/models/Ornith-1.0-35B",
device_map=device_map, # device_map=device_map,
torch_dtype=torch.bfloat16, # torch_dtype=torch.bfloat16,
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded successfully") # print(" ✓ Model loaded successfully")
#
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f"\n Pattern: {pattern}") # print(f"\n Pattern: {pattern}")
return True, pattern # return True, pattern
except Exception as e: # except Exception as e:
print(f"\n ✗ FAILED: {e}") # print(f"\n ✗ FAILED: {e}")
return False, str(e) # return False, str(e)
def test_strategy_4(): # def test_strategy_4():
"""Test 4: Load to CPU, then move to GPU manually""" # """Test 4: Load to CPU, then move to GPU manually"""
print("\n" + "=" * 80) # print("\n" + "=" * 80)
print("TEST 4: Load to CPU, then move to GPU") # print("TEST 4: Load to CPU, then move to GPU")
print("=" * 80) # print("=" * 80)
#
try: # try:
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print(" Loading bf16 model to CPU...") # print(" Loading bf16 model to CPU...")
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # "/data/models/Ornith-1.0-35B",
device_map="cpu", # device_map="cpu",
torch_dtype=torch.bfloat16, # torch_dtype=torch.bfloat16,
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded to CPU") # print(" ✓ Model loaded to CPU")
#
# Count params on CPU # # Count params on CPU
cpu_params = sum(p.numel() for p in model.parameters() if p.device.type == '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") # print(f" CPU parameters: {cpu_params / 1e9:.2f}B")
#
# Move to GPU 0 # # Move to GPU 0
print("\n Moving to GPU 0...") # print("\n Moving to GPU 0...")
model = model.to("cuda:0") # model = model.to("cuda:0")
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f" Pattern after move to GPU 0: {pattern}") # print(f" Pattern after move to GPU 0: {pattern}")
#
# Move to GPU 1 # # Move to GPU 1
print("\n Moving to GPU 1...") # print("\n Moving to GPU 1...")
model = model.to("cuda:1") # model = model.to("cuda:1")
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f" Pattern after move to GPU 1: {pattern}") # print(f" Pattern after move to GPU 1: {pattern}")
#
return True, "LOADED_TO_CPU_THEN_GPU" # return True, "LOADED_TO_CPU_THEN_GPU"
except Exception as e: # except Exception as e:
print(f"\n ✗ FAILED: {e}") # print(f"\n ✗ FAILED: {e}")
return False, str(e) # return False, str(e)
def test_strategy_5(): # def test_strategy_5():
"""Test 5: Sequential layer loading (manual distribution)""" # """Test 5: Sequential layer loading (manual distribution)"""
print("\n" + "=" * 80) # print("\n" + "=" * 80)
print("TEST 5: Sequential layer loading (manual distribution)") # print("TEST 5: Sequential layer loading (manual distribution)")
print("=" * 80) # print("=" * 80)
#
try: # try:
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print(" Loading model layer by layer...") # print(" Loading model layer by layer...")
#
# This is a simplified version - in reality would need more complex logic # # This is a simplified version - in reality would need more complex logic
# For now, just test if we can load to one GPU # # For now, just test if we can load to one GPU
print(" Loading bf16 to GPU 0 only...") # print(" Loading bf16 to GPU 0 only...")
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # "/data/models/Ornith-1.0-35B",
device_map={"": 0}, # device_map={"": 0},
torch_dtype=torch.bfloat16, # torch_dtype=torch.bfloat16,
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded to GPU 0") # print(" ✓ Model loaded to GPU 0")
#
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f"\n Pattern: {pattern}") # print(f"\n Pattern: {pattern}")
#
# Now try GPU 1 # # Now try GPU 1
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print("\n Loading bf16 to GPU 1 only...") # print("\n Loading bf16 to GPU 1 only...")
model = AutoModelForCausalLM.from_pretrained( # model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B", # "/data/models/Ornith-1.0-35B",
device_map={"": 1}, # device_map={"": 1},
torch_dtype=torch.bfloat16, # torch_dtype=torch.bfloat16,
trust_remote_code=True, # trust_remote_code=True,
low_cpu_mem_usage=True, # low_cpu_mem_usage=True,
) # )
print(" ✓ Model loaded to GPU 1") # print(" ✓ Model loaded to GPU 1")
#
pattern = check_gpu_memory() # pattern = check_gpu_memory()
print(f" Pattern: {pattern}") # print(f" Pattern: {pattern}")
#
return True, "SEQUENTIAL_LOAD" # return True, "SEQUENTIAL_LOAD"
except Exception as e: # except Exception as e:
print(f"\n ✗ FAILED: {e}") # print(f"\n ✗ FAILED: {e}")
return False, str(e) # return False, str(e)
def test_strategy_6(): def test_strategy_6():
"""Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU""" """Test 6: Load bf16 to CPU, quantize with BnB, then move to GPU"""
@@ -273,6 +273,196 @@ def test_strategy_6():
traceback.print_exc() traceback.print_exc()
return False, str(e) 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__": if __name__ == "__main__":
print("=" * 80) print("=" * 80)
print("Testing multiple model loading strategies") print("Testing multiple model loading strategies")
@@ -290,12 +480,16 @@ if __name__ == "__main__":
results = [] results = []
tests = [ tests = [
("Test 1: device_map=auto (no BnB)", test_strategy_1), # ("Test 1: device_map=auto (no BnB)", test_strategy_1),
("Test 2: device_map=auto + BnB 4-bit", test_strategy_2), # ("Test 2: device_map=auto + BnB 4-bit", test_strategy_2),
("Test 3: Explicit device_map", test_strategy_3), # ("Test 3: Explicit device_map", test_strategy_3),
("Test 4: Load to CPU then GPU", test_strategy_4), # ("Test 4: Load to CPU then GPU", test_strategy_4),
("Test 5: Sequential layer loading", test_strategy_5), # ("Test 5: Sequential layer loading", test_strategy_5),
("Test 6: bf16 to CPU → BnB 4-bit → GPU", test_strategy_6), ("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: for name, test_func in tests: