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

feat: test both strategies - device_map=auto and FSDP

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 12:19:05 -04:00
parent 7f498134f2
commit c5a3b87eed
1 changed files with 159 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

234
test_model_loading.py
View File

@@ -1,14 +1,145 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
""" """
Test model loading and GPU distribution without training. Test model loading and GPU distribution without training.
Tests multiple strategies to find what works.
""" """
import torch import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig from transformers import AutoModelForCausalLM
def test_model_loading(): 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) print("=" * 80)
print("Testing model loading and GPU distribution")
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")
# 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) print("=" * 80)
# Check GPU availability # Check GPU availability
@@ -19,76 +150,29 @@ def test_model_loading():
props = torch.cuda.get_device_properties(i) props = torch.cuda.get_device_properties(i)
print(f" GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)") print(f" GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
# Test 1: Load with device_map="auto" (distributed) # Test Strategy 1
print("\n2. Test 1: Load with device_map='auto' (should distribute across GPUs)") strategy1_ok = test_strategy_1()
try:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
print(" Loading model...")
model = AutoModelForCausalLM.from_pretrained(
"/data/models/Ornith-1.0-35B-4bit", # Use already-quantized 4-bit model
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=True,
low_cpu_mem_usage=True,
)
print(" ✓ Model loaded successfully")
# Check memory usage
print("\n3. Memory Usage:")
for i in range(torch.cuda.device_count()):
mem_allocated = torch.cuda.memory_allocated(i) / 1e9
mem_reserved = torch.cuda.memory_reserved(i) / 1e9
total = torch.cuda.get_device_properties(i).total_memory / 1e9
print(f" GPU {i}:")
print(f" Allocated: {mem_allocated:.2f} GB")
print(f" Reserved: {mem_reserved:.2f} GB")
print(f" Total: {total:.2f} GB")
print(f" Free: {total - mem_allocated:.2f} GB")
# Check if model is distributed
print("\n4. Distribution Check:")
# Get memory on each GPU
gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
print(f" GPU 0 memory: {gpu0_mem:.2f} GB")
print(f" GPU 1 memory: {gpu1_mem:.2f} GB")
# Determine distribution pattern
print("\n5. Distribution Pattern:")
if abs(gpu0_mem - gpu1_mem) < 1.0: # Within 1GB
if gpu0_mem < 10.0: # Less than 10GB each
print(" ✓ DISTRIBUTED: Model split across both GPUs")
print(f" Each GPU has ~{gpu0_mem:.2f}GB of the model")
else:
print(" ⚠ DUPLICATE: Same model loaded on both GPUs")
print(f" Each GPU has ~{gpu0_mem:.2f}GB (wasteful but fits)")
else:
print(" ✗ NOT DISTRIBUTED: Model on one GPU only")
if gpu0_mem > gpu1_mem:
print(f" GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
else:
print(f" GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
print("\n" + "=" * 80)
print("TEST PASSED: Model loaded and distributed across GPUs")
print("=" * 80)
except Exception as e:
print(f"\n✗ Test 1 FAILED: {e}")
print("\nThis means the model is NOT being distributed properly!")
print("It might be trying to fit the entire model on one GPU.")
return False
return True # Test Strategy 2 (FSDP)
strategy2_ok = test_strategy_2()
if __name__ == "__main__":
success = test_model_loading() # Summary
exit(0 if success else 1) 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)