A collection (in progress) of example high-performance large language model implementations, written with JAX.
Current contents include:
For multi-host cluster setup and distributed training, see multi_host_README.md and the tpu_toolkit.sh script.
-
Notifications
You must be signed in to change notification settings - Fork 19
Minimal yet performant LLM examples in pure JAX
License
jax-ml/jax-llm-examples
Folders and files
| Name | Name | Last commit message | Last commit date | |
|---|---|---|---|---|
Repository files navigation
About
Minimal yet performant LLM examples in pure JAX
Topics
Resources
License
Contributing
Stars
Watchers
Forks
Releases
No releases published
Packages 0
No packages published