xiaohei/model_garden
1
0
Fork 0
mirror of https://github.com/GoogleCloudPlatform/vertex-ai-samples.git synced 2026-09-26 14:42:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Damodar Panigrahi
2026-02-23 19:45:59 +00:00
committed by GitHub
parent 649800e646
commit 7750e83fbb
2 changed files with 77 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/alphagenome/cloudai_alphagenome_finetune.ipynb
+73 -20
View File
@@ -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": [
"! pip install --upgrade google-cloud-secret-manager \\\n",
" google-cloud-storage"
"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": [
"# 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
+2 -12
View File
@@ -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",
" 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",
")"
]
},