#!/usr/bin/env python3 """ Test script to verify DotsVLMForCausalLM model loading """ import os import json import shutil import torch from transformers import AutoConfig, AutoProcessor from .configuration_dots_vlm import DotsVLMConfig, DotsVLMProcessor from .modeling_dots_vlm import DotsVLMForCausalLM import traceback def create_debug_model_path(original_path, debug_layers=3): """ 创建一个debug版本的模型路径，只包含指定层数的配置和权重 Args: original_path: 原始模型路径 debug_layers: debug模式下使用的层数 Returns: debug模型的临时路径 """ debug_path = original_path + "_debug_temp" # 如果debug目录已存在，先删除 # if os.path.exists(debug_path): # shutil.rmtree(debug_path) # 创建debug目录 os.makedirs(debug_path, exist_ok=True) # 复制基本配置文件 files_to_copy = [ 'tokenizer_config.json', 'tokenizer.json', 'special_tokens_map.json', 'preprocessor_config.json', 'modeling_deepseek.py', 'configuration_deepseek.py' ] for file in files_to_copy: src = os.path.join(original_path, file) dst = os.path.join(debug_path, file) if os.path.exists(src): if not os.path.exists(dst): print(f"📁 Copying basefile {file}... from {src} to {dst}") shutil.copy2(src, dst) # # 复制ve目录（vision encoder） # ve_src = os.path.join(original_path, 've') # ve_dst = os.path.join(debug_path, 've') # if os.path.exists(ve_src): # shutil.copytree(ve_src, ve_dst) # 修改config.json - 减少层数 config_src = os.path.join(original_path, 'config.json') config_dst = os.path.join(debug_path, 'config.json') with open(config_src, 'r') as f: config = json.load(f) # 修改为debug层数 original_layers = config['num_hidden_layers'] config['num_hidden_layers'] = debug_layers print(f"🔧 DEBUG: Reducing num_hidden_layers from {original_layers} to {debug_layers}") with open(config_dst, 'w') as f: json.dump(config, f, indent=2) # 读取原始safetensor索引 index_src = os.path.join(original_path, 'model.safetensors.index.json') with open(index_src, 'r') as f: index_data = json.load(f) original_weight_map = index_data['weight_map'] # 确定需要的safetensor文件 # 前8个文件：包含嵌入层和前几层 front_files = [f"model-{i:05d}-of-00316.safetensors" for i in range(1, 9)] # 第2层的权重在model-00056中，需要特别包含 layer2_files = ["model-00056-of-00316.safetensors"] # 后2个文件：包含最后的输出层和vision_tower back_files = ["model-00315-of-00316.safetensors", "model-00316-of-00316.safetensors"] # model.norm所在文件 norm_files = ["model-00314-of-00316.safetensors"] needed_files = set(front_files + layer2_files + back_files + norm_files) print(f"🔧 DEBUG: Will load {len(needed_files)} safetensor files instead of 316:") for f in sorted(needed_files): print(f" - {f}") # 过滤权重映射，只保留需要的层和基础组件 new_weight_map = {} # 保留嵌入层和输出层 for key, file in original_weight_map.items(): # 保留基础组件 if any(key.startswith(prefix) for prefix in [ 'model.embed_tokens', 'model.norm', 'lm_head' ]): if file in needed_files: new_weight_map[key] = file # 保留vision_tower权重（这些在model-00315和model-00316中） elif key.startswith('vision_tower.'): if file in needed_files: new_weight_map[key] = file # 只保留前debug_layers层 elif key.startswith('model.layers.'): # 提取层号 layer_parts = key.split('.') if len(layer_parts) >= 3 and layer_parts[2].isdigit(): layer_num = int(layer_parts[2]) if layer_num < debug_layers and file in needed_files: new_weight_map[key] = file print(f"🔧 DEBUG: Filtered weight map from {len(original_weight_map)} to {len(new_weight_map)} entries") # 复制需要的safetensor文件 copied_files = set() for file in new_weight_map.values(): if file not in copied_files: src_file = os.path.join(original_path, file) dst_file = os.path.join(debug_path, file) if os.path.exists(src_file): if not os.path.exists(dst_file): print(f"📁 Copying Safetensor {file}... from {src_file} to {dst_file}") shutil.copy2(src_file, dst_file) copied_files.add(file) else: print(f"⚠️ File not found: {src_file}") # 创建新的索引文件 new_index_data = { "metadata": index_data.get("metadata", {}), "weight_map": new_weight_map } index_dst = os.path.join(debug_path, 'model.safetensors.index.json') with open(index_dst, 'w') as f: json.dump(new_index_data, f, indent=2) print(f"✅ DEBUG: Created debug model at {debug_path}") return debug_path def test_model_loading(): """Test loading the model from pretrained weights""" # Path to your model weights model_path = "." # Check if DEBUG mode is enabled debug_minimal_layers = os.getenv('DEBUG_MINIMAL_LAYERS', '0') == '1' if debug_minimal_layers: print("🔧 DEBUG MINIMAL LAYERS: Using only 3 layers and minimal safetensors") # 创建debug版本的模型 model_path = create_debug_model_path(model_path, debug_layers=3) print("Loading model configuration...") try: config = AutoConfig.from_pretrained(model_path) print(f"✓ Config loaded successfully: {config.__class__.__name__}") print(f" Model type: {config.model_type}") print(f" Architecture: {config.architectures}") print(f" Number of hidden layers: {config.num_hidden_layers}") # Check if quantization config exists and is fp8 if hasattr(config, 'quantization_config') and config.quantization_config is not None: quant_config = config.quantization_config if isinstance(quant_config, dict) and quant_config.get('quant_method') == 'fp8': print(" Detected FP8 quantization configuration") print(f" Format: {quant_config.get('fmt', 'unknown')}") print(f" Weight block size: {quant_config.get('weight_block_size', 'unknown')}") print(f" Activation scheme: {quant_config.get('activation_scheme', 'unknown')}") except Exception as e: print(f"✗ Error loading config: {e}") return False print("\nTesting processor loading...") try: # Test if processor can be loaded try: processor = AutoProcessor.from_pretrained(model_path) print(f"✓ AutoProcessor loaded successfully: {processor.__class__.__name__}") except Exception as e: print(f"⚠️ AutoProcessor failed, trying direct import: {e}") # Fallback to direct processor creation processor = DotsVLMProcessor.from_pretrained(model_path) print(f"✓ Direct processor loaded successfully: {processor.__class__.__name__}") # Check processor attributes if hasattr(processor, 'image_token'): print(f" Image token: {processor.image_token}") if hasattr(processor, 'tokenizer'): print(f" Tokenizer type: {processor.tokenizer.__class__.__name__}") if hasattr(processor, 'image_processor'): print(f" Image processor type: {processor.image_processor.__class__.__name__}") except Exception as e: print(f"⚠️ Processor loading failed (this is OK for now): {e}") processor = None print("\nLoading model with auto-fixing vision_tower quantization...") try: # Use the custom model class directly for elegant handling model = DotsVLMForCausalLM.from_pretrained( model_path, config=config, torch_dtype="auto", device_map="auto", trust_remote_code=True ) print(f"✓ Model loaded successfully: {model.__class__.__name__}") print(f" Number of parameters: {model.num_parameters():,}") # Check vision_tower dtype if hasattr(model, 'vision_tower') and model.vision_tower is not None: vision_sample_param = list(model.vision_tower.parameters())[0] print(f" Vision tower dtype: {vision_sample_param.dtype}") except Exception as e: traceback.print_exc() print(f"✗ Error loading model: {e}") return False # print("\nLoading vision weights...") # try: # vision_weights_path = f"{model_path}/ve/visual.pt" # vision_state_dict = torch.load(vision_weights_path, map_location='cpu', weights_only=False) # # Extract vision_encoder weights and remove vision_tower prefix # vision_encoder_weights = vision_state_dict['vision_encoder'] # cleaned_weights = {} # for key, value in vision_encoder_weights.items(): # if key.startswith('vision_tower.'): # new_key = key[len('vision_tower.'):] # Remove 'vision_tower.' prefix # cleaned_weights[new_key] = value # else: # cleaned_weights[key] = value # # Convert weights to match the vision_tower dtype # target_dtype = list(model.vision_tower.parameters())[0].dtype # for key, value in cleaned_weights.items(): # if value.dtype != target_dtype and value.dtype in [torch.float32, torch.float16, torch.bfloat16]: # cleaned_weights[key] = value.to(target_dtype) # # Use assign=True to avoid meta tensor copying warnings # model.vision_tower.load_state_dict(cleaned_weights, assign=True, strict=False) # print("✓ Vision weights loaded successfully") # print(f" Vision tower final dtype: {list(model.vision_tower.parameters())[0].dtype}") # except Exception as e: # print(f"✗ Error loading vision weights: {e}") # return False print("\nTesting model forward pass...") try: # Full model forward pass batch_size = 1 seq_len = 10 input_ids = torch.randint(0, 1000, (batch_size, seq_len)) # Move inputs to the same device as the model if hasattr(model, 'device'): device = next(model.parameters()).device input_ids = input_ids.to(device) with torch.no_grad(): outputs = model(input_ids=input_ids) print(f"✓ Forward pass successful") print(f" Output shape: {outputs.logits.shape}") print(f" Output dtype: {outputs.logits.dtype}") except Exception as e: print(f"✗ Error in forward pass: {e}") traceback.print_exc() return False # 清理临时debug目录 if debug_minimal_layers and model_path.endswith("_debug_temp"): try: # shutil.rmtree(model_path) print(f"🧹 Cleaned up temporary debug directory: {model_path}") except Exception as e: print(f"⚠️ Failed to clean up debug directory: {e}") elif debug_minimal_layers: print("\n🎉 DEBUG MINIMAL LAYERS: 3-layer model test completed!") print("💡 To test with full 61-layer model, run without DEBUG_MINIMAL_LAYERS=1") else: print("\n🎉 All tests passed! Model is working correctly.") return True if __name__ == "__main__": test_model_loading()