Add gemma support - Githubissues

google / jetstream-pytorch

PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"

Apache License 2.0

32 stars 14 forks source link

Closed qihqi closed 4 months ago

qihqi commented 4 months ago

Changes include:

refactor engine to not use hardcoded sharding annotations; but instead read from a file.
Replace GemmaAttention class used in Gemma with layers.Attention: they seems to be doing the same thing.
added model_name arg to create_engine

Run interactive and run server both works with random weights.