# kingoflolz/mesh-transformer-jax

Model parallel transformers in JAX and Haiku

Repository: https://github.com/kingoflolz/mesh-transformer-jax
Canonical: https://ross.abutalabs.com/products/mesh-transformer-jax
Language: Python
License: Apache-2.0
License Family: permissive
Last push: 2023-01-21T00:09:29+00:00

## Health v2 (maintenance only)
Score: 32/100 (v2, computed 2026-09-03T02:20:16.233290+00:00)
- activity 0, release rhythm 35, longevity 100
- inputs: {"age_days": 1999, "days_push": 1321, "days_rel": null, "gap_med": null, "n_releases_24m": 0}
- flags: no_releases
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 6380, forks 877 (observed 2026-08-28T04:09:42.952710+00:00)

## What it is
A JAX/Haiku library implementing model-parallel training and inference of transformer models using xmap/pjit operators, similar to Megatron-LM's parallelism scheme, optimized for TPUs. It is the codebase behind GPT-J-6B, a 6 billion parameter autoregressive language model trained on The Pile.

## Use cases
- train large transformer models on TPUs with model parallelism
- run inference with GPT-J-6B
- fine-tune GPT-J-6B on custom datasets
- generate text with a 6B parameter open-source language model
- experiment with Megatron-style tensor parallelism in JAX

## When to choose
- you need to train or fine-tune transformers up to ~40B parameters on TPU hardware
- you want to run or fine-tune GPT-J-6B
- you prefer JAX/Haiku over PyTorch for large-scale model parallelism

## When to avoid
- you need to scale beyond ~40B parameters (use GPT-NeoX or DeepSpeed instead)
- you are working on GPUs rather than TPUs
- you need a actively developed or general-purpose LLM training framework

## Facets
- artifact type: library
- maturity: maintenance
- function: llm-training, llm-inference, deep-learning, machine-learning
- domain: large-language-models, deep-learning, machine-learning, gpu-computing
- platform: python, cloud
- tags: jax, haiku, model-parallelism, tpu, gpt-j, transformers, megatron-style-parallelism, gpu

## Member repositories
- kingoflolz/mesh-transformer-jax (main) score 32

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:09:42.952710+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-29T17:45:09.540793+00:00, confidence not recorded.
  - readme: https://github.com/kingoflolz/mesh-transformer-jax (fetched 2026-08-28T04:09:42.952710+00:00, sha 6fc2f66d18e0)
- Data as of 2026-08-30T08:39:29.467469+00:00.