Colab与Apple M2 Max上WER结果差异问题求助

• 同步输入数据到MPS设备
  别只移动模型，音频张量、注意力掩码这类输入数据也得一起移到MPS上，不然模型在MPS运算但数据在CPU，会导致逻辑错乱输出空结果。代码示例：

  input_values = input_values.to(mps_device)
  attention_mask = attention_mask.to(mps_device)
  

• 关闭自动混合精度（AMP）
  MPS对AMP的支持还不完善，尤其是Wav2Vec2的CTC损失计算容易出问题。直接用FP32训练，删掉代码里的torch.cuda.amp.autocast()上下文，损失计算改为：

  loss = model(input_values, attention_mask=attention_mask, labels=labels).loss
  

• 解码前将模型输出移回CPU
  MPS张量直接参与CTC解码可能出现索引异常，导致空字符串输出。先把模型输出转到CPU再执行解码操作：

  logits = model(input_values, attention_mask=attention_mask).logits
  logits = logits.cpu().detach()
  # 后续执行argmax、词典映射等解码步骤
  

• 升级PyTorch与Transformers版本
  旧版本PyTorch对MPS的支持存在不少算子兼容bug，尤其针对语音模型。直接升级到最新稳定版：

  pip install --upgrade torch transformers datasets accelerate
  

  确保PyTorch版本≥2.0、Transformers版本≥4.28，能更好适配Apple Silicon的MPS后端。

• 验证模型权重一致性
  模型迁移到MPS时可能出现权重数值偏差，导致结果完全失效。可以对比Colab和M2上的关键层权重均值：

  # 在Colab上保存权重均值
  torch.save(model.feature_extractor.conv_layers[0].weight.mean().item(), "weight_mean.pt")
  # 在M2上加载后对比差值
  mean_colab = torch.load("weight_mean.pt")
  mean_mps = model.feature_extractor.conv_layers[0].weight.mean().item()
  print(abs(mean_colab - mean_mps))  # 差值应在1e-6以内，差太多就重新加载模型再迁移设备
  

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