Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Minimal library to train LLMs on TPU in JAX with pjit().
| Date | Stars |
|---|---|
| 2026-07-31 | 299 |
| 2026-08-02 | 299 |
| 2026-08-12 | 300 |
| 2026-09-14 | 300 |
| 2026-09-15 | 300 |
| 2026-09-18 | 300 |
| 2026-09-20 | 300 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# Jaxformer
JAX library for training of large language models with data and model parallelism based on the pjit() operator on TPU-v3/v4.
## Citation
Please cite:
```bibtex
@misc{Jaxformer,
title={Jaxformer: A minimal library for training LLMs on TPU},
author={Nijkamp, Erik},
howpublished = {\url{https://github.com/salesforce/jaxformer}},
year={2022}
}
```
Acknowledgments: Ben Wang, James Bradbury, Zak Stone, Bo Pang.
## Models
* CodeGen ([Paper](https://arxiv.org/abs/2203.13474)) ([Code](https://github.com/salesforce/codegen))
* ProGen2 ([Paper](https://arxiv.org/abs/2206.13517)) ([Code](https://github.com/salesforce/progen/tree/main/progen2))
### CodeGen
`350M`
```
gs://sfr-codegen-research/checkpoints/codegen-350M-nl/350000
gs://sfr-codegen-research/checkpoints/codegen-350M-multi/150000
gs://sfr-codegen-research/checkpoints/codegen-350M-mono/150000
```
`2B`
```
gs://sfr-codegen-research/checkpoints/codegen-2B-nl/350000
gs://sfr-codegen-research/checkpoints/codegen-2B-multi/150000
gs://sfr-codegen-research/checkpoints/codegen-2B-mono/100000
```
`6B`
```
gs://sfr-codegen-research/checkpoints/codegen-6B-nl/350000
gs://sfr-codegen-research/checkpoints/codegen-6B-multi/100000
gs://sfr-codegen-research/checkpoints/codegen-6B-mono/140000
```
## Sanity TPU
```sh
import jax
jax.devices()
device_count = jax.device_count()
local_device_count = jax.local_device_count()
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)
print('global device count:', jax.device_count())
print('local device count:', jax.local_device_count())
print('pmap result:', r)
gcloud compute tpus tpu-vm ssh [email protected] --zone=us-east1-d --internal-ip --worker=all --command="pip install 'jax[tpu]==0.3.16' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
gcloud compute tpus tpu-vm scp test.py [email protected]:/home/erik.nijkamp/ --zone=us-east1-d --internal-ip --worker=all
gcloud compute tpus tpu-vm ssh [email protected] --zone=us-east1-d --internal-ip --worker=all --command="python3 /home/erik.nijkamp/test.py"
```
## Training
### Mode 1: CPU local
```sh
brew install [email protected]
apt install --yes python3.9 python3.9-venv
git clone https://<username>:<secret>@github.com/salesforce/jaxformer.git/
cd jaxformer
python3.9 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip setuptools
pip install -r requirements.txt
python3 -m jaxformer.train --config config/debug_cpu.json
```
### Mode 2: TPU local
```sh
gcloud compute tpus list --zone=europe-west4-a
gcloud compute tpus tpu-vm delete sfr-erik.nijkamp-tpu-v3-8-europe-west4-d-1 --zone=europe-west4-a --quiet
gcloud compute tpus tpu-vm create sfr-erik.nijkamp-tpu-v3-8-europe-west4-d-1 --zone=europe-west4-a --accelerator-type=v3-8 --version=v2-alpha
gcloud compute tpus tpu-vm ssh sfr-erik.nijkamp-tpu-v3-8-europe-west4-d-1 --zone=europe-west4-a --project <project> --worker 0
export GOOGLE_APPLICATION_CREDENTIALS=~/.config/gcloud/legacy_credentials/<username>/adc.json
export GCLOUD_PROJECT=<project>
git clone https://<username>:<secret>@github.com/salesforce/jaxformer.git/
cd jaxformer
./jaxformer/env/env_tpu_v3.sh
pip install -r requirements.txt
source .venv/bin/activate
python3
import jax
jax.devices()
quit()
python3 -m jaxformer.train --config config/debug_tpu_v3_8.json
```
### Mode 3: TPU remote
```sh
gcloud beta compute --project=<project> instances create sfr-<username>-cpu-small-us-east1-d-1 --zone=us-east1-d --machine-type=e2-standard-4 --network-tier=PREMIUM --maintenance-policy=MIGRATE --service-account=<account> --scopes=https://www.googleapis.com/auth/devstorage.read_only,https://www.googleapis.com/auth/logging.write,https://www.googleapis.com/auth/monitoring.write,https://www.googleapis.com/auth/servicecontrol,https://www.googleapis.com/auth/service.management.readonly,httpExcerpt of 8,396 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:134a54066f9e786f, llm:description: 'Minimal library to train LLMs on TPU in JAX with pjit()'
matched fp:134a54066f9e786f, llm:description: 'Minimal library to train LLMs on TPU in JAX with pjit()'
matched fp:134a54066f9e786f, llm:description: 'Minimal library to train LLMs on TPU in JAX with pjit()'