RuntimeError: Неизвестное устройство при попытке запустить AlbertForMaskedLM на Colab TPU - PullRequest
Новая
       97

RuntimeError: Неизвестное устройство при попытке запустить AlbertForMaskedLM на Colab TPU

1 голос
/ 11 апреля 2020

Я выполняю следующий код на 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]

Я ничего не сделал с примером кода, кроме перемещения input_ids и model на устройстве TPU, используя .to(dev). Кажется, все перемещается в 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

Кто-нибудь знает, что происходит?

1 Ответ

1 голос
/ 15 апреля 2020

Решение здесь: 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)

Есть также некоторые проблемы с получением функции активации gelu, которую ALBERT использует для работы на 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

...