{ "cells": [ { "cell_type": "markdown", "id": "11ffd94a", "metadata": {}, "source": [ "# Getting Started with Distributed Data Parallel\n", "\n", "**Author**: [Shen Li](https://mrshenli.github.io/)\n", "\n", "**Edited by**: [Joe Zhu](https://github.com/gunandrose4u), [Chirag\n", "Pandya](https://github.com/c-p-i-o)\n", "\n", "Note\n", "\n", "View and edit this tutorial in\n", "[github](https://github.com/pytorch/tutorials/blob/main/intermediate_source/ddp_tutorial.rst).\n", "\n", "Prerequisites:\n", "\n", "- [PyTorch Distributed Overview](../beginner/dist_overview.html)\n", "- [DistributedDataParallel API\n", " documents](https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html)\n", "- [DistributedDataParallel\n", " notes](https://pytorch.org/docs/master/notes/ddp.html)\n", "\n", "[DistributedDataParallel](https://pytorch.org/docs/stable/nn.html#module-torch.nn.parallel)\n", "(DDP) is a powerful module in PyTorch that allows you to parallelize\n", "your model across multiple machines, making it perfect for large-scale\n", "deep learning applications. To use DDP, you'll need to spawn multiple\n", "processes and create a single instance of DDP per process.\n", "\n", "But how does it work? DDP uses collective communications from the\n", "[torch.distributed](https://pytorch.org/tutorials/intermediate/dist_tuto.html)\n", "package to synchronize gradients and buffers across all processes. This\n", "means that each process will have its own copy of the model, but\n", "they'll all work together to train the model as if it were on a single\n", "machine.\n", "\n", "To make this happen, DDP registers an autograd hook for each parameter\n", "in the model. When the backward pass is run, this hook fires and\n", "triggers gradient synchronization across all processes. This ensures\n", "that each process has the same gradients, which are then used to update\n", "the model.\n", "\n", "For more information on how DDP works and how to use it effectively, be\n", "sure to check out the [DDP design\n", "note](https://pytorch.org/docs/master/notes/ddp.html). With DDP, you can\n", "train your models faster and more efficiently than ever before!\n", "\n", "The recommended way to use DDP is to spawn one process for each model\n", "replica. The model replica can span multiple devices. DDP processes can\n", "be placed on the same machine or across machines. Note that GPU devices\n", "cannot be shared across DDP processes (i.e. one GPU for one DDP\n", "process).\n", "\n", "In this tutorial, we'll start with a basic DDP use case and then\n", "demonstrate more advanced use cases, including checkpointing models and\n", "combining DDP with model parallel.\n", "\n", "Note\n", "\n", "The code in this tutorial runs on an 8-GPU server, but it can be easily\n", "generalized to other environments.\n", "\n", "## Comparison between `DataParallel` and `DistributedDataParallel`\n", "\n", "Before we dive in, let's clarify why you would consider using\n", "`DistributedDataParallel` over `DataParallel`, despite its added\n", "complexity:\n", "\n", "- First, `DataParallel` is single-process, multi-threaded, but it only\n", " works on a single machine. In contrast, `DistributedDataParallel` is\n", " multi-process and supports both single- and multi- machine training.\n", " Due to GIL contention across threads, per-iteration replicated\n", " model, and additional overhead introduced by scattering inputs and\n", " gathering outputs, `DataParallel` is usually slower than\n", " `DistributedDataParallel` even on a single machine.\n", "- Recall from the [prior\n", " tutorial](https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html)\n", " that if your model is too large to fit on a single GPU, you must use\n", " **model parallel** to split it across multiple GPUs.\n", " `DistributedDataParallel` works with **model parallel**, while\n", " `DataParallel` does not at this time. When DDP is combined with\n", " model parallel, each DDP process would use model parallel, and all\n", " processes collectively would use data parallel.\n", "\n", "## Basic Use Case\n", "\n", "To create a DDP module, you must first set up process groups properly.\n", "More details can be found in [Writing Distributed Applications with\n", "PyTorch](https://pytorch.org/tutorials/intermediate/dist_tuto.html).\n", "\n", "``` python\n", "import os\n", "import sys\n", "import tempfile\n", "import torch\n", "import torch.distributed as dist\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "import torch.multiprocessing as mp\n", "\n", "from torch.nn.parallel import DistributedDataParallel as DDP\n", "\n", "# On Windows platform, the torch.distributed package only\n", "# supports Gloo backend, FileStore and TcpStore.\n", "# For FileStore, set init_method parameter in init_process_group\n", "# to a local file. Example as follow:\n", "# init_method=\"file:///f:/libtmp/some_file\"\n", "# dist.init_process_group(\n", "# \"gloo\",\n", "# rank=rank,\n", "# init_method=init_method,\n", "# world_size=world_size)\n", "# For TcpStore, same way as on Linux.\n", "\n", "def setup(rank, world_size):\n", " os.environ['MASTER_ADDR'] = 'localhost'\n", " os.environ['MASTER_PORT'] = '12355'\n", "\n", " # initialize the process group\n", " dist.init_process_group(\"gloo\", rank=rank, world_size=world_size)\n", "\n", "def cleanup():\n", " dist.destroy_process_group()\n", "```\n", "\n", "Now, let's create a toy module, wrap it with DDP, and feed it some\n", "dummy input data. Please note, as DDP broadcasts model states from rank\n", "0 process to all other processes in the DDP constructor, you do not need\n", "to worry about different DDP processes starting from different initial\n", "model parameter values.\n", "\n", "``` python\n", "class ToyModel(nn.Module):\n", " def __init__(self):\n", " super(ToyModel, self).__init__()\n", " self.net1 = nn.Linear(10, 10)\n", " self.relu = nn.ReLU()\n", " self.net2 = nn.Linear(10, 5)\n", "\n", " def forward(self, x):\n", " return self.net2(self.relu(self.net1(x)))\n", "\n", "\n", "def demo_basic(rank, world_size):\n", " print(f\"Running basic DDP example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " # create model and move it to GPU with id rank\n", " model = ToyModel().to(rank)\n", " ddp_model = DDP(model, device_ids=[rank])\n", "\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(rank)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " cleanup()\n", " print(f\"Finished running basic DDP example on rank {rank}.\")\n", "\n", "\n", "def run_demo(demo_fn, world_size):\n", " mp.spawn(demo_fn,\n", " args=(world_size,),\n", " nprocs=world_size,\n", " join=True)\n", "```\n", "\n", "As you can see, DDP wraps lower-level distributed communication details\n", "and provides a clean API as if it were a local model. Gradient\n", "synchronization communications take place during the backward pass and\n", "overlap with the backward computation. When the `backward()` returns,\n", "`param.grad` already contains the synchronized gradient tensor. For\n", "basic use cases, DDP only requires a few more lines of code to set up\n", "the process group. When applying DDP to more advanced use cases, some\n", "caveats require caution.\n", "\n", "## Skewed Processing Speeds\n", "\n", "In DDP, the constructor, the forward pass, and the backward pass are\n", "distributed synchronization points. Different processes are expected to\n", "launch the same number of synchronizations and reach these\n", "synchronization points in the same order and enter each synchronization\n", "point at roughly the same time. Otherwise, fast processes might arrive\n", "early and timeout while waiting for stragglers. Hence, users are\n", "responsible for balancing workload distributions across processes.\n", "Sometimes, skewed processing speeds are inevitable due to, e.g., network\n", "delays, resource contentions, or unpredictable workload spikes. To avoid\n", "timeouts in these situations, make sure that you pass a sufficiently\n", "large `timeout` value when calling\n", "[init\\_process\\_group](https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group).\n", "\n", "## Save and Load Checkpoints\n", "\n", "It's common to use `torch.save` and `torch.load` to checkpoint modules\n", "during training and recover from checkpoints. See [SAVING AND LOADING\n", "MODELS](https://pytorch.org/tutorials/beginner/saving_loading_models.html)\n", "for more details. When using DDP, one optimization is to save the model\n", "in only one process and then load it on all processes, reducing write\n", "overhead. This works because all processes start from the same\n", "parameters and gradients are synchronized in backward passes, and hence\n", "optimizers should keep setting parameters to the same values. If you use\n", "this optimization (i.e. save on one process but restore on all), make\n", "sure no process starts loading before the saving is finished.\n", "Additionally, when loading the module, you need to provide an\n", "appropriate `map_location` argument to prevent processes from stepping\n", "into others' devices. If `map_location` is missing, `torch.load` will\n", "first load the module to CPU and then copy each parameter to where it\n", "was saved, which would result in all processes on the same machine using\n", "the same set of devices. For more advanced failure recovery and\n", "elasticity support, please refer to\n", "[TorchElastic](https://pytorch.org/elastic).\n", "\n", "``` python\n", "def demo_checkpoint(rank, world_size):\n", " print(f\"Running DDP checkpoint example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " model = ToyModel().to(rank)\n", " ddp_model = DDP(model, device_ids=[rank])\n", "\n", "\n", " CHECKPOINT_PATH = tempfile.gettempdir() + \"/model.checkpoint\"\n", " if rank == 0:\n", " # All processes should see same parameters as they all start from same\n", " # random parameters and gradients are synchronized in backward passes.\n", " # Therefore, saving it in one process is sufficient.\n", " torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)\n", "\n", " # Use a barrier() to make sure that process 1 loads the model after process\n", " # 0 saves it.\n", " dist.barrier()\n", " # configure map_location properly\n", " map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}\n", " ddp_model.load_state_dict(\n", " torch.load(CHECKPOINT_PATH, map_location=map_location, weights_only=True))\n", "\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(rank)\n", "\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " # Not necessary to use a dist.barrier() to guard the file deletion below\n", " # as the AllReduce ops in the backward pass of DDP already served as\n", " # a synchronization.\n", "\n", " if rank == 0:\n", " os.remove(CHECKPOINT_PATH)\n", "\n", " cleanup()\n", " print(f\"Finished running DDP checkpoint example on rank {rank}.\")\n", "```\n", "\n", "## Combining DDP with Model Parallelism\n", "\n", "DDP also works with multi-GPU models. DDP wrapping multi-GPU models is\n", "especially helpful when training large models with a huge amount of\n", "data.\n", "\n", "``` python\n", "class ToyMpModel(nn.Module):\n", " def __init__(self, dev0, dev1):\n", " super(ToyMpModel, self).__init__()\n", " self.dev0 = dev0\n", " self.dev1 = dev1\n", " self.net1 = torch.nn.Linear(10, 10).to(dev0)\n", " self.relu = torch.nn.ReLU()\n", " self.net2 = torch.nn.Linear(10, 5).to(dev1)\n", "\n", " def forward(self, x):\n", " x = x.to(self.dev0)\n", " x = self.relu(self.net1(x))\n", " x = x.to(self.dev1)\n", " return self.net2(x)\n", "```\n", "\n", "When passing a multi-GPU model to DDP, `device_ids` and `output_device`\n", "must NOT be set. Input and output data will be placed in proper devices\n", "by either the application or the model `forward()` method.\n", "\n", "``` python\n", "def demo_model_parallel(rank, world_size):\n", " print(f\"Running DDP with model parallel example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " # setup mp_model and devices for this process\n", " dev0 = rank * 2\n", " dev1 = rank * 2 + 1\n", " mp_model = ToyMpModel(dev0, dev1)\n", " ddp_mp_model = DDP(mp_model)\n", "\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " # outputs will be on dev1\n", " outputs = ddp_mp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(dev1)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " cleanup()\n", " print(f\"Finished running DDP with model parallel example on rank {rank}.\")\n", "\n", "\n", "if __name__ == \"__main__\":\n", " n_gpus = torch.cuda.device_count()\n", " assert n_gpus >= 2, f\"Requires at least 2 GPUs to run, but got {n_gpus}\"\n", " world_size = n_gpus\n", " run_demo(demo_basic, world_size)\n", " run_demo(demo_checkpoint, world_size)\n", " world_size = n_gpus//2\n", " run_demo(demo_model_parallel, world_size)\n", "```\n", "\n", "## Initialize DDP with torch.distributed.run/torchrun\n", "\n", "We can leverage PyTorch Elastic to simplify the DDP code and initialize\n", "the job more easily. Let's still use the Toymodel example and create a\n", "file named `elastic_ddp.py`.\n", "\n", "``` python\n", "import torch\n", "import torch.distributed as dist\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "\n", "from torch.nn.parallel import DistributedDataParallel as DDP\n", "\n", "class ToyModel(nn.Module):\n", " def __init__(self):\n", " super(ToyModel, self).__init__()\n", " self.net1 = nn.Linear(10, 10)\n", " self.relu = nn.ReLU()\n", " self.net2 = nn.Linear(10, 5)\n", "\n", " def forward(self, x):\n", " return self.net2(self.relu(self.net1(x)))\n", "\n", "\n", "def demo_basic():\n", " torch.cuda.set_device(int(os.environ[\"LOCAL_RANK\"]))\n", " dist.init_process_group(\"nccl\")\n", " rank = dist.get_rank()\n", " print(f\"Start running basic DDP example on rank {rank}.\")\n", " # create model and move it to GPU with id rank\n", " device_id = rank % torch.cuda.device_count()\n", " model = ToyModel().to(device_id)\n", " ddp_model = DDP(model, device_ids=[device_id])\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(device_id)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", " dist.destroy_process_group()\n", " print(f\"Finished running basic DDP example on rank {rank}.\")\n", "\n", "if __name__ == \"__main__\":\n", " demo_basic()\n", "```\n", "\n", "One can then run a [torch\n", "elastic/torchrun](https://pytorch.org/docs/stable/elastic/quickstart.html)\n", "command on all nodes to initialize the DDP job created above:\n", "\n", "``` bash\n", "torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 elastic_ddp.py\n", "```\n", "\n", "In the example above, we are running the DDP script on two hosts and we\n", "run with 8 processes on each host. That is, we are running this job on\n", "16 GPUs. Note that `$MASTER_ADDR` must be the same across all nodes.\n", "\n", "Here `torchrun` will launch 8 processes and invoke `elastic_ddp.py` on\n", "each process on the node it is launched on, but user also needs to apply\n", "cluster management tools like slurm to actually run this command on 2\n", "nodes.\n", "\n", "For example, on a SLURM enabled cluster, we can write a script to run\n", "the command above and set `MASTER_ADDR` as:\n", "\n", "``` bash\n", "export MASTER_ADDR=$(scontrol show hostname ${SLURM_NODELIST} | head -n 1)\n", "```\n", "\n", "Then we can just run this script using the SLURM command:\n", "`srun --nodes=2 ./torchrun_script.sh`.\n", "\n", "This is just an example; you can choose your own cluster scheduling\n", "tools to initiate the `torchrun` job.\n", "\n", "For more information about Elastic run, please see the [quick start\n", "document](https://pytorch.org/docs/stable/elastic/quickstart.html)." ] }, { "cell_type": "code", "execution_count": 8, "id": "439cdd8a", "metadata": {}, "outputs": [], "source": [ "!mkdir temp" ] }, { "cell_type": "code", "execution_count": 1, "id": "c0a5d920", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting temp/ddp_tutorial.py\n" ] } ], "source": [ "%%writefile temp/ddp_tutorial.py\n", "import os\n", "import sys\n", "import tempfile\n", "import torch\n", "import torch.distributed as dist\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "import torch.multiprocessing as mp\n", "\n", "from torch.nn.parallel import DistributedDataParallel as DDP\n", "\n", "\n", "\n", "def setup(rank, world_size):\n", " # os.environ['MASTER_ADDR'] = 'localhost'\n", " # os.environ['MASTER_PORT'] = '12355'\n", " print(os.environ['MASTER_ADDR'])\n", " print(os.environ['MASTER_PORT'])\n", "\n", " # initialize the process group\n", " dist.init_process_group(\"gloo\", rank=rank, world_size=world_size)\n", "\n", "def cleanup():\n", " dist.destroy_process_group()\n", "\n", "\n", "# ``` python\n", "class ToyModel(nn.Module):\n", " def __init__(self):\n", " super(ToyModel, self).__init__()\n", " self.net1 = nn.Linear(10, 10)\n", " self.relu = nn.ReLU()\n", " self.net2 = nn.Linear(10, 5)\n", "\n", " def forward(self, x):\n", " return self.net2(self.relu(self.net1(x)))\n", "\n", "\n", "def demo_basic(rank, world_size):\n", " print(f\"Running basic DDP example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " # create model and move it to GPU with id rank\n", " \n", " print(\"Start creating model\")\n", " model = ToyModel().to(rank)\n", " ddp_model = DDP(model, device_ids=[rank])\n", "\n", " print(\"Model created\")\n", " print(\"Start creating loss function\")\n", " loss_fn = nn.MSELoss()\n", " print(\"Loss function created\")\n", " print(\"Start creating optimizer\")\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(rank)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " cleanup()\n", " print(f\"Finished running basic DDP example on rank {rank}.\")\n", "\n", "\n", "def run_demo(demo_fn, world_size):\n", " mp.spawn(demo_fn,\n", " args=(world_size,),\n", " nprocs=world_size,\n", " join=True)\n", "\n", "# ``` python\n", "def demo_checkpoint(rank, world_size):\n", " print(f\"Running DDP checkpoint example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " model = ToyModel().to(rank)\n", " ddp_model = DDP(model, device_ids=[rank])\n", "\n", "\n", " CHECKPOINT_PATH = tempfile.gettempdir() + \"/model.checkpoint\"\n", " if rank == 0:\n", " # All processes should see same parameters as they all start from same\n", " # random parameters and gradients are synchronized in backward passes.\n", " # Therefore, saving it in one process is sufficient.\n", " torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)\n", "\n", " # Use a barrier() to make sure that process 1 loads the model after process\n", " # 0 saves it.\n", " dist.barrier()\n", " # configure map_location properly\n", " map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}\n", " ddp_model.load_state_dict(\n", " torch.load(CHECKPOINT_PATH, map_location=map_location, weights_only=True))\n", "\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(rank)\n", "\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " # Not necessary to use a dist.barrier() to guard the file deletion below\n", " # as the AllReduce ops in the backward pass of DDP already served as\n", " # a synchronization.\n", "\n", " if rank == 0:\n", " os.remove(CHECKPOINT_PATH)\n", "\n", " cleanup()\n", " print(f\"Finished running DDP checkpoint example on rank {rank}.\")\n", " \n", "\n", "## Combining DDP with Model Parallelism\n", "\n", "\n", "# ``` python\n", "class ToyMpModel(nn.Module):\n", " def __init__(self, dev0, dev1):\n", " super(ToyMpModel, self).__init__()\n", " self.dev0 = dev0\n", " self.dev1 = dev1\n", " self.net1 = torch.nn.Linear(10, 10).to(dev0)\n", " self.relu = torch.nn.ReLU()\n", " self.net2 = torch.nn.Linear(10, 5).to(dev1)\n", "\n", " def forward(self, x):\n", " x = x.to(self.dev0)\n", " x = self.relu(self.net1(x))\n", " x = x.to(self.dev1)\n", " return self.net2(x)\n", "\n", "\n", "# ``` python\n", "def demo_model_parallel(rank, world_size):\n", " print(f\"Running DDP with model parallel example on rank {rank}.\")\n", " setup(rank, world_size)\n", "\n", " # setup mp_model and devices for this process\n", " dev0 = rank * 2\n", " dev1 = rank * 2 + 1\n", " mp_model = ToyMpModel(dev0, dev1)\n", " ddp_mp_model = DDP(mp_model)\n", "\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " # outputs will be on dev1\n", " outputs = ddp_mp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(dev1)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", "\n", " cleanup()\n", " print(f\"Finished running DDP with model parallel example on rank {rank}.\")\n", "\n", "\n", "if __name__ == \"__main__\":\n", " n_gpus = torch.cuda.device_count()\n", " assert n_gpus >= 2, f\"Requires at least 2 GPUs to run, but got {n_gpus}\"\n", " world_size = n_gpus\n", " run_demo(demo_basic, world_size)\n", " # run_demo(demo_checkpoint, world_size)\n", " # world_size = n_gpus//2\n", " # run_demo(demo_model_parallel, world_size)" ] }, { "cell_type": "code", "execution_count": 2, "id": "13ffecef", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "4" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import torch\n", "torch.cuda.device_count()" ] }, { "cell_type": "code", "execution_count": 54, "id": "4e702428", "metadata": {}, "outputs": [], "source": [ "# !torchrun --help" ] }, { "cell_type": "code", "execution_count": 57, "id": "18fcc4ca", "metadata": {}, "outputs": [], "source": [ "# --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \\\n", "\n", "\n", "command = '''\n", "# export MASTER_ADDR=localhost\n", "# export MASTER_PORT=12355\n", "# NVIDIA_VISIBLE_DEVICES=0,1 \n", "torchrun --nnodes=1 --nproc_per_node=4 \\\n", "--rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=localhost:12355 \\\n", " temp/ddp_tutorial.py\n", "echo 1\n", "'''" ] }, { "cell_type": "code", "execution_count": null, "id": "5a2f3748", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[WARNING] Reference not found for 'Key \"edit\"' at chunk line 2 column 7\n", "\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Converting the file to markdown...\n", "Markdown file saved as: ddp_tutorial.md\n", "Converting the markdown content to a Jupyter notebook...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Collecting jupytext\n", " Using cached jupytext-1.16.6-py3-none-any.whl.metadata (13 kB)\n", "Collecting markdown-it-py>=1.0 (from jupytext)\n", " Downloading markdown_it_py-3.0.0-py3-none-any.whl.metadata (6.9 kB)\n", "Collecting mdit-py-plugins (from jupytext)\n", " Using cached mdit_py_plugins-0.4.2-py3-none-any.whl.metadata (2.8 kB)\n", "Requirement already satisfied: nbformat in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (5.10.4)\n", "Requirement already satisfied: packaging in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (24.1)\n", "Requirement already satisfied: pyyaml in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (6.0.2)\n", "Requirement already satisfied: tomli in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (2.2.1)\n", "Collecting mdurl~=0.1 (from markdown-it-py>=1.0->jupytext)\n", " Downloading mdurl-0.1.2-py3-none-any.whl.metadata (1.6 kB)\n", "Requirement already satisfied: fastjsonschema>=2.15 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (2.21.1)\n", "Requirement already satisfied: jsonschema>=2.6 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (4.23.0)\n", "Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (5.7.2)\n", "Requirement already satisfied: traitlets>=5.1 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (5.14.3)\n", "Requirement already satisfied: attrs>=22.2.0 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (24.2.0)\n", "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (2023.12.1)\n", "Requirement already satisfied: referencing>=0.28.4 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (0.35.1)\n", "Requirement already satisfied: rpds-py>=0.7.1 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (0.20.0)\n", "Requirement already satisfied: platformdirs>=2.5 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupyter-core!=5.0.*,>=4.12->nbformat->jupytext) (4.3.6)\n", "Using cached jupytext-1.16.6-py3-none-any.whl (154 kB)\n", "Downloading markdown_it_py-3.0.0-py3-none-any.whl (87 kB)\n", "Using cached mdit_py_plugins-0.4.2-py3-none-any.whl (55 kB)\n", "Downloading mdurl-0.1.2-py3-none-any.whl (10.0 kB)\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Installing collected packages: mdurl, markdown-it-py, mdit-py-plugins, jupytext\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable.It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Successfully installed jupytext-1.16.6 markdown-it-py-3.0.0 mdit-py-plugins-0.4.2 mdurl-0.1.2\n", "[jupytext] Reading ddp_tutorial.md in format md\n", "[jupytext] Writing ddp_tutorial.ipynb\n", "Converted ddp_tutorial.md to jupyter notebook.\n", "Conversion complete.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages/nbformat/__init__.py:96: MissingIDFieldWarning: Cell is missing an id field, this will become a hard error in future nbformat versions. You may want to use `normalize()` on your notebooks before validations (available since nbformat 5.1.4). Previous versions of nbformat are fixing this issue transparently, and will stop doing so in the future.\n", " validate(nb)\n" ] } ], "source": [ "def convert_markdown_to_jupyter_notebook(filename):\n", " import os\n", " try: \n", " import jupytext\n", " except:\n", " os.system(\"pip install jupytext\")\n", " os.system(f\"jupytext --to notebook {filename}\")\n", " print(f\"Converted {filename} to jupyter notebook.\")\n", "\n", "def download_convert_pipeline(url):\n", " import requests\n", " response = requests.get(url)\n", " \n", " # dump the content of the response to a file\n", " # 1. get the file name from the url\n", " file_name = url.split('/')[-1]\n", " # 2. write the content to the file\n", " with open(file_name, 'w') as f:\n", " f.write(response.text)\n", " \n", " # 3. Convert the file to markdown\n", " print(\"Converting the file to markdown...\")\n", " md_content = convert_rst_to_md(response.text)\n", " md_file_name = file_name.replace('.rst', '.md')\n", " save_md_file(md_content, md_file_name)\n", " print(f\"Markdown file saved as: {md_file_name}\")\n", " \n", " # 4. Convert the markdown content to a notebook\n", " print(\"Converting the markdown content to a Jupyter notebook...\")\n", " convert_markdown_to_jupyter_notebook(md_file_name)\n", " print(\"Conversion complete.\")\n", " \n", " \n", " if response.status_code == 200:\n", " return response.text\n", " else:\n", " raise Exception(f\"Failed to download file: {response.status_code}\")\n", " \n", "def main():\n", " url = \"https://raw.githubusercontent.com/pytorch/tutorials/main/intermediate_source/ddp_tutorial.rst\"\n", " download_convert_pipeline(url)\n", " \n", "if __name__ == \"__main__\":\n", " main()" ] }, { "cell_type": "code", "execution_count": null, "id": "cf0f8591", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "W1226 06:19:17.846000 367977 site-packages/torch/distributed/run.py:793] \n", "W1226 06:19:17.846000 367977 site-packages/torch/distributed/run.py:793] *****************************************\n", "W1226 06:19:17.846000 367977 site-packages/torch/distributed/run.py:793] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. \n", "W1226 06:19:17.846000 367977 site-packages/torch/distributed/run.py:793] *****************************************\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Running basic DDP example on rank 1.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 3.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 0.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 0.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 0.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 3.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 2.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 3.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 3.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 1.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 1.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 0.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 2.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 2.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 2.\n", "localhost\n", "12355\n", "Running basic DDP example on rank 1.\n", "localhost\n", "12355\n" ] } ], "source": [ "import os\n", "os.system(command)" ] }, { "cell_type": "code", "execution_count": null, "id": "5a018f18", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[WARNING] Reference not found for 'Key \"edit\"' at chunk line 2 column 7\n", "\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Converting the file to markdown...\n", "Markdown file saved as: ddp_tutorial.md\n", "Converting the markdown content to a Jupyter notebook...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Collecting jupytext\n", " Using cached jupytext-1.16.6-py3-none-any.whl.metadata (13 kB)\n", "Collecting markdown-it-py>=1.0 (from jupytext)\n", " Downloading markdown_it_py-3.0.0-py3-none-any.whl.metadata (6.9 kB)\n", "Collecting mdit-py-plugins (from jupytext)\n", " Using cached mdit_py_plugins-0.4.2-py3-none-any.whl.metadata (2.8 kB)\n", "Requirement already satisfied: nbformat in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (5.10.4)\n", "Requirement already satisfied: packaging in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (24.1)\n", "Requirement already satisfied: pyyaml in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (6.0.2)\n", "Requirement already satisfied: tomli in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupytext) (2.2.1)\n", "Collecting mdurl~=0.1 (from markdown-it-py>=1.0->jupytext)\n", " Downloading mdurl-0.1.2-py3-none-any.whl.metadata (1.6 kB)\n", "Requirement already satisfied: fastjsonschema>=2.15 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (2.21.1)\n", "Requirement already satisfied: jsonschema>=2.6 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (4.23.0)\n", "Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (5.7.2)\n", "Requirement already satisfied: traitlets>=5.1 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from nbformat->jupytext) (5.14.3)\n", "Requirement already satisfied: attrs>=22.2.0 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (24.2.0)\n", "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (2023.12.1)\n", "Requirement already satisfied: referencing>=0.28.4 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (0.35.1)\n", "Requirement already satisfied: rpds-py>=0.7.1 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jsonschema>=2.6->nbformat->jupytext) (0.20.0)\n", "Requirement already satisfied: platformdirs>=2.5 in /dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages (from jupyter-core!=5.0.*,>=4.12->nbformat->jupytext) (4.3.6)\n", "Using cached jupytext-1.16.6-py3-none-any.whl (154 kB)\n", "Downloading markdown_it_py-3.0.0-py3-none-any.whl (87 kB)\n", "Using cached mdit_py_plugins-0.4.2-py3-none-any.whl (55 kB)\n", "Downloading mdurl-0.1.2-py3-none-any.whl (10.0 kB)\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Installing collected packages: mdurl, markdown-it-py, mdit-py-plugins, jupytext\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[33mWARNING: Ignoring invalid distribution -vidia-cublas-cu12 (/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable.It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n", "\u001b[0m" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Successfully installed jupytext-1.16.6 markdown-it-py-3.0.0 mdit-py-plugins-0.4.2 mdurl-0.1.2\n", "[jupytext] Reading ddp_tutorial.md in format md\n", "[jupytext] Writing ddp_tutorial.ipynb\n", "Converted ddp_tutorial.md to jupyter notebook.\n", "Conversion complete.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/dscilab_dungvo/workspace/bin/envs/vllm/lib/python3.10/site-packages/nbformat/__init__.py:96: MissingIDFieldWarning: Cell is missing an id field, this will become a hard error in future nbformat versions. You may want to use `normalize()` on your notebooks before validations (available since nbformat 5.1.4). Previous versions of nbformat are fixing this issue transparently, and will stop doing so in the future.\n", " validate(nb)\n" ] } ], "source": [ "def convert_markdown_to_jupyter_notebook(filename):\n", " import os\n", " try: \n", " import jupytext\n", " except:\n", " os.system(\"pip install jupytext\")\n", " os.system(f\"jupytext --to notebook {filename}\")\n", " print(f\"Converted {filename} to jupyter notebook.\")\n", "\n", "def download_convert_pipeline(url):\n", " import requests\n", " response = requests.get(url)\n", " \n", " # dump the content of the response to a file\n", " # 1. get the file name from the url\n", " file_name = url.split('/')[-1]\n", " # 2. write the content to the file\n", " with open(file_name, 'w') as f:\n", " f.write(response.text)\n", " \n", " # 3. Convert the file to markdown\n", " print(\"Converting the file to markdown...\")\n", " md_content = convert_rst_to_md(response.text)\n", " md_file_name = file_name.replace('.rst', '.md')\n", " save_md_file(md_content, md_file_name)\n", " print(f\"Markdown file saved as: {md_file_name}\")\n", " \n", " # 4. Convert the markdown content to a notebook\n", " print(\"Converting the markdown content to a Jupyter notebook...\")\n", " convert_markdown_to_jupyter_notebook(md_file_name)\n", " print(\"Conversion complete.\")\n", " \n", " \n", " if response.status_code == 200:\n", " return response.text\n", " else:\n", " raise Exception(f\"Failed to download file: {response.status_code}\")\n", " \n", "def main():\n", " url = \"https://raw.githubusercontent.com/pytorch/tutorials/main/intermediate_source/ddp_tutorial.rst\"\n", " download_convert_pipeline(url)\n", " \n", "if __name__ == \"__main__\":\n", " main()" ] }, { "cell_type": "code", "execution_count": 12, "id": "83bffe64", "metadata": {}, "outputs": [], "source": [ "!export MASTER_ADDR=localhost\n" ] }, { "cell_type": "markdown", "id": "2f67bbc3", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "id": "bd207ab0", "metadata": {}, "outputs": [], "source": [ "\n", "\n", "# ``` python\n", "import torch\n", "import torch.distributed as dist\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "\n", "from torch.nn.parallel import DistributedDataParallel as DDP\n", "\n", "class ToyModel(nn.Module):\n", " def __init__(self):\n", " super(ToyModel, self).__init__()\n", " self.net1 = nn.Linear(10, 10)\n", " self.relu = nn.ReLU()\n", " self.net2 = nn.Linear(10, 5)\n", "\n", " def forward(self, x):\n", " return self.net2(self.relu(self.net1(x)))\n", "\n", "\n", "def demo_basic():\n", " import os\n", " torch.cuda.set_device(int(os.environ[\"LOCAL_RANK\"]))\n", " dist.init_process_group(\"nccl\")\n", " rank = dist.get_rank()\n", " print(f\"Start running basic DDP example on rank {rank}.\")\n", " # create model and move it to GPU with id rank\n", " device_id = rank % torch.cuda.device_count()\n", " model = ToyModel().to(device_id)\n", " ddp_model = DDP(model, device_ids=[device_id])\n", " loss_fn = nn.MSELoss()\n", " optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n", "\n", " optimizer.zero_grad()\n", " outputs = ddp_model(torch.randn(20, 10))\n", " labels = torch.randn(20, 5).to(device_id)\n", " loss_fn(outputs, labels).backward()\n", " optimizer.step()\n", " dist.destroy_process_group()\n", " print(f\"Finished running basic DDP example on rank {rank}.\")\n", "\n", "if __name__ == \"__main__\":\n", " demo_basic()\n", "\n", "\n", "\n", "# ``` bash\n", "!torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 elastic_ddp.py\n", "# ```\n", "\n", "\n", "# ``` bash\n", "!export MASTER_ADDR=$(scontrol show hostname ${SLURM_NODELIST} | head -n 1)\n", "# ```\n" ] } ], "metadata": { "kernelspec": { "display_name": "vllm", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.14" } }, "nbformat": 4, "nbformat_minor": 5 }