使用CUDA微调Sentence-BERT生成model.safetensors无法加载问题求助

核心原因 #

你推测的没错，transformers库新版本（4.29+）在CUDA环境下默认使用safetensors格式保存模型权重，而部分旧版本的sentence-transformers对该格式支持不完善，导致加载时找不到pytorch_model.bin报错。以下是几种直接可行的解决方法：

方法1：代码层面强制保存PyTorch格式 #

在初始化Transformer模块时，添加model_save_kwargs参数禁用safetensors序列化，确保无论CPU/CUDA环境都生成pytorch_model.bin：

checkpoint = 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'

# 添加model_save_kwargs参数强制用PyTorch格式保存
word_embedding_model = models.Transformer(
    checkpoint, 
    cache_dir=f'model/{checkpoint}',
    model_save_kwargs={"safe_serialization": False}
)
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='mean')
model = SentenceTransformer(modules=[word_embedding_model, pooling_model], device='cuda')

# 后续训练代码不变
train_loss = losses.CosineSimilarityLoss(model)
evaluator = evaluation.EmbeddingSimilarityEvaluator.from_input_examples(val_examples, name='sbert')
model.fit(
    train_objectives=[(train_dataloader, train_loss)], 
    epochs=5, 
    evaluator=evaluator, 
    show_progress_bar=True, 
    output_path=f'model_FT/{checkpoint}', 
    save_best_model=True
)

方法2：降级transformers库版本 #

如果不想修改代码，可将transformers降级到默认不使用safetensors的版本（如4.28.x系列），避免版本兼容问题：

pip install transformers==4.28.1

注意降级后需确认与当前sentence-transformers版本兼容，可通过pip show sentence-transformers查看其依赖的transformers版本范围。

方法3：手动转换已生成的safetensors文件 #

如果已经得到了model.safetensors，可通过代码手动转换为pytorch_model.bin：

import torch
from safetensors.torch import load_file

# 加载safetensors权重
weights = load_file("model_FT/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/model.safetensors")
# 保存为PyTorch格式
torch.save(weights, "model_FT/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/pytorch_model.bin")

转换后即可被sentence-transformers正常加载。

额外建议 #

若环境允许，可直接升级sentence-transformers到2.2.0及以上版本，该版本已原生支持safetensors格式，无需修改保存逻辑即可正常加载：

pip install --upgrade sentence-transformers

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