liveBook · Manning

EasyLM

As explained in Deep Learning with JAX

A JAX/Flax library for pretraining, fine-tuning, evaluating, and serving large language models.

Need to train or serve a large language model without writing distributed-training plumbing? EasyLM is a JAX/Flax library that covers pretraining, fine-tuning, evaluation, and serving. It hides much of model and data parallelism while keeping core training and inference details available for customization. section 12.1.2

EasyLM is useful when you want to:

  • Scale training to hundreds of TPU or GPU accelerators through JAX pjit, without writing complicated distributed-training code. section 12.1.2
  • Customize LLM training and inference while retaining access to their essential implementation details. section 12.1.2
  • Move from training to deployment using provided pretraining and serving scripts. section 12.1.2

Concrete scale targets

The distribution includes an example for pretraining a 7B LLaMA model on a TPU v4-512 pod and another for serving a pretrained LLaMA model on a GPU machine or single TPU v3-8 VM. section 12.1.2

How Deep Learning with JAX covers it

How deep it goes
Hides distributed-parallelism complexity but exposes core training and inference details for customization.
Examples
Uses 7B LLaMA pretraining on TPU v4-512 and pretrained-LLaMA serving on GPU or TPU v3-8 examples.
What you can do after
Apply it via EasyLM to customize and scale LLM training or inference across TPU and GPU accelerators.
How it connects
Connects JAX, Flax, pjit, distributed model parallelism, data parallelism, LLM training, evaluation, and serving.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.