220 lines
7.8 KiB
Python
220 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test model loading and GPU distribution without training.
|
|
Tests multiple strategies to find what works.
|
|
"""
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
def test_strategy_1():
|
|
"""Test 1: Load with device_map='auto' (no FSDP)"""
|
|
print("\n" + "=" * 80)
|
|
print("TEST 1: Load with device_map='auto' (no FSDP)")
|
|
print("=" * 80)
|
|
|
|
try:
|
|
print(" Loading model...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"/data/models/Ornith-1.0-35B-4bit",
|
|
device_map="auto",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print(" ✓ Model loaded successfully")
|
|
|
|
# Test 1b: Try with load_in_4bit=True (force quantization)
|
|
print("\n Testing with load_in_4bit=True (force quantization)...")
|
|
torch.cuda.empty_cache()
|
|
|
|
from transformers import BitsAndBytesConfig
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.bfloat16,
|
|
)
|
|
|
|
model2 = AutoModelForCausalLM.from_pretrained(
|
|
"/data/models/Ornith-1.0-35B-4bit",
|
|
quantization_config=bnb_config,
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print(" ✓ Model loaded with BnB 4-bit")
|
|
|
|
print("\n Memory with BnB 4-bit:")
|
|
for i in range(torch.cuda.device_count()):
|
|
mem = torch.cuda.memory_allocated(i) / 1e9
|
|
print(f" GPU {i}: {mem:.2f} GB")
|
|
|
|
gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
|
|
gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
|
|
|
|
print("\n Distribution Pattern:")
|
|
if abs(gpu0_mem - gpu1_mem) < 1.0:
|
|
if gpu0_mem < 10.0:
|
|
print(" ✓ DISTRIBUTED: Model split across both GPUs")
|
|
return True
|
|
else:
|
|
print(" ⚠ DUPLICATE: Same model on both GPUs")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB")
|
|
return False
|
|
else:
|
|
print(" ✗ NOT DISTRIBUTED: Model on one GPU only")
|
|
return False
|
|
|
|
# Check memory usage
|
|
print("\n Memory Usage:")
|
|
for i in range(torch.cuda.device_count()):
|
|
mem_allocated = torch.cuda.memory_allocated(i) / 1e9
|
|
total = torch.cuda.get_device_properties(i).total_memory / 1e9
|
|
print(f" GPU {i}: {mem_allocated:.2f} GB / {total:.2f} GB")
|
|
|
|
gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
|
|
gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
|
|
|
|
print("\n Distribution Pattern:")
|
|
if abs(gpu0_mem - gpu1_mem) < 1.0:
|
|
if gpu0_mem < 10.0:
|
|
print(" ✓ DISTRIBUTED: Model split across both GPUs")
|
|
return True
|
|
else:
|
|
print(" ⚠ DUPLICATE: Same model on both GPUs")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB")
|
|
return False
|
|
else:
|
|
print(" ✗ NOT DISTRIBUTED: Model on one GPU only")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"\n ✗ Test 1 FAILED: {e}")
|
|
return False
|
|
|
|
def test_strategy_2():
|
|
"""Test 2: Load with device_map='auto' then wrap with FSDP"""
|
|
print("\n" + "=" * 80)
|
|
print("TEST 2: Load with device_map='auto' then wrap with FSDP")
|
|
print("=" * 80)
|
|
|
|
try:
|
|
import torch.distributed as dist
|
|
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
|
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
|
from functools import partial
|
|
|
|
# Initialize distributed process group
|
|
if not dist.is_initialized():
|
|
dist.init_process_group(backend="nccl")
|
|
print(" ✓ Distributed process group initialized")
|
|
|
|
# Clear GPU memory
|
|
torch.cuda.empty_cache()
|
|
|
|
print(" Loading model...")
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"/data/models/Ornith-1.0-35B-4bit",
|
|
device_map="auto",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
print(" ✓ Model loaded to GPU")
|
|
|
|
# Check memory before FSDP
|
|
print("\n Memory BEFORE FSDP:")
|
|
for i in range(torch.cuda.device_count()):
|
|
mem = torch.cuda.memory_allocated(i) / 1e9
|
|
print(f" GPU {i}: {mem:.2f} GB")
|
|
|
|
# Define auto wrap policy
|
|
def get_auto_wrap_policy(model):
|
|
from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer
|
|
return partial(
|
|
transformer_auto_wrap_policy,
|
|
transformer_layer_cls={Qwen3_5MoeDecoderLayer},
|
|
)
|
|
|
|
# Wrap with FSDP
|
|
print("\n Wrapping with FSDP...")
|
|
model = FSDP(
|
|
model,
|
|
auto_wrap_policy=get_auto_wrap_policy(model),
|
|
device_id=torch.cuda.current_device(),
|
|
mixed_precision=None,
|
|
sync_module_states=False,
|
|
use_orig_params=True,
|
|
)
|
|
print(" ✓ Model wrapped with FSDP")
|
|
|
|
# Check memory after FSDP
|
|
print("\n Memory AFTER FSDP (should be sharded):")
|
|
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
|
|
|
|
print("\n Distribution Pattern:")
|
|
if abs(gpu0_mem - gpu1_mem) < 2.0: # Within 2GB (more lenient for FSDP)
|
|
if gpu0_mem < 20.0: # Less than 20GB each (sharded)
|
|
print(" ✓ SUCCESSFULLY SHARDED: Model split across GPUs")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB (down from ~31GB)")
|
|
return True
|
|
else:
|
|
print(" ⚠ NOT SHARDED: Still duplicate loading")
|
|
print(f" Each GPU has ~{gpu0_mem:.2f}GB")
|
|
return False
|
|
else:
|
|
print(" ✗ FAILED: Uneven distribution")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"\n ✗ Test 2 FAILED: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 80)
|
|
print("Testing 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)")
|
|
|
|
# Test Strategy 1
|
|
strategy1_ok = test_strategy_1()
|
|
|
|
# Test Strategy 2 (FSDP)
|
|
strategy2_ok = test_strategy_2()
|
|
|
|
# Summary
|
|
print("\n" + "=" * 80)
|
|
print("SUMMARY")
|
|
print("=" * 80)
|
|
if strategy1_ok:
|
|
print("✓ Strategy 1 (device_map=auto): WORKS - model distributed")
|
|
else:
|
|
print("✗ Strategy 1 (device_map=auto): FAILED - model not distributed")
|
|
|
|
if strategy2_ok:
|
|
print("✓ Strategy 2 (FSDP): WORKS - model successfully sharded")
|
|
else:
|
|
print("✗ Strategy 2 (FSDP): FAILED - model not sharded")
|
|
|
|
if strategy1_ok or strategy2_ok:
|
|
print("\n✓ At least one strategy works!")
|
|
exit(0)
|
|
else:
|
|
print("\n✗ No strategy works - model cannot be distributed")
|
|
exit(1)
|