add mps runtime (rasbt#223)

agragui · Jun 19, 2024 · 1ba1bb4 · 1ba1bb4
1 parent fba2059
commit 1ba1bb4
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/ch07/01_main-chapter-code/ch07.ipynb b/ch07/01_main-chapter-code/ch07.ipynb
@@ -1111,6 +1111,15 @@
     "from functools import partial\n",
     "\n",
     "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "# If you have a Mac with Apple Silicon chip, you can uncomment the next lines of code\n",
+    "# to train the model on the Mac's GPU cores. However, as of this writing, this results in\n",
+    "# larger numerical deviations from the results shown in this chapter, because Apple Silicon\n",
+    "# support in PyTorch is still experimental\n",
+    "\n",
+    "# if torch.backends.mps.is_available():\n",
+    "#     device = torch.device(\"mps\")\n",
+    "\n",
     "print(\"Device:\", device)\n",
     "\n",
     "customized_collate_fn = partial(custom_collate_fn, device=device, allowed_max_length=1024)"
@@ -1743,9 +1752,11 @@
     "| Model              | Device                | Runtime for 2 Epochs |\n",
     "|--------------------|-----------------------|----------------------|\n",
     "| gpt2-medium (355M) | CPU (M3 MacBook Air)  | 15.78 minutes        |\n",
+    "| gpt2-small (124M)  | GPU (M3 MacBook Air)  | 10.77 minutes        |\n",
     "| gpt2-medium (355M) | GPU (L4)              | 1.83 minutes         |\n",
     "| gpt2-medium (355M) | GPU (A100)            | 0.86 minutes         |\n",
     "| gpt2-small (124M)  | CPU (M3 MacBook Air)  | 5.74 minutes         |\n",
+    "| gpt2-small (124M)  | GPU (M3 MacBook Air)  | 3.73 minutes         |\n",
     "| gpt2-small (124M)  | GPU (L4)              | 0.69 minutes         |\n",
     "| gpt2-small (124M)  | GPU (A100)            | 0.39 minutes         |\n",
     "\n",