{
"cells": [
{
"cell_type": "markdown",
"id": "8506306b",
"metadata": {},
"source": [
"# Tutorial 4: Up Next\n",
"\n",
"[](https://github.com/sshkhr/MinText/blob/main/docs/tutorials/3_Future_Work.ipynb)\n",
"[](https://colab.research.google.com/github/sshkhr/MinText/blob/main/docs/tutorials/3_Future_Work.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "92111afa",
"metadata": {},
"source": [
"We have started to scratch the surface of distributed training with our first three tutorials. Why is that?\n",
"\n",
"We have been using JAX's underlying GSPMD compiler to distribute our model across multiple devices. This is a powerful tool, but it abstracts away most of the complexities of distributed training such as implementing the collectives operations for us. For finer-grained control over distributed training, we would often need to implement our own sharded layers, collectives, training loops, and even backward passes. \n",
"\n",
"There are several libraries that provide a lot of underlying primitives for distributed training, the two most popular ones being [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) and [DeepSpeed](https://github.com/deepspeedai/DeepSpeed). If you want to continue working with JAX, you can look at the [MaxText](https://github.com/AI-Hypercomputer/maxtext) library, which was used for traiing the PaLM language model. A Pytorch native library in active development for distributed training is [Torch Titan](https://github.com/pytorch/torchtitan). \n",
"\n",
"While all of these libraries are great, they might be a bit overwhelming to get started with if you are new to distributed deep learning. At the end of this tutorial is a list of resources to help you get started with more accessible resources."
]
},
{
"cell_type": "markdown",
"id": "fa91009d",
"metadata": {},
"source": [
"In the rest of this notebook, we will get a sense of what lies ahead in the landscape of distributed training and where you can read more about it."
]
},
{
"cell_type": "markdown",
"id": "43299d3b",
"metadata": {},
"source": [
"## 1. Other Forms of Parallelism in Distributed Deep Learning"
]
},
{
"cell_type": "markdown",
"id": "303994dc",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "bedaf34e",
"metadata": {},
"source": [
" Image Source: [UvA Deep Learning Tutorials](https://uvadlc-notebooks.readthedocs.io/) "
]
},
{
"cell_type": "markdown",
"id": "a3f2902b",
"metadata": {},
"source": [
"### 1.1 Pipeline Parallelism"
]
},
{
"cell_type": "markdown",
"id": "cbcdb488",
"metadata": {},
"source": [
"Data and Tensor parallelism can be grouped together as **intra-layer parallelism**, where each layer is sharded across different devices. In contrast, pipeline parallelism is a form of **inter-layer parallelism**. In this approach, the model is divided into multiple stages (several groups of layers), and each stage is assigned to a different device. The input data flows through these stages sequentially, and each accelerator handles one stage (figure below). It effectively reduces memory usage of each accelerator.\n",
"\n",
"Both tensor and pipeline parallelism involve splitting the model across multiple devices, and are grouped together as **model parallelism**."
]
},
{
"cell_type": "markdown",
"id": "92ee47ef",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "dfdcd157",
"metadata": {},
"source": [
" Image Source: [Google GPipe](https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html) "
]
},
{
"cell_type": "markdown",
"id": "3404cae5",
"metadata": {},
"source": [
"Pipeline parallelism is very common for distirbuted training on GPUs, but not used as much on TPUs."
]
},
{
"cell_type": "markdown",
"id": "71d2481f",
"metadata": {},
"source": [
"Training with a combination of data, tensor, and pipeline parallelism is often referred to as 3D parallelism."
]
},
{
"cell_type": "markdown",
"id": "126effd2",
"metadata": {},
"source": [
"### 1.2 Sequence Parallelism"
]
},
{
"cell_type": "markdown",
"id": "71592601",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "7cdca5e2",
"metadata": {},
"source": [
" Image Source: [Reducing Activation Recomputation in Large Transformer Models](https://proceedings.mlsys.org/paper_files/paper/2023/hash/80083951326cf5b35e5100260d64ed81-Abstract-mlsys2023.html) "
]
},
{
"cell_type": "markdown",
"id": "873f9bc6",
"metadata": {},
"source": [
"Sequence parallelism was introduced with tensor parallelism by NVIDIA. When we covered tensor parallelism in the last tutorial, we only applied it to the MLP and Attention layers. But it cannot be applied to other layers which involve access to feature level statistics, such as the LayerNor layer. Sequence parallelism is a technique that allows us to apply tensor parallelism to these layers as well. These operations are **independent along the sequence dimension** (i.e. along token dimension), so they can be parallelized along the sequence dimension."
]
},
{
"cell_type": "markdown",
"id": "ddbaf12f",
"metadata": {},
"source": [
"### 1.3 Context Parallelism"
]
},
{
"cell_type": "markdown",
"id": "1afab08c",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "45a25f82",
"metadata": {},
"source": [
" Image Source: [xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism](https://arxiv.org/abs/2411.01738) "
]
},
{
"cell_type": "markdown",
"id": "a7e70bef",
"metadata": {},
"source": [
"Context parallelism can be thought of as a form of sequence parallelism, but it is applied independently of tensor parallelism (unlike sequence parallelism).It parallelizes training/inference prefill on a long sequence by partitioning the sequence into multiple subsets of tokens which are distributed across devices. To maintain the context and relationship between tokens, each device needs to communicate to compute attention exactly the same as local attention computation."
]
},
{
"cell_type": "markdown",
"id": "177647ad",
"metadata": {},
"source": [
"### 1.4 Expert Parallelism"
]
},
{
"cell_type": "markdown",
"id": "51ed8056",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "markdown",
"id": "2f9652ff",
"metadata": {},
"source": [
"In expert parallelism, the weights of each expert in a Mixture of Experts (MoE) model are distributed across different GPUs. This is still a very new area of research, and overlapping communication and computation for MoEs is still an open problem."
]
},
{
"cell_type": "markdown",
"id": "e717f6af",
"metadata": {},
"source": [
"## 2. What to Read Next"
]
},
{
"cell_type": "markdown",
"id": "7024e350",
"metadata": {},
"source": [
"These are some accessible resources to get you started with distributed training and inference: "
]
},
{
"cell_type": "markdown",
"id": "6313ef58",
"metadata": {},
"source": [
"- HuggingFace [Nanotron](https://github.com/huggingface/nanotron) Pytorch library for minimalistic large language model 3D-parallelism training, and [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - a comprehensive guide to efficiently scale LLM training with Nanotron\n",
"- [How to Scale Your Model](https://jax-ml.github.io/scaling-book/) book and [High Performance LLMs Course](https://github.com/rwitten/HighPerfLLMs2024) for an equivalent in the JAX ecosystem\n",
"- [Reference] [Awesome Distributed Machine Learning Systems](https://github.com/Shenggan/awesome-distributed-ml) - a curated list of resources on distributed machine learning systems."
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 5
}