liveBook · Manning

Jaxformer

As explained in Deep Learning with JAX

A minimal JAX library for training large language models on TPUs using data and model parallelism.

Jaxformer is a minimal Salesforce library for training large language models in JAX on TPUs. It combines data parallelism and model parallelism through pjit(). section 12.1.2

Jaxformer fits when you want:

  • A focused JAX-based starting point for TPU language-model training. section 12.1.2
  • Both data and model parallelism rather than single-device training alone. section 12.1.2

How Deep Learning with JAX covers it

How deep it goes
Defines the library’s purpose and parallelism approach but does not detail model architectures, training steps, or implementation internals.
What you can do after
Evaluate/understand Jaxformer as a minimal option for JAX-based, TPU-oriented LLM training.
How it connects
Connects JAX and pjit() with data parallelism, model parallelism, and large-language-model training.
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.