feat: dynamically detect layer names from model config

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")