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.
forked from jax-ml/jax-llm-examples
- Notifications
You must be signed in to change notification settings - Fork 0
Minimal yet performant LLM examples in pure JAX
License
martin-marek/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
Resources
License
Contributing
Stars
Watchers
Forks
Releases
No releases published
Packages 0
No packages published
Languages
- Python 98.6%
- Shell 1.4%