diff --git a/test_model_loading.py b/test_model_loading.py index e853f0a..1898243 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -1,16 +1,39 @@ #!/usr/bin/env python3 """ -Test model loading and GPU distribution without training. -Tests multiple strategies to find what works. +Test multiple model loading strategies to find what works. +Each strategy is tested independently. """ import torch -from transformers import AutoModelForCausalLM +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 test_strategy_1(): - """Test 1: Load with device_map='auto' (no FSDP)""" + """Test 1: device_map='auto' (no quantization config)""" print("\n" + "=" * 80) - print("TEST 1: Load with device_map='auto' (no FSDP)") + print("TEST 1: device_map='auto' (no BnB)") print("=" * 80) try: @@ -24,163 +47,175 @@ def test_strategy_1(): ) print(" ✓ Model loaded successfully") - # Test 1b: Try with load_in_4bit=True (force quantization) - print("\n Testing with load_in_4bit=True (force quantization)...") + 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: device_map='auto' with BnB 4-bit""" + print("\n" + "=" * 80) + print("TEST 2: device_map='auto' + BnB 4-bit") + print("=" * 80) + + try: torch.cuda.empty_cache() - - from transformers import BitsAndBytesConfig + print(" Loading model with BnB 4-bit...") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, ) - - model2 = AutoModelForCausalLM.from_pretrained( + model = 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 + print(" ✓ Model loaded successfully") + pattern = check_gpu_memory() + print(f"\n Pattern: {pattern}") + return True, pattern except Exception as e: - print(f"\n ✗ Test 1 FAILED: {e}") - return False + print(f"\n ✗ FAILED: {e}") + return False, str(e) -def test_strategy_2(): - """Test 2: Load with device_map='auto' then wrap with FSDP""" +def test_strategy_3(): + """Test 3: device_map with explicit GPU assignment""" print("\n" + "=" * 80) - print("TEST 2: Load with device_map='auto' then wrap with FSDP") + print("TEST 3: device_map with explicit GPU assignment") 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 with explicit device_map...") - print(" Loading model...") + # Get model config to determine layers + from transformers import AutoConfig + config = AutoConfig.from_pretrained("/data/models/Ornith-1.0-35B-4bit", 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-4bit", - device_map="auto", + device_map=device_map, 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 + print(" ✓ Model loaded successfully") + pattern = check_gpu_memory() + print(f"\n Pattern: {pattern}") + return True, pattern except Exception as e: - print(f"\n ✗ Test 2 FAILED: {e}") - import traceback - traceback.print_exc() - return False + 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 model to CPU...") + model = AutoModelForCausalLM.from_pretrained( + "/data/models/Ornith-1.0-35B-4bit", + device_map="cpu", + torch_dtype=torch.float16, + 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 to GPU 0 only...") + model = AutoModelForCausalLM.from_pretrained( + "/data/models/Ornith-1.0-35B-4bit", + device_map={"": 0}, + torch_dtype=torch.float16, + 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 to GPU 1 only...") + model = AutoModelForCausalLM.from_pretrained( + "/data/models/Ornith-1.0-35B-4bit", + device_map={"": 1}, + torch_dtype=torch.float16, + 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) if __name__ == "__main__": print("=" * 80) - print("Testing model loading strategies") + print("Testing multiple model loading strategies") print("=" * 80) # Check GPU availability @@ -191,29 +226,43 @@ if __name__ == "__main__": 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() + # Run all tests + results = [] - # Test Strategy 2 (FSDP) - strategy2_ok = test_strategy_2() + 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), + ] + + 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) - 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") + for name, success, pattern in results: + status = "✓ PASS" if success else "✗ FAIL" + print(f"{status}: {name}") + print(f" Pattern: {pattern}") - if strategy1_ok or strategy2_ok: - print("\n✓ At least one strategy works!") - exit(0) + # 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 strategy works - model cannot be distributed") - exit(1) + print("\n✗ No strategies work!")