diff --git a/test_model_loading.py b/test_model_loading.py index 68b3bc8..6b48e84 100644 --- a/test_model_loading.py +++ b/test_model_loading.py @@ -4,14 +4,41 @@ Test multiple model loading strategies to find what works. Each strategy is tested independently. Model: deepreinforce-ai/Ornith-1.0-35B (Qwen3_5Moe architecture) -Model class: Qwen3_5MoeForCausalLM -Layer classes: Qwen3_5MoeDecoderLayer, Qwen3_5MoeSparseMoeBlock - -Note: Model does NOT have quantize_4bit() method - need manual quantization """ import torch -from transformers import AutoModelForCausalLM, BitsAndBytesConfig +from transformers import AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig + +def get_layer_names(model_path): + """Detect decoder layer class names from model config""" + print(" Detecting layer names from config...") + config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + + # Common layer name patterns + layer_names = [] + + # Check for decoder layer + if hasattr(config, 'decoder_layer'): + layer_names.append(config.decoder_layer) + + # Check for common patterns + if hasattr(config, 'hidden_act'): + # Some configs have layer info in different fields + pass + + # If no standard field, try to infer from model type + if not layer_names: + model_type = config.model_type + if 'moe' in model_type.lower(): + layer_names.append(f"{model_type.title().replace('_', '')}DecoderLayer") + layer_names.append(f"{model_type.title().replace('_', '')}SparseMoeBlock") + elif 'qwen' in model_type.lower(): + layer_names.append("Qwen2DecoderLayer") + else: + layer_names.append("DecoderLayer") + + print(f" Detected layers: {layer_names}") + return layer_names def check_gpu_memory(): """Check memory usage on all GPUs.""" @@ -308,7 +335,12 @@ def test_strategy_7(): try: torch.cuda.empty_cache() - print(" Step 1: Load bf16 model to CPU with BnB 4-bit...") + + # Detect layer names dynamically + print(" Detecting layer names...") + layer_names = get_layer_names("/data/models/Ornith-1.0-35B") + + print("\n Step 1: Load bf16 model to CPU with BnB 4-bit...") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", @@ -330,7 +362,7 @@ def test_strategy_7(): device_map = infer_auto_device_map( model, max_memory={0: "15GB", 1: "15GB"}, - no_split_module_classes=["Qwen3_5MoeDecoderLayer"], + no_split_module_classes=layer_names, ) print(f" Created device_map with {len(device_map)} entries")