Maxtext: A simple, performant and scalable Jax LLM

Re:Maxtext: A simple, performant and scalable Jax LLM by bionhoward on Wednesday April 24, 2024 @05:19AM (#40142231)

Is t5x an encoder/decoder architecture?

Some more general options.

The Flax ecosystem

https://github.com/google/flax?tab=readme-ov-file

or dm-haiku

https://github.com/google-deepmind/dm-haiku

were some of the best developed communities in the Jax AI field

Perhaps the “trax” repo? https://github.com/google/trax

Some HF examples https://github.com/huggingface/transformers/tree/main/exampl...

Sadly it seems much of the work is proprietary these days, but one example could be Grok-1, if you customize the details. https://github.com/xai-org/grok-1/blob/main/run.py

[ Parent | on HN ]

Re:Maxtext: A simple, performant and scalable Jax LLM by terafo on Wednesday April 24, 2024 @07:37AM (#40143104)

t5 is an architecture, t5x is a framework for training models that was created with that architecture in mind, but can be used to train other architectures, including decoder-only ones(there is one in examples).

[ Parent | on HN ]

Re:Maxtext: A simple, performant and scalable Jax LLM by ma2rten on Wednesday April 24, 2024 @11:01AM (#40145186)

t5x was used to train PaLM 1.

[ Parent | on HN ]