WIP

graphcore-research · Oct 1, 2024 · 990d731 · 990d731
1 parent 4dda0a3
commit 990d731
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/docs/JAX FP8 matmul tutorial.ipynb b/docs/JAX FP8 matmul tutorial.ipynb
@@ -118,7 +118,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2024-09-25 14:11:11.258006: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.82). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n"
+      "2024-10-01 14:13:43.092200: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.5.82). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n"
      ]
     },
     {
@@ -206,7 +206,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 9,
    "id": "72d805ea-89b6-457d-9558-ff31fdd23d35",
    "metadata": {},
    "outputs": [
@@ -229,7 +229,7 @@
           "1"
          ]
         },
-        "epilogue": "DEFAULT",
+        "epilogue": "RELU",
         "grad_x": false,
         "grad_y": false,
         "lhs_stride": "2048",
@@ -241,7 +241,7 @@
          ]
         },
         "rhs_stride": "8192",
-        "selected_algorithm": "5"
+        "selected_algorithm": "2"
        },
        "operation_queue_id": "0",
        "wait_on_operation_queues": []
@@ -250,23 +250,23 @@
        "<IPython.core.display.JSON object>"
       ]
      },
-     "execution_count": 4,
      "metadata": {
       "application/json": {
        "expanded": true,
        "root": "root"
       }
      },
-     "output_type": "execute_result"
+     "output_type": "display_data"
     }
    ],
    "source": [
     "from IPython.display import JSON\n",
+    "from IPython.display import display, JSON\n",
     "\n",
     "hlo_module = parse_hlo_module(fn_compiled)\n",
     "# Let's extract the `backend_config` dict from the FP8 matmul call.\n",
     "backend_config = next((m.backend_config for m in hlo_module if \"__cublas$lt$matmul$f8\" in m.cmd))\n",
-    "JSON(backend_config, expanded=True)"
+    "display(JSON(backend_config, expanded=True))"
    ]
   },
   {