使用PEFT的LoRA微调NV-Embed-v2时遇inputs_embeds参数错误咨询

使用PEFT库的LoRA微调Hugging Face上的nvidia/NV-Embed-v2模型时，触发错误：

TypeError: NVEmbedModel.forward() got an unexpected keyword argument 'inputs_embeds'

已确认模型本身应支持inputs_embeds参数，当前使用FEATURE_EXTRACTION任务类型，通过AutoModel以8bit量化加载模型，相关代码如下：

from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, TaskType
import pandas as pd
from datasets import Dataset

# load the base model
model_name = "nvidia/NV-Embed-v2"
base_model = AutoModel.from_pretrained(model_name, trust_remote_code=True, load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# define lora config
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type=TaskType.FEATURE_EXTRACTION
)

peft_model = get_peft_model(base_model, lora_config)

# load the data
df = pd.DataFrame([['tomato', 'tomato ketchup heinz', 1], ['tomato', '500gm salted chips', 0], ['tomato', 'tomato 500gm', 1], ['strawberry', 'strawberry ripe 200gm', 1]], columns=['search_term', 'product_string', 'score'])
data = Dataset.from_pandas(df)
data = data.train_test_split(test_size=0.25)

train_dataset = data["train"]
test_dataset = data["test"]


def preprocess_function(examples):
    return tokenizer(examples["product_string"], examples["search_term"], truncation=True, padding="max_length", max_length=128)

train_dataset_final = train_dataset.map(preprocess_function, batched=True)

train_dataset_final.set_format(
    type="torch", columns=["input_ids", "attention_mask", "score"]
)

# Train the model
training_args = TrainingArguments(
    output_dir='lora-nvembedv2',
    auto_find_batch_size=True,
    learning_rate= 3e-2,
    num_train_epochs=1
)

trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=train_dataset_final,
)

trainer.train()

1. 修改PEFT任务类型 #

FEATURE_EXTRACTION任务类型会触发PEFT和Trainer的特定逻辑，自动传递inputs_embeds参数，而NVEmbedModel的forward方法不兼容该参数。将任务类型改为SEQ_CLS（序列分类），匹配你的打分任务本质：

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type=TaskType.SEQ_CLS  # 替换任务类型
)

2. 包装模型过滤不支持的参数 #

如果必须保留FEATURE_EXTRACTION类型，通过包装PEFT模型，在forward时移除inputs_embeds参数：

from peft import PeftModel

class WrappedNVEmbedModel(PeftModel):
    def forward(self, **kwargs):
        # 移除模型不支持的inputs_embeds参数
        kwargs.pop("inputs_embeds", None)
        return super().forward(**kwargs)

# 替换原peft_model初始化代码
peft_model = WrappedNVEmbedModel(base_model, lora_config)

3. 自定义Trainer的损失计算逻辑 #

重写Trainer的compute_loss方法，直接使用input_ids和attention_mask调用模型，避免Trainer自动生成inputs_embeds：

import torch

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # 分离任务标签
        labels = inputs.pop("score")
        # 用模型支持的参数调用forward
        outputs = model(**inputs)
        # 根据任务类型选择损失函数，这里用MSE适配0/1打分任务
        loss = torch.nn.functional.mse_loss(outputs.last_hidden_state[:, 0, :].mean(dim=-1), labels.float())
        return (loss, outputs) if return_outputs else loss

# 使用自定义Trainer替代原Trainer
trainer = CustomTrainer(
    model=peft_model,
    args=training_args,
    train_dataset=train_dataset_final,
)

4. 确认模型forward参数 #

查看NVEmbedModel的forward方法签名，确保只传递模型明确支持的参数（如input_ids、attention_mask），避免Trainer触发自动生成inputs_embeds的逻辑。

内容的提问来源于stack exchange，提问作者yaksh