Purpose
Mesh Transformer JAX is a model-parallel Transformer library built using JAX and Haiku, utilizing JAX's xmap / pjit operators for model parallelism of transformers, suitable for large-scale model training and inference on TPUs. The library also provides the pretrained model GPT-J-6B along with accompanying weights, evaluations, and usage guidance.
Main Capabilities
- Implements model parallelism for Transformers. Its parallelism scheme is similar to the original Megatron-LM and is highly efficient due to the TPU's high-speed 2D mesh network.
- Provides an experimental model version that implements ZeRO-style sharding.
- The library is designed to scale to approximately 40 billion parameters on TPUv3; beyond this scale, other parallelism strategies should be considered, such as GPT-NeoX or DeepSpeed.
- One future research direction is to integrate this codebase with swarm-jax to achieve further scaling through pipeline parallelism.
Pretrained Models
Provides GPT-J-6B: a 6-billion-parameter autoregressive text generation model trained on The Pile.
- Download slim weights (bf16 weights only, for inference, 9GB)
- Download full weights (including optimizer parameters, 61GB)
- Partially trained checkpoints are also provided
Model architecture details: The model consists of 28 layers, with a model dimension of 4096 and a feedforward dimension of 16384. The model dimension is split into 16 attention heads, each with a head dimension of 256. Rotary position embeddings (RoPE) are applied to 64 dimensions of each head. The model was trained using a tokenizer with a vocabulary size of 50257, using the same BPE set as GPT-2/GPT-3. Each layer consists of a feedforward block and a self-attention block. The license for the GPT-J-6B weights is the Apache License 2.0.
Usage
Most scripts in the repository are designed to run on TPUs. Under the TPU-VM architecture, the TPU is a virtual machine capable of running arbitrary code. The design flow of most scripts is: start a TPU, log in via SSH to set up dependencies and copy the code from the local directory over, then start a connectableRay worker invoked by RPC.
TPUVM is responsible for running model training steps and evaluation, checkpoint saving and loading, while the driver-side Python program is responsible for data loading and overall orchestration (for example, when to save checkpoints, etc.). This means that most scripts (such as train.py, eval_harness.py, etc.) are expected to run on a GCE virtual machine in the same region as the TPU, to minimize RPC latency and data transfer costs. Other scripts (usually scripts without the --tpu argument, such as device_sample.py, device_serve.py, or device_train.py) are expected to run directly on the TPUVM. The device_* scripts can only run on v3-8 and cannot run on larger pods.
In addition, the repository provides an example (resharding_example.py) showing how to convert a provided checkpoint (GPT-J-6B as 8 shards) to a smaller number of shards, for example when running on GPU.
Fine-tuning: Run device_train.py on a TPU VM to fine-tune the model. With TPU v3-8, fine-tuning speed is about 5000 tokens per second, which should be sufficient for small and medium-sized datasets. The repository provides a step-by-step guide to explain the detailed fine-tuning process.
JAX version dependency: This library has specific requirements for the JAX version. Specifically, using v1 models (including GPT-J 6B) requires jax==0.2.12, which in turn depends on jaxlib==0.1.68; if this is not satisfied, incomprehensible xmap errors will occur. To use the v2 model code (no public weights), the latest JAX version can be used.
Applicable Scenarios
Suitable for research and engineering scenarios that need to train or infer large-scale Transformers on TPU in a model-parallel manner, especially for using or fine-tuning the GPT-J-6B model. It is also suitable for studying large-scale model implementations by referring to its parallel scheme (similar to Megatron-LM and ZeRO-style sharding).
Limitations
- This library is designed to scale to about 40 billion parameters on TPUv3; beyond that scale, other parallel strategies are recommended.
- Most scripts are designed to run on TPU, and most scripts are expected to run on a GCE virtual machine in the same region as the TPU; the device_* scriptscan only run on v3-8 and cannot run on larger pods.
- Using v1 models (including GPT-J 6B) requires jax==0.2.12 and jaxlib==0.1.68, otherwise xmap errors will occur.
- This library does not provide information related to fees, pricing, etc.
