Running the model on TPUs?

salesforce / ctrl

Conditional Transformer Language Model for Controllable Generation

https://arxiv.org/abs/1909.05858

BSD 3-Clause "New" or "Revised" License

1.87k stars 208 forks source link

Running the model on TPUs? #55

Open vessenes opened 5 years ago

vessenes commented 5 years ago

Hi,

I have the 256 and 512 models working on GCP with a Tesla V100. Text generates, but slowly, and I'm wanting to get faster generation out of the system. I thought running CTRL on TPUs could get me faster text, but I have no idea how to do that.

Do you have an incantation or pointer that would let me point CTRL at a TPU?

dimitri320 commented 5 years ago

Second this!

keskarnitish commented 5 years ago

I haven't quite figured out how to get TPUs to be faster than GPUs for inference. I'll probably look into this soon. It's especially more complicated with top-k/nucleus sampling and other add-ons. Seems like others have found the same behavior.