Apple M2 Pro运行Sentence-Transformers训练报MPS设备RuntimeError

环境信息 #

• 设备：Apple M2 Pro
• Python版本：3.11
• 依赖版本：
  • sentence-transformers 3.0.1
  • accelerate 0.32.1
  • torch 2.3.1

训练代码 #

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, losses
from datasets import Dataset

path = "/Users/austin/Documents/Career/huggingface"
model_directory = path + "/hub/all-mpnet-base-v2"

model = SentenceTransformer(model_directory)
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
    "negative": ["It's quite rainy, sadly.", "She walked to the store."],
})

loss = losses.TripletLoss(model=model)

args = SentenceTransformerTrainingArguments(
    output_dir="test_trainer",
    use_mps_device=True,
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

报错信息 #

运行时触发错误：

RuntimeError: Placeholder storage has not been allocated on MPS device!

MPS设备验证 #

已通过以下代码确认MPS设备可用：

import torch
if torch.backends.mps.is_available():
    mps_device = torch.device("mps")
    x = torch.ones(1, device=mps_device)
    print (x)
else:
    print ("MPS device not found.")

输出结果：

tensor([1.], device='mps:0')

解决方案 #

该错误源于模型或部分张量未正确迁移至MPS设备，可通过以下步骤修复：

1. 手动将模型迁移至MPS设备 #

初始化模型后显式调用to("mps")，确保所有权重加载到MPS设备：

model = SentenceTransformer(model_directory)
model = model.to("mps")  # 添加此行

2. 关闭自动混合精度（AMP） #

MPS对FP16自动混合精度的支持存在兼容性问题，在训练参数中禁用该功能：

args = SentenceTransformerTrainingArguments(
    output_dir="test_trainer",
    use_mps_device=True,
    fp16=False  # 添加此行
)

3. 基于MPS设备的模型初始化损失函数 #

调整代码顺序，先迁移模型到MPS，再初始化损失函数，避免损失内部张量留在CPU：

model = SentenceTransformer(model_directory)
model = model.to("mps")
loss = losses.TripletLoss(model=model)  # 模型迁移后再初始化损失

修改后的完整代码 #

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, losses
from datasets import Dataset

path = "/Users/austin/Documents/Career/huggingface"
model_directory = path + "/hub/all-mpnet-base-v2"

model = SentenceTransformer(model_directory)
model = model.to("mps")  # 手动迁移模型到MPS设备

train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
    "negative": ["It's quite rainy, sadly.", "She walked to the store."],
})

loss = losses.TripletLoss(model=model)  # 基于MPS模型初始化损失

args = SentenceTransformerTrainingArguments(
    output_dir="test_trainer",
    use_mps_device=True,
    fp16=False  # 关闭自动混合精度
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

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