Gemma 3 1b model loading asyncio error on kaggle notebook

[kamalkraj]

Expected Behavior
The model loading without any asyncio error.
Actual Behavior
The model is loaded properly. But the notebook cell execution is stopped after the model load cell. Need to manually run after the model loading cells.

Steps to Reproduce the Problem

tunix-error.ipynb

Environment

• OS: Kaggle TPU Notebook
• Project Version:

Package Name	Version
jax[tpu]	0.8.1
flax	0.12.1
optax	0.2.6
qwix	0.1.5
orbax-checkpoint	0.11.31
grain	0.2.15
google-metrax	0.2.4
google-tunix[prod]	0.1.5

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided all the required information in the "Environment" section.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?

[selamw1]

Hi - The gemma-3-1b-it model's loading process has been verified as successful on both Kaggle and Colab-TPUs using the latest library versions. Please refer to the accompanying gist, which outlines the recommended model loading procedure as well.

[mdda]

Your tests don't prove that this is a non-issue : They are implicitly relying on extra instructions in the same cell so that the async issue get hidden in other latencies/async calls.

One way I found to cleanly let orbax do its loading in a Jupyter cell is to import asyncio at the top of the notebook, and then the loader can used ~stand-alone in a cell:

gemma = params.create_model_from_checkpoint(MODEL_CP_PATH, config)

# Add this:
await asyncio.sleep(0)  # This is required to ensure that we can move on to next cell cleanly

This allows for an extra cycle of callback mop-up, which respects the notebook's ownership of the async loop.

[kamalkraj]

Your tests don't prove that this is a non-issue : They are implicitly relying on extra instructions in the same cell so that the async issue get hidden in other latencies/async calls.

One way I found to cleanly let orbax do its loading in a Jupyter cell is to import asyncio at the top of the notebook, and then the loader can used ~stand-alone in a cell:

gemma = params.create_model_from_checkpoint(MODEL_CP_PATH, config)

Add this:

await asyncio.sleep(0) # This is required to ensure that we can move on to next cell cleanly
This allows for an extra cycle of callback mop-up, which respects the notebook's ownership of the async loop.

Thanks @mdda