Inference on TPU's

[Arya Manjaramkar]

Is it possible to setup CodeGen for inference on TPU's?

According to Huggingface, Pytorch-XLA generation is not yet supported: huggingface/transformers#12322

Is there another way to use this model on TPU's for inference?

Thanks,
Arya

[Erik Nijkamp]

Hi @Aryagm,

The log-likelihood and sampling can be run on TPUs using JAX.

I'm currently verifying the training code on TPU-v3 and then will release the implementation.

I can include sampling code on TPU, if this helps.

[Arya Manjaramkar]

Hi @Aryagm,

The log-likelihood and sampling can be run on TPUs using JAX.

I'm currently verifying the training code on TPU-v3 and then will release the implementation.

I can include sampling code on TPU, if this helps.

Hi,

Including the code for inference would be of a lot of help!

I need it to upgrade an open-source project I am currently working on.

Thanks,
Arya

[Arya Manjaramkar]

Hi!

Any update on when the TPU sampling code will be available?

I do not require the fine-tuning code, just the code for sampling.

Have a great day!

Thanks,
Arya

[Xander Dunn]

The jax code that @enijkamp is referring to has been published here: https://github.com/salesforce/jaxformer
Unfortunately it's focused on training and not immediately clear how to use the code for inference on a pre-trained model

[Erik Nijkamp]

@xanderdunn @Aryagm I'm working on adding the pjit'ed sampling code as we speak, sorry for those delays.

[Bryan]

Thanks for the great work! am also interested in the sampling code, would be really useful!

(btw for those who really need it, I think you could load it into gpt-J and then use sampling code in that repository, but I haven't tried it myself)

[Arya Manjaramkar]

Hi @enijkamp!

Any update on this?

Thanks,
Arya