https://github.com/kingoflolz/mesh-transformer-jax/#gpt-j-6b
They even have a google colab notebook thingie. Not much setup needed, just an account.
https://colab.research.google.com/github/kingoflolz/mesh-tra...