fix: inference key change, finetuning jaxlib update (#4447)

Author: dpanigra

notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb

@@ -106,11 +106,12 @@
         "id": "8c8f59a6c426"
       },
       "source": [
-        "### Prerequisites\n",
+        "## 0: Prerequisites\n",
         "\n",
         "- Install AlphaGenome Research and Google Cloud Platform packages.\n",
         "- (*) Choose either H100 or A100 specific vm notebook runtime.\n",
         "- (**) Save Huggingface credentials in Google Cloud Secret manager\n",
+        "- (***) Install 0.9.0 jax libraries\n",
         "\n",
         "Notebook launch:\n",
         "- Launch the Notebook in Google Cloud Enterprise Colab.\n",
@@ -130,7 +131,11 @@
         "You will be downloading weights from Huggingface.\n",
         "Ensure that:\n",
         "- You create a token that has 'Read access to contents of all public gated repos you can access' (under Finegrained control)\n",
-        "- You accept the T&C of the [model](https://huggingface.co/google/alphagenome-fold-0)."
+        "- You accept the T&C of the [model](https://huggingface.co/google/alphagenome-fold-0).\n",
+        "\n",
+        "\n",
+        "(***):\n",
+        "Upgrading to 0.9.0 will require multiple runtime restarts."
       ]
     },
     {
@@ -151,14 +156,65 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "ictFNXdeQ4Cf",
       "metadata": {
-        "id": "2bdbb159823b"
+        "id": "834dcef76adf"
+      },
+      "outputs": [],
+      "source": [
+        "import jax\n",
+        "# We need >0.9.0 jax libs.\n",
+        "# Check the jax version.\n",
+        "import jaxlib\n",
+        "\n",
+        "print(f\"{jax.__version__=}\")\n",
+        "print(f\"{jaxlib.__version__=}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9a3295c80eb9"
       },
       "outputs": [],
       "source": [
-        "! pip install --upgrade google-cloud-secret-manager \\\n",
-        "                              google-cloud-storage"
+        "# Uninstall the previous version.\n",
+        "# Run only if version < 0.9.0.\n",
+        "# Restart runtime/kernel after uninstalling.\n",
+        "# Run from next cell after the kernel restart\n",
+        "if jax.__version__ != \"0.9.0\":\n",
+        "    print(f\"Unistalling {jax.__version__}\")\n",
+        "    ! pip uninstall -y jax jaxlib jax_cuda12_plugin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3e0625b2f13d"
+      },
+      "outputs": [],
+      "source": [
+        "# Install specific vesion.\n",
+        "# Run only once after unistalling the jax packages.\n",
+        "# Restart the runtime/kernel.\n",
+        "# Run from next cell after the kernel restart.\n",
+        "!pip install --upgrade jax[cuda12_pip]==0.9.0 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b666617f1088"
+      },
+      "outputs": [],
+      "source": [
+        "# Install specific vesion.\n",
+        "# Run only once after the upgrade.\n",
+        "# Restart the runtime/kernel.\n",
+        "# Run from next cell after the kernel restart.\n",
+        "!pip install jax_cuda12_plugin==0.9.0"
       ]
     },
     {
@@ -195,7 +251,6 @@
         "\n",
         "import haiku as hk\n",
         "import huggingface_hub\n",
-        "import jax\n",
         "import numpy as np\n",
         "import optax\n",
         "import orbax.checkpoint as ocp\n",
@@ -247,12 +302,13 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "id": "TDvrVjiCzaYJ",
       "metadata": {
-        "id": "534dc48681bd"
+        "id": "8f34c4142dd0"
       },
       "outputs": [],
       "source": [
+        "import jax\n",
+        "# we need >0.9.0 jax libs\n",
         "# check the jax versions\n",
         "import jaxlib\n",
         "\n",
@@ -809,10 +865,7 @@
       },
       "outputs": [],
       "source": [
-        "forward_fn = finetune.get_forward_fn(\n",
-        "    output_metadata, jmp_policy=\"params=float32,compute=float32,output=float32\"\n",
-        ")\n",
-        "# forward_fn = finetune.get_forward_fn(output_metadata)\n",
+        "forward_fn = finetune.get_forward_fn(output_metadata)\n",
         "with jax.set_mesh(mesh):\n",
         "    batch = jax.device_put(batch, data_sharding)\n",
         "    params_ft, state_ft = jax.jit(\n",
@@ -952,10 +1005,9 @@
       },
       "outputs": [],
       "source": [
-        "loss = []\n",
-        "step = 0\n",
-        "start_time = time.monotonic()\n",
+        "loss, times = [], []\n",
         "for step in range(NUM_TRAIN_STEPS):\n",
+        "    start_time = time.time()\n",
         "    try:\n",
         "        batch = next(ds_iter)\n",
         "    except StopIteration:\n",
@@ -965,11 +1017,12 @@
         "        batch = jax.device_put(batch, data_sharding)\n",
         "        params, state, opt_state, scalars = train_step(params, state, opt_state, batch)\n",
         "    loss.append(scalars[\"loss\"])\n",
+        "    times.append(time.time() - start_time)\n",
         "    if step % 10 == 1:\n",
-        "        print(\"loss\", step, loss[-1])\n",
-        "end_time = time.monotonic()\n",
-        "duration = end_time - start_time\n",
-        "print(f\"Training took: {duration:.4f} seconds\")\n",
+        "        print(\"loss\", step, loss[-1], f\"SPS: {1./np.mean(times[1:]):.4f}\")\n",
+        "\n",
+        "print(f\"Total Training time: {np.sum(times[1:]):.4f} seconds\")\n",
+        "print(f\"Average Training time per step: {np.mean(times[1:]):.4f} seconds\")\n",
         "ckpt_path = save((params, state), step + 1)"
       ]
     },

notebooks/community/alphagenome/cloudai_alphagenome_vai_quickstart.ipynb

@@ -4458,24 +4458,14 @@
       "outputs": [],
       "source": [
         "def extract_k562(adata):\n",
-        "    if \"ontologyTerm\" in adata.var.columns and not adata.var[\"ontologyTerm\"].empty:\n",
-        "        mask = adata.var[\"ontologyTerm\"].apply(\n",
-        "            lambda x: isinstance(x, dict)\n",
-        "            and x.get(\"ontologyType\") == \"ONTOLOGY_TYPE_EFO\"\n",
-        "            and x.get(\"id\") == \"2067\"\n",
-        "        )\n",
-        "        values = adata.X[:, mask]\n",
-        "    else:\n",
-        "        raise ValueError(\n",
-        "            \"Expected 'ontologyTerm' column with dictionary values not found in adata.var\"\n",
-        "        )\n",
-        "    assert values.size == 1\n",
-        "    return values.flatten()[0]\n",
+        "  values = adata.X[:, adata.var['ontology_curie'] == 'EFO:0002067']\n",
+        "  assert values.size == 1\n",
+        "  return values.flatten()[0]\n",
         "\n",
         "\n",
         "ism_result = ism.ism_matrix(\n",
         "    [extract_k562(x[0]) for x in variant_scores],\n",
-        "    variants=[v[0].uns[\"variant\"] for v in variant_scores],\n",
+        "    variants=[v[0].uns['variant'] for v in variant_scores],\n",
         ")"
       ]
     },