diff --git a/test_model_loading.py b/test_model_loading.py index 72859e8..d08fbf9 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -1,14 +1,145 @@ #!/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, 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("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) # Check GPU availability @@ -19,76 +150,29 @@ def test_model_loading(): props = torch.cuda.get_device_properties(i) print(f" GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)") - # Test 1: Load with device_map="auto" (distributed) - print("\n2. Test 1: Load with device_map='auto' (should distribute across GPUs)") - 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 + # Test Strategy 1 + strategy1_ok = test_strategy_1() - return True - -if __name__ == "__main__": - success = test_model_loading() - exit(0 if success else 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)