salesforce / CodeGen

CodeGen is an open-source model for program synthesis. Trained on TPU-v4. Competitive with OpenAI Codex.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Inference on TPU's

Aryagm opened this issue · comments

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

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 @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

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

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

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

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)

Hi @enijkamp!

Any update on this?

Thanks,
Arya