ola/model/builder_back.py

import os
import warnings
import shutil

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig, AutoProcessor
import torch
from ola.model import *
from ola.model.speech_encoder.builder import build_speech_encoder

def load_pretrained_model(model_path, model_type, model_base, is_lora=False, s2s=False, load_8bit=False, load_4bit=False, device="cuda", use_flash_attn=False, **kwargs):
    device = "cuda"
    if load_8bit:
        kwargs['load_in_8bit'] = True
    elif load_4bit:
        kwargs['load_in_4bit'] = True
        kwargs['quantization_config'] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type='nf4'
        )
    else:
        kwargs['torch_dtype'] = torch.bfloat16

    if use_flash_attn:
        kwargs['attn_implementation'] = 'flash_attention_2'

    if model_type == 'ola_internvl':
        model_cls = OlaQwen3ForCausalLM
        print('Loading OlaQwen3ForCausalLM model...')
    else:
        model_cls = OlaQwenForCausalLM

    # Load Ola model
    if is_lora:
        assert model_base is not None, "model_base is required for LoRA models."
        from ola.model.language_model.ola_qwen import OlaConfigQwen
        lora_cfg_pretrained = OlaConfigQwen.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
        print('Loading Ola from base model...')
        model = model_cls.from_pretrained(model_base, low_cpu_mem_usage=False, config=lora_cfg_pretrained, **kwargs)
        print('Loading additional Ola weights...')
        if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
            non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu')
        non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}
        if any(k.startswith('model.model.') for k in non_lora_trainables):
            non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}
        model.load_state_dict(non_lora_trainables, strict=False, assign=True)

        from peft import PeftModel
        print('Loading LoRA weights...')
        model = PeftModel.from_pretrained(model, model_path)
        print('Merging LoRA weights...')
        model = model.merge_and_unload()
        print('Model is loaded...')
    elif model_base is not None:
        print('Loading Ola from base model...')
        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
        cfg_pretrained = AutoConfig.from_pretrained(model_path)
        model = model_cls.from_pretrained(model_base, low_cpu_mem_usage=False, config=cfg_pretrained, **kwargs)
        
        speech_projector_weights = torch.load(os.path.join(model_path, 'speech_projector.bin'), map_location='cpu')
        speech_projector_weights = {k: v.to(torch.float16) for k, v in speech_projector_weights.items()}
        model.load_state_dict(speech_projector_weights, strict=False, assign=True)
        model = model.to(device=device)
    elif model_type == 'ola_internvl':
        cfg = AutoConfig.from_pretrained("/data1/cxy/plm-v/modeling/old_ola", trust_remote_code=True)
        # breakpoint()
        tokenizer = AutoTokenizer.from_pretrained("/data1/cxy/plm-v/modeling/internvl3_5-2B", use_fast=False)
        with torch.device("cpu"):
            # model = model_cls.from_pretrained("/data1/cxy/plm-v/modeling/internvl3_5-2B", low_cpu_mem_usage=False, attn_implementation="eager", config=cfg, **kwargs)
            # model = model_cls.from_config(config=cfg)
            model = model_cls(cfg)
        # breakpoint()
        # model.model.layers[1].self_attn.q_proj.weight
    else:
        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
        with torch.device("cpu"):
            model = model_cls.from_pretrained(
                model_path,
                **kwargs,
            )
        model = model.to(device=device)
    # model.resize_token_embeddings(len(tokenizer))
    from safetensors.torch import load_file
    partial_state_dict = load_file(f"/data1/cxy/plm-v/modeling/internvl3_5-2B/model.safetensors") # 替换为你的部分权重路径
    mapping = {
    "mlp1.0.weight": "model.mm_projector.layer_norm.weight",
    "mlp1.0.bias": "model.mm_projector.layer_norm.bias",
    "mlp1.1.weight": "model.mm_projector.linear_1.weight",
    "mlp1.1.bias": "model.mm_projector.linear_1.bias",
    "mlp1.3.weight": "model.mm_projector.linear_2.weight",
    "mlp1.3.bias": "model.mm_projector.linear_2.bias",
    }

# 遍历 state_dict 并重命名
    def remap_keys(state_dict, mapping):
        new_state_dict = {}
        for k, v in state_dict.items():
            if k in mapping:
                new_state_dict[mapping[k]] = v
            else:
                new_state_dict[k] = v
        return new_state_dict
    # merged_state_dict = {**partial_state_dict, **partial_state_dict2}
    # 2. 重命名 key：multi_modal_projector -> mm_projector
    # breakpoint()
    rename_dict = {}
    for k in list(partial_state_dict.keys()):
        if k.startswith("language_model"):
            new_k = k.replace("language_model.", "", 1)
            rename_dict[k] = new_k
        if k.startswith("vision_model"):
            new_k = k.replace("vision_model", "model.vision_tower", 1)
            rename_dict[k] = new_k

    # 应用重命名
    for old_k, new_k in rename_dict.items():
        partial_state_dict[new_k] = partial_state_dict.pop(old_k)
    partial_state_dict = remap_keys(partial_state_dict, mapping)

    whisper_state_dict = torch.load("/data1/cxy/model/THUdyh/Ola-7b/large-v3.pt", map_location='cpu')
    # breakpoint()
    whisper_state_dict = whisper_state_dict["model_state_dict"]
    
    # Filter to keep only encoder weights
    whisper_encoder_dict = {}
    for key, value in whisper_state_dict.items():
        if key.startswith('encoder.'):
            whisper_encoder_dict[key] = value
    
    print(f"Original Whisper keys: {len(whisper_state_dict)}")
    print(f"Filtered encoder keys: {len(whisper_encoder_dict)}")
    print("Sample encoder keys:")
    for i, key in enumerate(list(whisper_encoder_dict.keys())[:5]):
        print(f"  {key}")
    
    # Create mapping for Whisper parameters to OLA format
    def create_whisper_mapping():
        mapping = {}
        
        # Base encoder components
        base_mappings = {
            'encoder.positional_embedding': 'model.speech_encoder.whisper_model.positional_embedding',
            'encoder.conv1.weight': 'model.speech_encoder.whisper_model.conv1.weight',
            'encoder.conv1.bias': 'model.speech_encoder.whisper_model.conv1.bias',
            'encoder.conv2.weight': 'model.speech_encoder.whisper_model.conv2.weight',
            'encoder.conv2.bias': 'model.speech_encoder.whisper_model.conv2.bias',
            'encoder.ln_post.weight': 'model.speech_encoder.whisper_model.ln_post.weight',
            'encoder.ln_post.bias': 'model.speech_encoder.whisper_model.ln_post.bias',
        }
        mapping.update(base_mappings)
        
        # Encoder blocks (32 blocks: 0-31)
        for block_idx in range(32):
            # Attention components
            attn_components = [
                'attn.query.weight', 'attn.query.bias',
                'attn.key.weight', 'attn.key.bias', 
                'attn.value.weight', 'attn.value.bias',
                'attn.out.weight', 'attn.out.bias',
                'attn_ln.weight', 'attn_ln.bias'
            ]
            
            for component in attn_components:
                source_key = f'encoder.blocks.{block_idx}.{component}'
                target_key = f'model.speech_encoder.whisper_model.blocks.{block_idx}.{component}'
                mapping[source_key] = target_key
            
            # MLP components
            mlp_components = [
                'mlp.0.weight', 'mlp.0.bias',
                'mlp.2.weight', 'mlp.2.bias',
                'mlp_ln.weight', 'mlp_ln.bias'
            ]
            
            for component in mlp_components:
                source_key = f'encoder.blocks.{block_idx}.{component}'
                target_key = f'model.speech_encoder.whisper_model.blocks.{block_idx}.{component}'
                mapping[source_key] = target_key
        
        return mapping
    
    # Apply mapping to whisper_encoder_dict
    whisper_mapping = create_whisper_mapping()
    mapped_whisper_dict = {}
    unmapped_whisper_keys = []
    
    for key, value in whisper_encoder_dict.items():
        if key in whisper_mapping:
            mapped_key = whisper_mapping[key]
            mapped_whisper_dict[mapped_key] = value
        else:
            unmapped_whisper_keys.append(key)
            print(f"Warning: No mapping found for Whisper encoder key '{key}'")
    
    if unmapped_whisper_keys:
        print(f"Total unmapped Whisper encoder keys: {len(unmapped_whisper_keys)}")
        print("First 10 unmapped Whisper encoder keys:")
        for key in unmapped_whisper_keys[:10]:
            print(f"  {key}")
    
    print(f"Successfully mapped {len(mapped_whisper_dict)} encoder parameters")
    
    beat_state_dict = torch.load("/data1/cxy/model/THUdyh/Ola-7b//BEATs_iter3_plus_AS2M_finetuned_on_AS2M_cpt2.pt", map_location='cpu')
    beat_state_dict = beat_state_dict['model']
    beat_state_dict = {"model.speech_encoder.beats_model."+k: v for k, v in beat_state_dict.items()}
    
    # 处理 BEATs 模型中的参数化权重映射 (先pop后添加)
    keys_to_process = list(beat_state_dict.keys())
    breakpoint()
    processed_count = 0
    
    # for key in keys_to_process:
    #     if 'weight_g' in key:
    #         # pop 原始权重并添加为 weight_g
    #         weight_tensor = beat_state_dict.pop(key)
    #         new_key = key.replace('weight_g','parametrizations.weight.original0')
    #         beat_state_dict[new_key] = weight_tensor
    #         processed_count += 1
    #     elif 'weight_v' in key:
    #         # pop 原始权重并添加为 weight_v
    #         weight_tensor = beat_state_dict.pop(key)
    #         new_key = key.replace('weight_v', 'parametrizations.weight.original1')
    #         beat_state_dict[new_key] = weight_tensor
    #         processed_count += 1
    
    print(f"Processed {processed_count} parametrized weight keys in BEATs model (pop and add)")
    breakpoint()
    # breakpoint()
    partial_state_dict = {**partial_state_dict, **mapped_whisper_dict, **beat_state_dict}
    
    # Ensure all tensors in the state dict are on CPU and have proper device information
    print("Moving all state dict tensors to CPU...")
    for key, tensor in partial_state_dict.items():
        if torch.is_tensor(tensor):
            # Ensure tensor has device information and move to CPU
            if not tensor.device.type:
                print(f"Warning: Tensor {key} has no device, creating on CPU")
                partial_state_dict[key] = torch.tensor(tensor.detach().numpy()).cpu()
            else:
                partial_state_dict[key] = tensor.cpu()
    
    # Ensure model is on CPU before loading state dict to avoid device mismatches
    print("Moving model to CPU before loading state dict...")
    model = model.cpu()
    
    print("Loading state dict...")
    breakpoint()
    missing, unexpected = model.load_state_dict(partial_state_dict, strict=False, assign=True)

    print("Missing keys:", missing)
    print("Unexpected keys:", unexpected)
    
    # Convert model to bfloat16 before saving
    print("Converting model to bfloat16...")
    model = model.to(torch.bfloat16)
    model = model.to("cpu")
    
    # Save model in bfloat16 format
    print("Saving model in bfloat16 format...")
    model.save_pretrained("/data1/cxy/plm-v/modeling/plm_internvl3_ola", safe_serialization=False, torch_dtype=torch.bfloat16)
    print("Model saved successfully in bfloat16 format!")
    breakpoint()
    # model.model.mm_projector.linear_1.weight:-0.0106 multi_modal_projector.linear_1.weight model.mm_projector.linear_2.bias
    # model.vision_tower.encoder.layers.7.attn.proj.bias
    # model.model.vision_tower.encoder.layers[0].attn.qkv.weight: -6.5613e-03 dui
    # 
    # breakpoint()
    # model.get_model().speech_encoder.load_model("")
    # language_model.model.layers.9.mlp.up_proj.weight vision_model.encoder.layers
    # model.layers.14.self_attn.q_proj.weight model.vision_tower.encoder.layers.23.attn.proj.bias
    # model.get_model().speech_encoder = build_speech_encoder(model.config)
    # model.get_model().speech_encoder.to(device=device, dtype=torch.float16)
    image_processor = None
    model.resize_token_embeddings(len(tokenizer))
    vision_tower = model.get_vision_tower()
    print("Loading vision tower...")
    # if not vision_tower.is_loaded:
    #     vision_tower.load_model(device_map=device)
    # if device != "auto":
    #     vision_tower.to(device="cuda", dtype=torch.bfloat16)
    # else:
    #     vision_tower.to(device="cuda:0", dtype=torch.bfloat16)
    # image_processor = vision_tower.image_processor
    print("Loading vision tower succeeded.")
    
    if hasattr(model.config, "max_sequence_length"):
        context_len = model.config.max_sequence_length
    else:
        context_len = 16384
    image_processor = AutoProcessor.from_pretrained("/data1/cxy/plm-v/modeling/internvl3_5-2B-HF")
    # breakpoint()
    return tokenizer, model, image_processor, context_len