Support GPT-NeoX 20B (LLM)

[pure-rgb]

GPT-NeoX

Official pytorch: https://github.com/EleutherAI/gpt-neox Unofficial jax: https://github.com/kingoflolz/mesh-transformer-jax

[abheesht17]

+1

This is the open source version of GPT-3, right?

[pure-rgb]

yes, as far as I know. @abheesht17 wdyt? Is keras nlp ready to support this type of model? It's a very big model and unlikely to train for all.

[pure-rgb]

Minimal code pytorch and flax

• https://github.com/zphang/minimal-gpt-neox-20b

Some reading.

People use jax/flax to leverage TPU:

• https://github.com/kingoflolz/mesh-transformer-jax