2 ms·
Did you consider using https://github.com/AI-Hypercomputer/maxtext https://github.com/AI-Hypercomputer/maxtext ? It has a Jax llama implementation, and gets dec
by logicchains 2y ago
Did you consider using https://github.com/AI-Hypercomputer/maxtext https://github.com/AI-Hypercomputer/maxtext ? It has a Jax llama implementation, and gets decent MFU on TPU and GPU (I've only tried it on NVidia GPU, not AMD).
- lostmsu 2y agoCould you share performance so we could compare?