{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# A look behind the curtain: some internals for the experts \n", "\n", "In this section, we'll go through some Heat-specific functionalities that simplify the implementation of a data-parallel application in Python. We'll demonstrate them on small arrays and 4 processes on a single cluster node, but the functionalities are indeed meant for a multi-node set up with huge arrays that cannot be processed on a single node." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Your IPython cluster should still be running. Let's check it out." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "4 engines found\n" ] } ], "source": [ "from ipyparallel import Client\n", "rc = Client(profile=\"default\")\n", "rc.ids\n", "\n", "if len(rc.ids) == 0:\n", " print(\"No engines found\")\n", "else:\n", " print(f\"{len(rc.ids)} engines found\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "If no engines are found, go back to the `0_setup` notebook for instructions." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We already mentioned that the DNDarray object is \"MPI-aware\". Each DNDarray is associated to an MPI communicator, it is aware of the number of processes in the communicator, and it knows the rank of the process that owns it. \n", "\n", "We will use the `%%px` magic in every cell that executes MPI code." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[1:11]: \u001b[0m" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "import torch\nimport heat as ht\n\na = ht.random.randn(7,4,3, split=0)\na.comm\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 11, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:21.227261Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:11]: \u001b[0m" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "import torch\nimport heat as ht\n\na = ht.random.randn(7,4,3, split=0)\na.comm\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 11, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:21.227047Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[2:11]: \u001b[0m" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "import torch\nimport heat as ht\n\na = ht.random.randn(7,4,3, split=0)\na.comm\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 11, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:21.227323Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[3:11]: \u001b[0m" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "import torch\nimport heat as ht\n\na = ht.random.randn(7,4,3, split=0)\na.comm\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 11, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:21.227372Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "import torch\n", "import heat as ht\n", "\n", "a = ht.random.randn(7,4,3, split=0)\n", "a.comm" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:2] a is distributed over 4 processes\n", "a is a distributed 3-dimensional array with global shape (7, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] a is distributed over 4 processes\n", "a is a distributed 3-dimensional array with global shape (7, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:0] a is distributed over 4 processes\n", "a is a distributed 3-dimensional array with global shape (7, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] a is distributed over 4 processes\n", "a is a distributed 3-dimensional array with global shape (7, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "# MPI size = total number of processes \n", "size = a.comm.size\n", "\n", "print(f\"a is distributed over {size} processes\")\n", "print(f\"a is a distributed {a.ndim}-dimensional array with global shape {a.shape}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:0] Rank 0 holds a slice of DNDarray 'a' with local shape (2, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:2] Rank 2 holds a slice of DNDarray 'a' with local shape (2, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] Rank 1 holds a slice of DNDarray 'a' with local shape (2, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] Rank 3 holds a slice of DNDarray 'a' with local shape (1, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "# MPI rank = rank of each process\n", "rank = a.comm.rank\n", "# Local shape = shape of the data on each process\n", "local_shape = a.lshape\n", "print(f\"Rank {rank} holds a slice of DNDarray 'a' with local shape {local_shape}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Distribution map\n", "\n", "On many occasions, when building a memory-distributed pipeline, it will be convenient for each rank to have information on which slice of the distributed array each rank holds. \n", "\n", "The `lshape_map` attribute of a DNDarray gathers (or, if possible, calculates) this info from all processes and stores it as metadata of the DNDarray. Because it is meant for internal use, it is stored in a torch tensor, not a DNDarray. \n", "\n", "The `lshape_map` tensor is a 2D tensor, where the first dimension is the number of processes and the second dimension is the number of dimensions of the array. Each row of the tensor contains the local shape of the array on a process. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[2:14]: \u001b[0m\n", "tensor([[2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [1, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "lshape_map = a.lshape_map\nlshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[2, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [1, 4, 3]])" }, "execution_count": 14, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:29.696699Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[3:14]: \u001b[0m\n", "tensor([[2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [1, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "lshape_map = a.lshape_map\nlshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[2, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [1, 4, 3]])" }, "execution_count": 14, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:29.697780Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[1:14]: \u001b[0m\n", "tensor([[2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [1, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "lshape_map = a.lshape_map\nlshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[2, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [1, 4, 3]])" }, "execution_count": 14, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:29.696572Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:14]: \u001b[0m\n", "tensor([[2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [1, 4, 3]])" ] }, "metadata": { "after": [], "completed": "2025-11-21T14:39:29.728858Z", "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "lshape_map = a.lshape_map\nlshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[2, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [1, 4, 3]])" }, "execution_count": 14, "metadata": {} }, "follow": [], "is_broadcast": false, "is_coalescing": false, "msg_id": "c9193665-68c4e705ebe8469470239198_2382638_13", "outputs": [], "received": "2025-11-21T14:39:29.732557Z", "started": "2025-11-21T14:39:29.699881Z", "status": "ok", "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:29.695557Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "lshape_map = a.lshape_map\n", "lshape_map" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Go back to where we created the DNDarray and create `a` with a different split axis. See how the `lshape_map` changes." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Modifying the DNDarray distribution\n", "\n", "In a distributed pipeline, it is sometimes necessary to change the distribution of a DNDarray, when the array is not distributed in the most convenient way for the next operation / algorithm.\n", "\n", "Depending on your needs, you can choose between:\n", "- `DNDarray.redistribute_()`: This method keeps the original split axis, but redistributes the data of the DNDarray according to a \"target map\".\n", "- `DNDarray.resplit_()`: This method changes the split axis of the DNDarray. This is a more expensive operation, and should be used only when absolutely necessary. Depending on your needs and available resources, in some cases it might be wiser to keep a copy of the DNDarray with a different split axis.\n", "\n", "Let's see some examples." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[1:15]: \u001b[0m\n", "tensor([[1, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "#redistribute\ntarget_map = a.lshape_map\ntarget_map[:, a.split] = torch.tensor([1, 2, 2, 2])\n# in-place redistribution (see ht.redistribute for out-of-place)\na.redistribute_(target_map=target_map)\n\n# new lshape map after redistribution\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[1, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [2, 4, 3]])" }, "execution_count": 15, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:34.626127Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[2:15]: \u001b[0m\n", "tensor([[1, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "#redistribute\ntarget_map = a.lshape_map\ntarget_map[:, a.split] = torch.tensor([1, 2, 2, 2])\n# in-place redistribution (see ht.redistribute for out-of-place)\na.redistribute_(target_map=target_map)\n\n# new lshape map after redistribution\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[1, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [2, 4, 3]])" }, "execution_count": 15, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:34.626201Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:15]: \u001b[0m\n", "tensor([[1, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "#redistribute\ntarget_map = a.lshape_map\ntarget_map[:, a.split] = torch.tensor([1, 2, 2, 2])\n# in-place redistribution (see ht.redistribute for out-of-place)\na.redistribute_(target_map=target_map)\n\n# new lshape map after redistribution\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[1, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [2, 4, 3]])" }, "execution_count": 15, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:34.625859Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[3:15]: \u001b[0m\n", "tensor([[1, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3],\n", " [2, 4, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "#redistribute\ntarget_map = a.lshape_map\ntarget_map[:, a.split] = torch.tensor([1, 2, 2, 2])\n# in-place redistribution (see ht.redistribute for out-of-place)\na.redistribute_(target_map=target_map)\n\n# new lshape map after redistribution\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[1, 4, 3],\n [2, 4, 3],\n [2, 4, 3],\n [2, 4, 3]])" }, "execution_count": 15, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:34.626272Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "#redistribute\n", "target_map = a.lshape_map\n", "target_map[:, a.split] = torch.tensor([1, 2, 2, 2])\n", "# in-place redistribution (see ht.redistribute for out-of-place)\n", "a.redistribute_(target_map=target_map)\n", "\n", "# new lshape map after redistribution\n", "a.lshape_map" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[3:27]: \u001b[0m\n", "tensor([[[ 0.5730, -1.0918, -0.8577],\n", " [ 0.6610, -0.4874, 0.9850],\n", " [ 1.0930, -0.8518, -0.7061],\n", " [-0.7625, 0.6767, 0.1940]],\n", "\n", " [[-1.1230, 0.2482, 0.7127],\n", " [-0.3202, -0.3510, -1.2052],\n", " [-1.0595, -0.5830, 0.4192],\n", " [ 0.5600, -1.2777, -0.1323]]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "b9f6f6e8-01c224a4024814eaffce2266", "error": null, "execute_input": "# local arrays after redistribution\na.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[ 0.5730, -1.0918, -0.8577],\n [ 0.6610, -0.4874, 0.9850],\n [ 1.0930, -0.8518, -0.7061],\n [-0.7625, 0.6767, 0.1940]],\n\n [[-1.1230, 0.2482, 0.7127],\n [-0.3202, -0.3510, -1.2052],\n [-1.0595, -0.5830, 0.4192],\n [ 0.5600, -1.2777, -0.1323]]])" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-05-19T19:20:48.893023Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[2:27]: \u001b[0m\n", "tensor([[[ 1.6286, 0.4707, -0.5730],\n", " [ 0.3841, -0.4789, -0.8033],\n", " [ 0.1299, -0.6602, -2.0182],\n", " [ 0.5541, -0.1653, -0.4314]],\n", "\n", " [[ 1.1544, -0.8126, -0.7634],\n", " [-0.0817, -1.5430, -0.6341],\n", " [ 0.0291, 0.9677, 0.1294],\n", " [-0.3747, -1.4987, -0.1063]]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "e3e9e719-1b11a826b66969f71d179e21", "error": null, "execute_input": "# local arrays after redistribution\na.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[ 1.6286, 0.4707, -0.5730],\n [ 0.3841, -0.4789, -0.8033],\n [ 0.1299, -0.6602, -2.0182],\n [ 0.5541, -0.1653, -0.4314]],\n\n [[ 1.1544, -0.8126, -0.7634],\n [-0.0817, -1.5430, -0.6341],\n [ 0.0291, 0.9677, 0.1294],\n [-0.3747, -1.4987, -0.1063]]])" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-05-19T19:20:48.892765Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[1:27]: \u001b[0m\n", "tensor([[[-0.0919, -0.7646, 0.1660],\n", " [-0.9814, 0.9445, -1.8339],\n", " [-1.0218, 0.8454, -0.6050],\n", " [-0.4161, -0.0764, 0.4383]],\n", "\n", " [[ 0.3151, -2.1761, 0.9970],\n", " [ 0.9423, 0.7667, 0.6834],\n", " [ 1.9586, -0.0994, 0.0186],\n", " [-0.0961, -0.3901, 1.2133]]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "4a6ffcbf-4b7c9961beb0aa49f4f299a5", "error": null, "execute_input": "# local arrays after redistribution\na.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[-0.0919, -0.7646, 0.1660],\n [-0.9814, 0.9445, -1.8339],\n [-1.0218, 0.8454, -0.6050],\n [-0.4161, -0.0764, 0.4383]],\n\n [[ 0.3151, -2.1761, 0.9970],\n [ 0.9423, 0.7667, 0.6834],\n [ 1.9586, -0.0994, 0.0186],\n [-0.0961, -0.3901, 1.2133]]])" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-05-19T19:20:48.892426Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:45]: \u001b[0m\n", "tensor([[[-0.1776, -0.8116, -0.6636],\n", " [ 0.3238, 2.4110, 0.4005],\n", " [-0.7808, -2.0984, 1.7691],\n", " [ 0.9370, 0.0141, 0.6934]]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "26ba0021-35d3d060b50582f7d11d6ead", "error": null, "execute_input": "# local arrays after redistribution\na.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[-0.1776, -0.8116, -0.6636],\n [ 0.3238, 2.4110, 0.4005],\n [-0.7808, -2.0984, 1.7691],\n [ 0.9370, 0.0141, 0.6934]]])" }, "execution_count": 45, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-05-19T19:20:48.891730Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "# local arrays after redistribution\n", "a.larray" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[3:16]: \u001b[0m\n", "tensor([[7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "# resplit\na.resplit_(axis=1)\n\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[7, 1, 3],\n [7, 1, 3],\n [7, 1, 3],\n [7, 1, 3]])" }, "execution_count": 16, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:37.734646Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:16]: \u001b[0m\n", "tensor([[7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "# resplit\na.resplit_(axis=1)\n\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[7, 1, 3],\n [7, 1, 3],\n [7, 1, 3],\n [7, 1, 3]])" }, "execution_count": 16, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:37.733951Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[1:16]: \u001b[0m\n", "tensor([[7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "# resplit\na.resplit_(axis=1)\n\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[7, 1, 3],\n [7, 1, 3],\n [7, 1, 3],\n [7, 1, 3]])" }, "execution_count": 16, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:37.734513Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[2:16]: \u001b[0m\n", "tensor([[7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3],\n", " [7, 1, 3]])" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "# resplit\na.resplit_(axis=1)\n\na.lshape_map\n", "execute_result": { "data": { "text/plain": "tensor([[7, 1, 3],\n [7, 1, 3],\n [7, 1, 3],\n [7, 1, 3]])" }, "execution_count": 16, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:37.734590Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "# resplit\n", "a.resplit_(axis=1)\n", "\n", "a.lshape_map" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can use the `resplit_` method (in-place), or `ht.resplit` (out-of-place) to change the distribution axis, but also to set the distribution axis to None. The latter corresponds to an MPI.Allgather operation that gathers the entire array on each process. This is useful when you've achieved a small enough data size that can be processed on a single device, and you want to avoid communication overhead." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[0:17]: \u001b[0m\n", "DNDarray([[[ 0.5357, -1.9890, -0.4163],\n", " [ 0.0807, 0.0103, -0.8651],\n", " [ 0.5091, -0.9029, -0.2220],\n", " [-1.3162, 0.4313, 2.0407]],\n", "\n", " [[ 0.5421, -1.1846, -1.2136],\n", " [-0.2742, 1.2303, 0.3469],\n", " [ 2.3878, 0.6656, -0.0703],\n", " [-0.7895, 0.2615, 1.4271]],\n", "\n", " [[ 0.4421, 0.8812, 0.0758],\n", " [-0.0802, -1.0479, 1.0972],\n", " [ 0.3463, 0.7280, -0.4285],\n", " [-0.9275, -2.3051, -0.4425]],\n", "\n", " [[-0.3667, 1.9827, 1.5161],\n", " [ 0.5978, -0.1475, -0.4847],\n", " [-0.8409, -1.1828, -0.4474],\n", " [-0.3909, -0.5872, -0.2087]],\n", "\n", " [[-1.5586, 0.7291, -0.8386],\n", " [ 0.1649, -0.7093, 0.0767],\n", " [ 0.5224, -1.7311, -0.6501],\n", " [-1.1856, -0.5394, 1.0237]],\n", "\n", " [[ 0.9554, 2.0103, 1.3122],\n", " [ 0.1958, -0.4657, -0.4318],\n", " [-1.0655, -1.0135, -0.7660],\n", " [-0.6943, 0.2370, 1.6383]],\n", "\n", " [[ 0.2506, -0.7750, -0.7626],\n", " [ 1.4696, 1.1283, 0.2312],\n", " [ 0.3658, 0.3560, -0.8484],\n", " [ 0.2760, -0.7415, 0.4204]]], dtype=ht.float32, device=cpu:0, split=None)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "# \"un-split\" distributed array\na.resplit_(axis=None)\n# each process now holds a copy of the entire array\n", "execute_result": { "data": { "text/plain": "DNDarray([[[ 0.5357, -1.9890, -0.4163],\n [ 0.0807, 0.0103, -0.8651],\n [ 0.5091, -0.9029, -0.2220],\n [-1.3162, 0.4313, 2.0407]],\n\n [[ 0.5421, -1.1846, -1.2136],\n [-0.2742, 1.2303, 0.3469],\n [ 2.3878, 0.6656, -0.0703],\n [-0.7895, 0.2615, 1.4271]],\n\n [[ 0.4421, 0.8812, 0.0758],\n [-0.0802, -1.0479, 1.0972],\n [ 0.3463, 0.7280, -0.4285],\n [-0.9275, -2.3051, -0.4425]],\n\n [[-0.3667, 1.9827, 1.5161],\n [ 0.5978, -0.1475, -0.4847],\n [-0.8409, -1.1828, -0.4474],\n [-0.3909, -0.5872, -0.2087]],\n\n [[-1.5586, 0.7291, -0.8386],\n [ 0.1649, -0.7093, 0.0767],\n [ 0.5224, -1.7311, -0.6501],\n [-1.1856, -0.5394, 1.0237]],\n\n [[ 0.9554, 2.0103, 1.3122],\n [ 0.1958, -0.4657, -0.4318],\n [-1.0655, -1.0135, -0.7660],\n [-0.6943, 0.2370, 1.6383]],\n\n [[ 0.2506, -0.7750, -0.7626],\n [ 1.4696, 1.1283, 0.2312],\n [ 0.3658, 0.3560, -0.8484],\n [ 0.2760, -0.7415, 0.4204]]], dtype=ht.float32, device=cpu:0, split=None)" }, "execution_count": 17, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:40.839690Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "# \"un-split\" distributed array\na.resplit_(axis=None)\n# each process now holds a copy of the entire array\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 17, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:40.840185Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "# \"un-split\" distributed array\na.resplit_(axis=None)\n# each process now holds a copy of the entire array\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 17, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:40.839956Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "# \"un-split\" distributed array\na.resplit_(axis=None)\n# each process now holds a copy of the entire array\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 17, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:40.840251Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "# \"un-split\" distributed array\n", "a.resplit_(axis=None)\n", "# each process now holds a copy of the entire array" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The opposite is not true, i.e. you cannot use `resplit_` to distribute an array with split=None. In that case, you must use the `ht.array()` factory function:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%px\n", "# make `a` split again\n", "a = ht.array(a, split=0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Making disjoint data into a global DNDarray\n", "\n", "Another common occurrence in a data-parallel pipeline: you have addressed the embarassingly-parallel part of your algorithm with any array framework, each process working independently from the others. You now want to perform a non-embarassingly-parallel operation on the entire dataset, with Heat as a backend.\n", "\n", "You can use the `ht.array` factory function with the `is_split` argument to create a DNDarray from a disjoint (on each MPI process) set of arrays. The `is_split` argument indicates the axis along which the disjoint data is to be \"joined\" into a global, distributed DNDarray." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[3:19]: \u001b[0m(12, 4)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee", "error": null, "execute_input": "# create some random local arrays on each process\nimport numpy as np\nlocal_array = np.random.rand(3, 4)\n\n# join them into a distributed array\na_0 = ht.array(local_array, is_split=0)\na_0.shape\n", "execute_result": { "data": { "text/plain": "(12, 4)" }, "execution_count": 19, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:46.335100Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:19]: \u001b[0m(12, 4)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27", "error": null, "execute_input": "# create some random local arrays on each process\nimport numpy as np\nlocal_array = np.random.rand(3, 4)\n\n# join them into a distributed array\na_0 = ht.array(local_array, is_split=0)\na_0.shape\n", "execute_result": { "data": { "text/plain": "(12, 4)" }, "execution_count": 19, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:46.334675Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[1:19]: \u001b[0m(12, 4)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "a1070193-d3a142f960197ef957f774ee", "error": null, "execute_input": "# create some random local arrays on each process\nimport numpy as np\nlocal_array = np.random.rand(3, 4)\n\n# join them into a distributed array\na_0 = ht.array(local_array, is_split=0)\na_0.shape\n", "execute_result": { "data": { "text/plain": "(12, 4)" }, "execution_count": 19, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:46.334963Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[2:19]: \u001b[0m(12, 4)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd", "error": null, "execute_input": "# create some random local arrays on each process\nimport numpy as np\nlocal_array = np.random.rand(3, 4)\n\n# join them into a distributed array\na_0 = ht.array(local_array, is_split=0)\na_0.shape\n", "execute_result": { "data": { "text/plain": "(12, 4)" }, "execution_count": 19, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:39:46.335026Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "# create some random local arrays on each process\n", "import numpy as np\n", "local_array = np.random.rand(3, 4)\n", "\n", "# join them into a distributed array\n", "a_0 = ht.array(local_array, is_split=0)\n", "a_0.shape" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Change the cell above and join the arrays along a different axis. Note that the shapes of the local arrays must be consistent along the non-split axes. They can differ along the split axis." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `ht.array` function takes any data object as an input that can be converted to a torch tensor. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Once you've made your disjoint data into a DNDarray, you can apply any Heat operation or algorithm to it and exploit the cumulative RAM of all the processes in the communicator. " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can access the MPI communication functionalities of the DNDarray through the `comm` attribute, i.e.:\n", "\n", "```python\n", "# these are just examples, this cell won't do anything\n", "a.comm.Allreduce(a, b, op=MPI.SUM)\n", "\n", "a.comm.Allgather(a, b)\n", "a.comm.Isend(a, dest=1, tag=0)\n", "```\n", "\n", "etc." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise: Comparing Redistribution and Resplit (solution below)\n", "\n", "This exercise explores the difference between `DNDarray.redistribute_()` and `DNDarray.resplit_()`.\n", "\n", "* **Task**:\n", " 1. Create a $10 \\times 4$ $\\mathbf{DNDarray}$ ($\\mathbf{A}$) distributed along $\\text{axis}=0$.\n", " 2. Use `DNDarray.resplit_()` to redistribute $\\mathbf{A}$ along $\\text{axis}=1$ ($\\mathbf{B}$).\n", " 3. Create a **target map** to redistribute $\\mathbf{A}$ along $\\text{axis}=0$ such that rank 0 holds 1 row, and all other ranks hold 3 rows. Apply `DNDarray.redistribute_()` ($\\mathbf{C}$).\n", "* **Check**: Print the `lshape` and `split` for $\\mathbf{A}$, $\\mathbf{B}$, and $\\mathbf{C}$ on each rank to confirm the new distributions." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "UsageError: Line magic function `%%px` not found.\n" ] } ], "source": [ "%%px\n", "import heat as ht\n", "import torch\n", "\n", "# 1. Create original DNDarray A (split=0)\n", "A = ht.random.randn(10, 4, split=0)\n", "rank = A.comm.rank\n", "\n", "# 2. Resplit A to split=1 (creating B)\n", "B = A.copy() # Use copy to keep A untouched for redistribution test\n", "B.resplit_(axis=1)\n", "\n", "# 3. Create target map and redistribute A (creating C)\n", "# Original size: 10 rows. Split=0. Ranks: 0, 1, 2, 3.\n", "# Target distribution (rows along axis 0): [1, 3, 3, 3]\n", "target_map = A.lshape_map.copy()\n", "# The original split axis is 0 (index 0 of lshape_map rows)\n", "target_map[:, A.split] = torch.tensor([1, 3, 3, 3])\n", "\n", "C = A.copy()\n", "C.redistribute_(target_map=target_map)\n", "\n", "# Check: Print properties\n", "print(f\"--- Rank {rank} ---\")\n", "print(f\"A (Original): Split={A.split}, lshape={A.lshape}\")\n", "print(f\"B (Resplit A to axis=1): Split={B.split}, lshape={B.lshape}\")\n", "print(f\"C (Redistribute A): Split={C.split}, lshape={C.lshape}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Check out how we use Heat's distributed-array infrastructure to scale complex data analysis workflows to large datasets and high-performance computing resources.\n", "\n", "- [Data loading and preprocessing](2_loading_preprocessing.ipynb)\n", "- [Linear algebra operations](Linear_algebra.ipynb)\n", "- [Clustering and PCA algorithms](Clustering_and_PCA.ipynb)\n", "- [Dynamic Mode Decomposition](DMD.ipynb)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.12" } }, "nbformat": 4, "nbformat_minor": 4 }