""" Qwen 模型配置和加载 """ from dataclasses import dataclass from typing import List import os import torch import torch.nn as nn import torch.nn.functional as F import logging from transformers import AutoModelForCausalLM, AutoTokenizer from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from finetunex.models.base import BaseModelConfig logger = logging.getLogger(__name__) @dataclass class QwenConfig(BaseModelConfig): """Qwen 模型专用配置""" model_name: str = "Qwen/Qwen3.5-0.5B" target_modules: List[str] = None def __post_init__(self): if self.target_modules is None: self.target_modules = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] def is_qwen3_5(self) -> bool: return "qwen3.5" in self.model_name.lower() or "qwen3_5" in self.model_name.lower() def is_qwen3(self) -> bool: name = self.model_name.lower() return "qwen3" in name and "qwen3.5" not in name and "qwen3_5" not in name class _NPUConv1d(nn.Module): """华为升腾 NPU 兼容的 Conv1d 实现 使用 unfold + einsum 替代 F.conv1d， 避免 NPU 上 Conv2D 算子编译失败的问题。 """ def __init__(self, original_conv1d: nn.Conv1d): super().__init__() self.stride = original_conv1d.stride[0] self.padding = original_conv1d.padding[0] self.dilation = original_conv1d.dilation[0] self.groups = original_conv1d.groups self.kernel_size = original_conv1d.kernel_size[0] self.in_channels = original_conv1d.in_channels self.out_channels = original_conv1d.out_channels self.weight = original_conv1d.weight self.bias = original_conv1d.bias def forward(self, input: torch.Tensor) -> torch.Tensor: if self.padding > 0: input = F.pad(input, (self.padding, 0)) if self.kernel_size == 1 and self.stride == 1 and self.dilation == 1: output = F.linear(input.transpose(1, 2), self.weight.squeeze(-1), self.bias) return output.transpose(1, 2) unfolded = input.unfold(2, self.kernel_size, self.stride) weight = self.weight output = torch.einsum('bci,oci->bo', unfolded, weight) if self.bias is not None: output = output + self.bias.unsqueeze(0) return output def _patch_conv1d_for_npu(): """Monkey-patch Conv1d 使其在华为升腾 NPU 上使用纯 PyTorch 实现""" original_forward = nn.Conv1d.forward if hasattr(nn.Conv1d, '_npu_patched'): return def npu_conv1d_forward(self, input): try: return original_forward(self, input) except RuntimeError as e: if "Conv2D" in str(e) or "500001" in str(e): logger.info(f"Conv1d 在 NPU 上失败，回退到纯 PyTorch 实现: {e}") fallback = _NPUConv1d(self) return fallback(input) raise nn.Conv1d.forward = npu_conv1d_forward nn.Conv1d._npu_patched = True logger.info("已应用 Conv1d NPU 兼容补丁") def _patch_qwen3_5_for_npu(model_path: str): """修补 Qwen3.5 模型配置以兼容华为升腾 NPU""" import json config_path = os.path.join(model_path, "config.json") if not os.path.exists(config_path): logger.warning(f"未找到模型配置文件：{config_path}") return with open(config_path, "r", encoding="utf-8") as f: model_config = json.load(f) if model_config.get("model_type") != "qwen3_5": return changed = False if "linear_attn" in model_config: logger.info("检测到 linear_attn 配置，NPU 不支持，将替换为 sdpa attention") del model_config["linear_attn"] changed = True if model_config.get("_attn_implementation", "") == "linear": logger.info("检测到 _attn_implementation=linear，将替换为 eager") model_config["_attn_implementation"] = "eager" changed = True attn_layers = model_config.get("attention_layers", None) if attn_layers: if any(v == "linear_attn" for v in attn_layers.values()): logger.info("检测到 attention_layers 中包含 linear_attn，将替换为 eager") model_config["attention_layers"] = { k: "eager" if v == "linear_attn" else v for k, v in attn_layers.items() } changed = True if changed: backup_path = config_path + ".bak" if not os.path.exists(backup_path): import shutil shutil.copy2(config_path, backup_path) logger.info(f"原始配置已备份到：{backup_path}") with open(config_path, "w", encoding="utf-8") as f: json.dump(model_config, f, indent=2, ensure_ascii=False) logger.info("模型配置已修改，linear attention 已替换为 sdpa attention") def load_qwen_model(config: QwenConfig): """加载 Qwen 模型""" print(f"正在加载模型：{config.model_name}") is_npu = hasattr(torch, 'npu') and torch.npu.is_available() is_qwen3_5 = config.is_qwen3_5() is_qwen3 = config.is_qwen3() if is_npu: logger.info("检测到华为升腾 NPU，应用兼容性补丁...") if is_qwen3_5: _patch_conv1d_for_npu() _patch_qwen3_5_for_npu(config.model_name) elif is_qwen3: logger.info("Qwen3 模型使用标准 attention，NPU 兼容性良好，无需额外补丁") compute_dtype = config.get_compute_dtype() use_quantization = config.use_4bit if use_quantization: try: from transformers import BitsAndBytesConfig bnb_config = BitsAndBytesConfig( load_in_4bit=config.use_4bit, bnb_4bit_quant_type=config.bnb_4bit_quant_type, bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=config.use_nested_quant, ) print("使用 4bit 量化加载模型（需要 NVIDIA GPU）") except (ImportError, Exception) as e: logger.warning(f"无法使用 4bit 量化: {e}") logger.warning("将使用 bf16/fp16 加载模型") use_quantization = False bnb_config = None else: bnb_config = None print(f"使用 {compute_dtype} 精度加载模型") tokenizer = AutoTokenizer.from_pretrained( config.model_name, trust_remote_code=config.trust_remote_code, padding_side="right", ) tokenizer.pad_token = tokenizer.eos_token model_kwargs = { "quantization_config": bnb_config if use_quantization else None, "device_map": "auto", "trust_remote_code": config.trust_remote_code, "torch_dtype": compute_dtype, } if is_npu: model_kwargs["attn_implementation"] = "eager" model = AutoModelForCausalLM.from_pretrained( config.model_name, **model_kwargs, ) if use_quantization: model = prepare_model_for_kbit_training(model) peft_config = LoraConfig( lora_alpha=config.lora_alpha, lora_dropout=config.lora_dropout, r=config.lora_r, bias="none", task_type="CAUSAL_LM", target_modules=config.target_modules, ) model = get_peft_model(model, peft_config) print(f"模型加载完成！可训练参数：{model.print_trainable_parameters()}") return model, tokenizer, peft_config