Research Software
Scalax
Scalax is a collection of utilties for helping developers to easily scale up JAX based machine learning models. The main idea of scalax is pretty simple: users write model and training code for a single GPU/TPU, and rely on scalax to automatically scale it up to hundreds of GPUs/TPUs. This is made possible by the JAX jit compiler, and scalax provides a set of utilities to help the users obtain the sharding annotations required by the jit compiler. Because scalax wraps around the jit compiler, existing JAX code can be easily scaled up using scalax with minimal changes.
EasyLM
Easy to use model parallel large language models training and evaluation in JAX/Flax using pjit on cloud TPU pods, with support for popular language models such as GPT-J, OPT and Roberta.
MLXU
Machine Learning eXperiment Utilities: convenient utilities for running machine learning experiments, parsing experiment configurations and logging results.
Open Source Models
OpenLLaMA: An Open Reproduction of LLaMA
OpenLLaMA is a permissively licensed open source reproduction of Meta AI's LLaMA large language model. We provide PyTorch and JAX weights of our pre-trained OpenLLaMA 7B model, as well as evaluation results and comparison against the original LLaMA models. The OpenLLaMA model weights can serve as a drop in replacement for the original LLaMA in downstream applications.
