.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated/examples/distributed_tensors.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_distributed_tensors.py: Distributed Tensors in Monarch ------------------------------ .. GENERATED FROM PYTHON SOURCE LINES 11-20 .. code-block:: default import monarch import torch import torch.nn as nn from monarch.actor import this_host torch.set_default_device("cuda") .. GENERATED FROM PYTHON SOURCE LINES 21-25 Meshes ------ All computation is done on a 'mesh' of devices. Here we create a mesh composed of the machine running the notebook: .. GENERATED FROM PYTHON SOURCE LINES 25-29 .. code-block:: default mesh = this_host().spawn_procs({"gpu": 8}) print(mesh.to_table()) .. GENERATED FROM PYTHON SOURCE LINES 30-31 Without a mesh active, torch runs locally. .. GENERATED FROM PYTHON SOURCE LINES 31-34 .. code-block:: default torch.rand(3, 4) .. GENERATED FROM PYTHON SOURCE LINES 35-36 Once active, torch runs on every device in the mesh. .. GENERATED FROM PYTHON SOURCE LINES 36-42 .. code-block:: default with mesh.activate(): t = torch.rand(3, 4, device="cuda") t .. GENERATED FROM PYTHON SOURCE LINES 43-44 Inspect moves rank0's copy of t to the notebook for debugging. .. GENERATED FROM PYTHON SOURCE LINES 44-49 .. code-block:: default monarch.inspect(t) monarch.show(t) .. GENERATED FROM PYTHON SOURCE LINES 50-51 Providing coordinates lets us inspect other ranks copies. .. GENERATED FROM PYTHON SOURCE LINES 51-54 .. code-block:: default monarch.show(t, gpu=1) .. GENERATED FROM PYTHON SOURCE LINES 55-60 Tensor Commands --------------- Any command done on the controller, such as multiplying these tensors, performs that action to all of the tensors in the collection. .. GENERATED FROM PYTHON SOURCE LINES 60-65 .. code-block:: default with mesh.activate(): obj = t @ t.T monarch.show(obj) .. GENERATED FROM PYTHON SOURCE LINES 66-68 If a command fails, the workers stay alive and can execute future commands: .. GENERATED FROM PYTHON SOURCE LINES 68-84 .. code-block:: default try: with mesh.activate(): # too big big_w = torch.rand(4, 1024 * 1024 * 1024 * 1024 * 8, device="cuda") v = t @ big_w monarch.show(v) except Exception: import traceback traceback.print_exc() del big_w print("RECOVERED!") .. GENERATED FROM PYTHON SOURCE LINES 85-87 Since monarch recovers from errors, you can search for what works: .. GENERATED FROM PYTHON SOURCE LINES 87-100 .. code-block:: default N = 1 while True: try: with mesh.activate(): batch = torch.rand(N, 1024 * 1024 * 1024, device="cuda") monarch.inspect(batch.sum()) N = 2 * N print(f"at least 2**{N} elements work") except Exception: print(f"max is 2**{N} elements") break .. GENERATED FROM PYTHON SOURCE LINES 101-106 Collectives ----------- Each machine has its own copy of the tensor, similar to torch.distributed. To compute across tensors in the mesh, we use special communication operators, analogous to collectives. .. GENERATED FROM PYTHON SOURCE LINES 106-117 .. code-block:: default with mesh.activate(): a = torch.rand(3, 4, device="cuda") r = a.reduce("gpu", "sum") monarch.show(a, gpu=0) # try monarch.show(a, gpu=1) # try monarch.show(r, gpu=0) # try monarch.show(r, gpu=1) # try .. GENERATED FROM PYTHON SOURCE LINES 118-122 Remote GPUs --------- We can also connect to remote GPUs reserved from some scheduler .. GENERATED FROM PYTHON SOURCE LINES 122-133 .. code-block:: default # NYI: schedule public API based on config, just fake it locally remote_mesh = this_host().spawn_procs({"host": 4, "gpu": 4}) print(remote_mesh.to_table()) with remote_mesh.activate(): eg = torch.rand(3, 4, device="cuda") rgpu = eg.reduce("gpu", "sum") rhost = eg.reduce("host", "sum") .. GENERATED FROM PYTHON SOURCE LINES 134-138 Device Mesh Dimensions ---------------------- Meshes can be renamed and reshaped to fit the parallelism desired. .. GENERATED FROM PYTHON SOURCE LINES 138-145 .. code-block:: default mesh_2d_parallel = remote_mesh.rename(host="dp", gpu="tp") print(mesh_2d_parallel.to_table()) mesh_3d_parallel = remote_mesh.split(host=("dp", "pp"), gpu=("tp",), pp=2) print(mesh_3d_parallel.to_table()) .. GENERATED FROM PYTHON SOURCE LINES 146-151 Pipelining ---------- Pipelining is accomplished by slicing the mesh, and copying tensors from one mesh to another. .. GENERATED FROM PYTHON SOURCE LINES 151-157 .. code-block:: default pipeline_mesh = remote_mesh.rename(host="pp") meshes = [pipeline_mesh.slice(pp=i) for i in range(pipeline_mesh.size("pp"))] print(meshes[0].to_table()) .. GENERATED FROM PYTHON SOURCE LINES 158-159 Initialize a model across multiple meshes .. GENERATED FROM PYTHON SOURCE LINES 159-186 .. code-block:: default layers_per_stage = 2 stages = [] for stage_mesh in meshes: with stage_mesh.activate(): layers = [] for _ in range(layers_per_stage): layers.extend([nn.Linear(4, 4), nn.ReLU()]) stages.append(nn.Sequential(*layers)) def forward_pipeline(x): with torch.no_grad(): for stage_mesh, stage in zip(meshes, stages): x = x.to_mesh(stage_mesh) with stage_mesh.activate(): x = stage(x) return x with meshes[0].activate(): input = torch.rand(3, 4, device="cuda") output = forward_pipeline(input) monarch.show(output) print(output.mesh.to_table()) .. GENERATED FROM PYTHON SOURCE LINES 187-198 DDP Example ----------- The next sections will use an example of writing DDP to illustrate a typical way to develop code in monarch. Let's interleave the backward pass with the gradient reductions and parameter updates. We use monarch.grad_generator to incrementally run the backward pass. It returns an iterator that computes the grad parameters one at a time. .. GENERATED FROM PYTHON SOURCE LINES 198-214 .. code-block:: default def train(model, input, target): loss = model(input, target) rparameters = list(reversed(list(model.parameters()))) grads = monarch.grad_generator(loss, rparameters) with torch.no_grad(): it = iter(zip(rparameters, grads)) todo = next(it, None) while todo is not None: param, grad = todo grad.reduce_("dp", "sum") todo = next(it, None) param += 0.01 * grad .. GENERATED FROM PYTHON SOURCE LINES 215-222 Simulation of DDP ----------------- We can use a simulator to check for expected behavior of code before running it for real. It is another kind of mesh, which simulates rather than computes results for real. .. GENERATED FROM PYTHON SOURCE LINES 222-251 .. code-block:: default class Net(nn.Module): def __init__(self): super().__init__() layers = [] for x in range(8): layers.append(nn.Linear(4, 4)) layers.append(nn.ReLU()) self.layers = nn.Sequential(*layers) def forward(self, input, target): output = self.layers(input) return torch.nn.functional.cross_entropy(output, target) def simulate(): simulator = monarch.Simulator(hosts=1, gpus=4, trace_mode="stream_only") mesh = simulator.mesh.rename(gpu="dp") with mesh.activate(): model = Net() train(model, torch.rand(3, 4), torch.full((3,), 1, dtype=torch.int64)) simulator.display() simulate() .. GENERATED FROM PYTHON SOURCE LINES 252-258 Overlapping Comms/Compute ------------------------- Commands on different devices run in parallel, but by default commands on a single device run sequentially. We introduce parallelism on a device via stream objects. .. GENERATED FROM PYTHON SOURCE LINES 258-262 .. code-block:: default main = monarch.get_active_stream() comms = monarch.Stream("comms") .. GENERATED FROM PYTHON SOURCE LINES 263-265 The ``main`` stream runs computation sequentially, while the ``comms`` stream runs communication (e.g. allreduce) in parallel on the same device. .. GENERATED FROM PYTHON SOURCE LINES 267-272 To use a tensor from one stream on another we borrow it. The borrow API ensures deterministic memory usage, and eliminates the race conditions in the torch.cuda.stream API. A borrow transfers a tensor from one stream to another. The original stream cannot use the tensor until ``borrow.drop()`` is called, ensuring no races. .. GENERATED FROM PYTHON SOURCE LINES 274-275 The DDP example again, but using multiple streams. .. GENERATED FROM PYTHON SOURCE LINES 275-302 .. code-block:: default def train(model, input, target): loss = model(input, target) rparameters = list(reversed(list(model.parameters()))) grads = monarch.grad_generator(loss, rparameters) with torch.no_grad(): # NEW: iter also produces the tensor borrowed # to the comm stream it = iter( (param, grad, *comms.borrow(grad, mutable=True)) for param, grad in zip(rparameters, grads) ) todo = next(it, None) while todo is not None: param, grad, comm_grad, borrow = todo # NEW: compute the reduce on the comm stream with comms.activate(): comm_grad.reduce_("dp", "sum") borrow.drop() todo = next(it, None) param += 0.01 * grad simulate() .. GENERATED FROM PYTHON SOURCE LINES 303-306 The simulation result showed the results did not overlap much due to where the borrow.drop was placed. The reduce on the comms stream completed before the next backward step started on the main stream. .. GENERATED FROM PYTHON SOURCE LINES 308-311 The goal is to overlap the reduce (comms stream) with the next backward step (main stream). We can achieve this by ending the borrow after the grad step but before we update the param. .. GENERATED FROM PYTHON SOURCE LINES 311-335 .. code-block:: default def train(model, input, target): loss = model(input, target) rparameters = list(reversed(list(model.parameters()))) grads = monarch.grad_generator(loss, rparameters) with torch.no_grad(): it = iter( (param, grad, *comms.borrow(grad, mutable=True)) for param, grad in zip(rparameters, grads) ) todo = next(it, None) while todo is not None: param, grad, comm_grad, borrow = todo with comms.activate(): comm_grad.reduce_("dp", "sum") todo = next(it, None) # NEW: delay the borrow as late as possible borrow.drop() param += 0.01 * grad simulate() .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.000 seconds) .. _sphx_glr_download_generated_examples_distributed_tensors.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: distributed_tensors.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: distributed_tensors.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_