AI-Hypercomputer/JetStream

Support on Huggingface transformers

Closed this issue · 2 comments

ImKeTT commented

Hi, great work! Do you have any future plans on supporting the Flax/Jax implemented Huggingface transformer models?

vipannalla commented

JetStream is agnostic to particular model frameworks and should work for all LLM models if the engine supports it. Currently, we have two reference engine implementations -- one for JAX models and one for Pytorch models.

Can you try/check if your model is supported in Maxtext repo?

  • 👍1
ImKeTT commented

Thanks for your reply, seems that we might need a lot more efforts if we want to use LLMs that are currently not supported in Maxtext.