Training Deep Networks with Data Parallelism in Jax
mishalaskin.com · 1,590 words · saved by 1 readers
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
- Pipeline-Parallelism: Distributed Training via Model Partitioningsiboehm.com
- How To Scale Your Modeljax-ml.github.io
- How to Parallelize a Transformer for Training — an explorable explanationezyang.github.io
- 5D parallelism in LLM training - gdymind's Bloggdymind.com
- Transformer Inference Arithmetic | kipply's blogkipp.ly
- Transformer Math 101 | EleutherAI Blogblog.eleuther.ai
- Async RL in Pure JAXdivyamakkar0.github.io
- Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
- Paradigms of Parallelism | Colossal-AIcolossalai.org
- How to Parallelize Deep Learning on GPUs Part 1/2: Data Parallelism - Tim Dettmerstimdettmers.com
- The Best GPUs for Deep Learning in 2023 — An In-depth Analysistimdettmers.com