Transient Gemma-1B Loading Error in Colab GPU Runtime#9
Open
boiled-darvari wants to merge 5 commits into
Open
Conversation
…les. - fix grammatical error in README.md: "The eight courses are in the curriculum are:" -> "The eight courses in the curriculum are:" - fix misspelling in ai_foundations/visualizations/plots.py: "correclty" -> "correctly" - remove word repetition "and and" in course_2/gdm_lab_2_1_preprocess_data.ipynb - remove word repetition "to to" in course_2/gdm_lab_2_2_tokenize_texts_into_characters_and_words.ipynb - fix misspelling "occurence" -> "occurrence" in course_3/gdm_lab_3_1_distinguish_between_signal_and_noise.ipynb - remove word repetition "at at" in course_3/gdm_lab_3_6_mitigate_overfitting.ipynb - fix "fo" -> "of" in course_3/gdm_lab_3_6_mitigate_overfitting.ipynb - remove word repetition "a a" in course_4/gdm_lab_4_5_reflection_on_trainable_parameters.ipynb
Fix grammatical errors in course materials
docs: fixed a small typo in the notebook
Transient Gemma-1B Loading Error in Colab GPU Runtime
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Author
|
I think I found the main issue: Changing the JAX version to But Gemma-1B model loads successfully in the first try! You can reproduce the issue with this code cell alone: !pip install orbax-checkpoint==0.11.21 jax[cuda12]==0.6.2
!pip install "git+https://github.com/google-deepmind/ai-foundations.git@main"
# Packages used.
import os # For setting a variable needed to load the model onto the GPU.
import pandas as pd # For loading the Africa Galore dataset.
# Functions for clearing outputs and formatting.
from IPython.display import clear_output, display, HTML
# Functions for generating texts with a language model, visualizing probability
# distributions, and loading an n-gram model.
from ai_foundations import generation
from ai_foundations import visualizations
from ai_foundations.ngram import model as ngram_model
# Set the full GPU memory usage for JAX.
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "1.00"Then: import time
import jax
# waiting for CUDA backend to be initialized
max_check_retries = 30
print("Waiting for CUDA backend initialization...")
for i in range(max_check_retries):
try:
print(jax.devices())
print("CUDA backend initialized successfully.")
break
except Exception as e:
print(".")
if i < max_check_retries - 1:
time.sleep(1)
else:
print("CUDA backend initialization failed, try again.")
raiseWith failure on the first try, same as loading Gemma-1B model if the JAX version is So it's not about waiting for the CUDA backend to be initialized, as I used It seems more complex than just simple PR. Let me know if you need to make any changes to all the notebooks. And I don't mind if you change yourself by even closing this one. |
3 tasks
Author
|
Google should build a Gemini Dependency Hell Update Agent just for this |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When running the "Compare N-Gram Models and Transformer Language Models" lab in Google Colab with a GPU runtime, the initial attempt to load the Gemma-1B model using
ai_foundations.generation.load_gemma()occasionally fails with aRuntimeErrorrelated to CUDA backend initialization.The error observed is:
Subsequent attempts to run the cell and load the model typically succeed without error. This suggests a potential timing or transient issue during the initial setup of the JAX-CUDA backend in the Colab environment.
Steps to Reproduce:
generation.load_gemma().RuntimeErroron the first execution of the cell loading Gemma.Expected Behavior:
The Gemma-1B model should load successfully on the first attempt in a correctly configured Colab GPU runtime.
Observed Behavior:
The initial attempt to load the Gemma-1B model sometimes fails with a CUDA backend initialization error.
Environment:
ai_foundations,jax,orbax-checkpoint, etc.)Proposed Solution/Workaround:
A temporary workaround is to implement a retry mechanism around the
generation.load_gemma()call, catching theRuntimeErrorand retrying the load after a short delay. This is the approach currently implemented in the notebook to mitigate the issue for users.Note on GPU Initialization:
While a retry mechanism was implemented to address a transient issue observed during the initial loading of the Gemma model on Colab's GPU runtime, a definitive, quick, and correct method to explicitly "wait" for the underlying JAX CUDA backend to be fully ready before attempting the load was not identified. The current approach provides a practical workaround for the observed behavior, where the first loading attempt might fail due to initialization timing.
After fix:
P.S.: If the retry logic seems confusing for beginners, another option is to insert a single
sleep(n)before loading the model. Alternatively, you can simply let learners decide whether to rerun the code cell if they encounter a loading error on their first attempt.