Training Deep Networks with Data Parallelism in Jax

Misha Laskin personal website. Includes a blog and projects focused on artifical intelligence.

Training Deep Networks with Data Parallelism in Jax You , Sun Feb 19 2023 • large model engineering Back One of the main challenges in training large neural networks, whether they are LLMs or VLMs, is that they are too large to fit on a single GPU. To address this issue, their training can be parallelized across multiple GPUs. This means either parallelizing the data or model to distribute computation across several devices. In this post, we'll cover batch splitting, also known as data parallelism, and show how to use JAX's pmap function to parallelize computations across multiple devices. Par

related reading

• irhum.github.io - Tensor Parallelism with jax.pjitirhum.github.io
• How to Parallelize a Transformer for Training — an explorable explanationezyang.github.io
• Paradigms of Parallelism | Colossal-AIcolossalai.org
• How To Scale Your Modeljax-ml.github.io
• Pipeline-Parallelism: Distributed Training via Model Partitioningsiboehm.com
• Parallelism in Distributed Deep Learning · Better Tomorrow with Computer Scienceinsujang.github.io
• How to Parallelize Deep Learning on GPUs Part 1/2: Data Parallelism - Tim Dettmerstimdettmers.com
• 5D parallelism in LLM training - gdymind's Bloggdymind.com
• Parallelism methods · Hugging Facehuggingface.co
• Why You Should (or Shouldn't) be Using Google's JAX in 2023assemblyai.com
• How to think in JAX — JAX documentationdocs.jax.dev
• The Little Book of Deep Learningfleuret.org