Xla: TypeError: can't pickle torch._C.ScriptFunction objects when calling xm.send_cpu_data_to_device(model, dev)

[goggoloid]

Related issue: https://github.com/pytorch/xla/issues/1909

In the linked issue above I was running the following example code, which returned an Unknown device error:

```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]

The problem was that the model had lots of cpu default tensors which couldn't be moved onto the TPU device with ```.to(dev)```--this was solved by using ```xm.send_cpu_data_to_device(model, xm.xla_device())```

But now when I use ```xm.send_cpu_data_to_device(model, xm.xla_device())``` in the same code above, I get the following error:

---

TypeError Traceback (most recent call last)
in ()
11 tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
12 model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
---> 13 model = xm.send_cpu_data_to_device(model, dev)
14 model = model.to(dev)
15 input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) # Batch size 1

18 frames
/usr/lib/python3.6/copy.py in copy(x)
94 reductor = getattr(x, "__reduce_ex__", None)
95 if reductor:
---> 96 rv = reductor(4)
97 else:
98 reductor = getattr(x, "__reduce__", None)

TypeError: can't pickle torch._C.ScriptFunction objects

Edited version of the above code using ```xm.send_cpu_data_to_device(model, xm.xla_device())``` which returns the TypeError:

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

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')
model = xm.send_cpu_data_to_device(model, dev)
model = model.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]
print(f'loss: {loss}, prediction_scores: {prediction_scores}')
```

[dlibenzi]

It is the same underline issue.
The data-copy-to-device internally uses copy.copy().
So there is currently no way to transfer to XLA device models which has torch-script inside, unfortunately.

[goggoloid]

Unfortunate! How come it was working a few days ago? It no longer works because of a huggingface update?

[dlibenzi]

Yeah, that is likely the culprit.
We do not mess up with pickling in our code.
@jysohn23 any idea, since you followed huggingface?

[jysohn23]

Hi @goggoloid sorry for the late reply! I'll try to repro this today and see why it recently broke. Thanks.

[NaxAlpha]

I have faced similar issue with GPT-2 model. And apparently it is being caused by jitted activation functions. Here is a quick fix which worked in my case: (put this code before creating your model)

import math
from transformers import activations

def gelu_new(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

activations.ACT2FN['gelu_new'] = gelu_new

[khalilRhouma]

@NaxAlpha I am facing the same problem right now, I am working with XLNet model, I have tried your code but still facing the same error:
TypeError: can't pickle torch._C.ScriptFunction objects

[NaxAlpha]

@khalilRhouma After looking at xlnet config, they are using gelu not gelu_new. You may try this for gelu:

import math
from transformers import activations

def _gelu_python(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

activations.ACT2FN['gelu'] = _gelu_python

[khalilRhouma]

@khalilRhouma After looking at xlnet config, they are using gelu not gelu_new. You may try this for gelu:

import math
from transformers import activations

def _gelu_python(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

activations.ACT2FN['gelu'] = _gelu_python

@NaxAlpha thanks, I will try that

[khalilRhouma]

I was able to fix this issue when I delete transformers package and clone it from git and installed locally.

[goggoloid]

I was able to fix this issue when I delete transformers package and clone it from git and installed locally.

This seems to have worked for me too. I guess transformers has been fixed!