Add GPU usage in helloworld ipynb

Bluefog-Lib · Jan 19, 2021 · c25967b · c25967b
1 parent f88efcd
commit c25967b
Showing 1 changed file with 53 additions and 16 deletions.
diff --git a/examples/interactive_bluefog_helloworld.ipynb b/examples/interactive_bluefog_helloworld.ipynb
@@ -142,10 +142,10 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[stdout:0] Hello, I am 2 among 4 processes\n",
-      "[stdout:1] Hello, I am 3 among 4 processes\n",
-      "[stdout:2] Hello, I am 1 among 4 processes\n",
-      "[stdout:3] Hello, I am 0 among 4 processes\n"
+      "[stdout:0] Hello, I am 0 among 4 processes\n",
+      "[stdout:1] Hello, I am 1 among 4 processes\n",
+      "[stdout:2] Hello, I am 3 among 4 processes\n",
+      "[stdout:3] Hello, I am 2 among 4 processes\n"
      ]
     }
    ],
@@ -182,15 +182,15 @@
      "output_type": "stream",
      "text": [
       "==========================================\n",
-      "[stdout:0] Hello, I am 2 among 4 processes\n",
-      "[stdout:1] Hello, I am 3 among 4 processes\n",
-      "[stdout:2] Hello, I am 1 among 4 processes\n",
-      "[stdout:3] Hello, I am 0 among 4 processes\n",
+      "[stdout:0] Hello, I am 0 among 4 processes\n",
+      "[stdout:1] Hello, I am 1 among 4 processes\n",
+      "[stdout:2] Hello, I am 3 among 4 processes\n",
+      "[stdout:3] Hello, I am 2 among 4 processes\n",
       "==========================================\n",
-      "[stdout:0] Hello, I am 2 among 4 processes\n",
-      "[stdout:1] Hello, I am 3 among 4 processes\n",
-      "[stdout:2] Hello, I am 1 among 4 processes\n",
-      "[stdout:3] Hello, I am 0 among 4 processes\n",
+      "[stdout:0] Hello, I am 0 among 4 processes\n",
+      "[stdout:1] Hello, I am 1 among 4 processes\n",
+      "[stdout:2] Hello, I am 3 among 4 processes\n",
+      "[stdout:3] Hello, I am 2 among 4 processes\n",
       "==========================================\n"
      ]
     }
@@ -257,7 +257,7 @@
     {
      "data": {
       "text/plain": [
-       "tensor([2.])"
+       "tensor([0.])"
       ]
      },
      "execution_count": 6,
@@ -292,7 +292,7 @@
     {
      "data": {
       "text/plain": [
-       "[tensor([3.]), tensor([1.])]"
+       "[tensor([1.]), tensor([3.])]"
       ]
      },
      "execution_count": 7,
@@ -318,7 +318,7 @@
     {
      "data": {
       "text/plain": [
-       "[tensor([2.]), tensor([3.]), tensor([1.]), tensor([0.])]"
+       "[tensor([0.]), tensor([1.]), tensor([3.]), tensor([2.])]"
       ]
      },
      "execution_count": 8,
@@ -390,6 +390,43 @@
     "print(\"I received seed as value: \", seed)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## (Optional) Working with multiple GPUs\n",
+    "\n",
+    "If you have multiple GPUs, it is typical that each worker is pinned to one device.\n",
+    "Then, you can let the tensor wihin each worker set on different GPUs as shown in following code.\n",
+    "\n",
+    "*Note: If you want to pin multiple workers in one device, you cannot use NCCL as communication backend.*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[stdout:0] tensor([0.], device='cuda:0')\n",
+      "[stdout:1] tensor([1.], device='cuda:1')\n",
+      "[stdout:2] tensor([3.], device='cuda:3')\n",
+      "[stdout:3] tensor([2.], device='cuda:2')\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%px\n",
+    "import torch\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.set_device(bf.local_rank())\n",
+    "    x = torch.FloatTensor([bf.rank()]).cuda()\n",
+    "    print(x)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -754,7 +791,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.5"
+   "version": "3.7.9"
   },
   "toc": {
    "base_numbering": 1,