MacOS M2上微调HuggingFace模型遇MPS生成器设备循环报错

问题根源 #

你碰到的是PyTorch 2.0.1在MPS设备上的生成器设备标识不兼容问题：设置PYTORCH_ENABLE_MPS_FALLBACK=1后，DataLoader期望生成器设备为cpu，但fallback逻辑会尝试调用MPS；指定generator=torch.Generator(device='mps')时，PyTorch内部又要求完整设备名mps:0，导致出现*"Expected a 'cpu' generator device type but found 'mps'"和"Expected a 'mps:0' generator device but found 'mps'"*的循环报错。

解决方案 #

1. 使用完整MPS设备标识 #

保持PYTORCH_ENABLE_MPS_FALLBACK=1环境变量，同时将生成器设备指定为mps:0：

train_dataloader = DataLoader(
    small_train_dataset, 
    shuffle=True, 
    batch_size=8, 
    generator=torch.Generator(device='mps:0').manual_seed(1)
)
eval_dataloader = DataLoader(
    small_eval_dataset, 
    batch_size=8, 
    generator=torch.Generator(device='mps:0')
)

该写法匹配PyTorch内部对MPS设备的标识规则，fallback机制会自动处理aten::random_这类不支持的算子。

2. 强制使用CPU生成器（调试用） #

若暂时不需要MPS加速，可关闭MPS fallback，全程用CPU生成器：

• 取消设置PYTORCH_ENABLE_MPS_FALLBACK=1环境变量
• 保持DataLoader生成器指定为CPU：

train_dataloader = DataLoader(
    small_train_dataset, 
    shuffle=True, 
    batch_size=8, 
    generator=torch.Generator(device='cpu').manual_seed(1)
)
eval_dataloader = DataLoader(
    small_eval_dataset, 
    batch_size=8, 
    generator=torch.Generator(device='cpu')
)

3. 升级PyTorch版本 #

PyTorch 2.1及以上版本已修复MPS生成器的设备标识问题，建议升级到最新稳定版：

pip install --upgrade torch

升级后无需手动指定生成器设备，PyTorch会自动处理MPS与fallback的兼容逻辑。

额外提示 #

• 模型训练时，需将DataLoader输出的batch张量移至MPS设备（若模型在MPS上）：

for batch in train_dataloader:
    batch = {key: tensor.to('mps') for key, tensor in batch.items()}
    # 执行模型前向传播、损失计算等逻辑

• 避免在dataset.shuffle()和DataLoader的shuffle=True重复设置种子，保持种子一致即可保证数据打乱的可复现性。

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