Skip to content
forked from apple/axlearn

An Extensible Deep Learning Library

License

Notifications You must be signed in to change notification settings

Ethanlm/axlearn

 
 

Repository files navigation

The AXLearn Library for Deep Learning

This library is under active development and the API is subject to change.

Table of Contents

Section Description
Introduction What is AXLearn?
Getting Started Getting up and running with AXLearn.
Concepts Core concepts and design principles.
CLI User Guide How to use the CLI.
Infrastructure Core infrastructure components.

Introduction

AXLearn is a library built on top of JAX and XLA to support the development of large-scale deep learning models.

AXLearn takes an object-oriented approach to the software engineering challenges that arise from building, iterating, and maintaining models. The configuration system of the library lets users compose models from reusable building blocks and integrate with other libraries such as Flax and Hugging Face transformers.

AXLearn is built to scale. It supports the training of models with up to hundreds of billions of parameters across thousands of accelerators at high utilization. It is also designed to run on public clouds and provides tools to deploy and manage jobs and data. Built on top of GSPMD, AXLearn adopts a global computation paradigm to allow users to describe computation on a virtual global computer rather than on a per-accelerator basis.

AXLearn supports a wide range of applications, including natural language processing, computer vision, and speech recognition and contains baseline configurations for training state-of-the-art models.

Please see Concepts for more details on the core components and design of AXLearn, or Getting Started if you want to get your hands dirty.

About

An Extensible Deep Learning Library

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 96.8%
  • Jupyter Notebook 2.9%
  • Other 0.3%