RuntimeError：尝试在 colab tpu 上运行 AlbertForMaskedLM 时出现未知设备答案

【问题标题】：RuntimeError: Unknown device when trying to run AlbertForMaskedLM on colab tpuRuntimeError：尝试在 colab tpu 上运行 AlbertForMaskedLM 时出现未知设备
【发布时间】：2020-07-24 04:58:29
【问题描述】：

我在 colab 上运行以下代码，取自此处的示例：https://huggingface.co/transformers/model_doc/albert.html#albertformaskedlm

import os
import torch
import torch_xla
import torch_xla.core.xla_model as xm

assert os.environ['COLAB_TPU_ADDR']

dev = xm.xla_device()

from transformers import AlbertTokenizer, AlbertForMaskedLM
import torch

tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
model = AlbertForMaskedLM.from_pretrained('albert-base-v2').to(dev)
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1

data = input_ids.to(dev)

outputs = model(data, masked_lm_labels=data)
loss, prediction_scores = outputs[:2]

除了使用.to(dev) 将input_ids 和model 移动到TPU 设备上之外，我没有对示例代码做任何事情。似乎一切都移到了 TPU 没有问题，因为当我输入 data 时，我得到以下输出：tensor([[ 2, 10975, 15, 51, 1952, 25, 10901, 3]], device='xla:1')

但是，当我运行此代码时，出现以下错误：

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-f756487db8f7> in <module>()
      1 
----> 2 outputs = model(data, masked_lm_labels=data)
      3 loss, prediction_scores = outputs[:2]

9 frames
/usr/local/lib/python3.6/dist-packages/transformers/modeling_albert.py in forward(self, hidden_states, attention_mask, head_mask)
    277         attention_output = self.attention(hidden_states, attention_mask, head_mask)
    278         ffn_output = self.ffn(attention_output[0])
--> 279         ffn_output = self.activation(ffn_output)
    280         ffn_output = self.ffn_output(ffn_output)
    281         hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])

RuntimeError: Unknown device

有人知道怎么回事吗？

【问题讨论】：

标签： nlp pytorch tpu huggingface-transformers tensorflow-xla

【解决方案1】：

解决方案在这里：https://github.com/pytorch/xla/issues/1909

在调用model.to(dev)之前，需要先调用xm.send_cpu_data_to_device(model, xm.xla_device())：

model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
model = xm.send_cpu_data_to_device(model, dev)
model = model.to(dev)

获取 ALBERT 使用的 Gelu 激活函数在 TPU 上工作也存在一些问题，因此在 TPU 上工作时需要使用以下转换器分支：https://github.com/huggingface/transformers/tree/fix-jit-tpu

请参阅以下 colab 笔记本（https://github.com/jysohn23）以获取完整解决方案：https://colab.research.google.com/gist/jysohn23/68d620cda395eab66289115169f43900/getting-started-with-pytorch-on-cloud-tpus.ipynb

【讨论】：