Bug: Error loading the Gemma-1B model in Course 1 Lab 3 (gdm_lab_1_3_compare_n_gram_models_and_transformer_language_models.ipynb)

[nicolasguillard]

Description
Running the notebook on Colab with T4 GPU (Neither paid plan nor Comput unit).

The execution of this statement gemma_model = generation.load_gemma() (Section "Load the models", first Python block, line 14) raises this error message :

Loading Gemma-1B model...
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in backends()
    811 
--> 812         backend = _init_backend(platform)
    813         _backends[platform] = backend

13 frames
[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in _init_backend(platform)
    895   logger.debug("Initializing backend '%s'", platform)
--> 896   backend = registration.factory()
    897   # TODO(skye): consider raising more descriptive errors directly from backend

[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in make_pjrt_c_api_client(plugin_name, options)
    548   if distributed.global_state.client is None:
--> 549     return xla_client.make_c_api_client(plugin_name, updated_options, None)
    550 

[/usr/local/lib/python3.12/dist-packages/jaxlib/xla_client.py](https://localhost:8080/#) in make_c_api_client(plugin_name, options, distributed_client, transfer_server_factory)
    155     options = {}
--> 156   return _xla.get_c_api_client(
    157       plugin_name,

XlaRuntimeError: INVALID_ARGUMENT: Unexpected option name passed to PJRT_Client_Create: abort_collectives_on_failure

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
[/tmp/ipython-input-2532588917.py](https://localhost:8080/#) in <cell line: 0>()
      1 print("Loading Gemma-1B model...")
----> 2 gemma_model = generation.load_gemma()
      3 print("Loaded Gemma-1B model.")

[/usr/local/lib/python3.12/dist-packages/ai_foundations/generation/loaders.py](https://localhost:8080/#) in load_gemma(model_name)
     57     else:
     58       model = attention.AttentionWeightGemma3_1B()
---> 59     params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA3_1B_PT)
     60   elif model_name == "Gemma-4B":
     61     tokenizer = gm.text.Gemma3Tokenizer()

[/usr/local/lib/python3.12/dist-packages/gemma/gm/ckpts/_checkpoint.py](https://localhost:8080/#) in load_params(path, params, donate, text_only, sharding, quantize)
    184     raise ValueError('`sharding` and `params` are mutually exclusive.')
    185 
--> 186   ckpt = ocp.StandardCheckpointer()
    187 
    188   metadata, path = _get_metadata_and_path(ckpt, path)

[/usr/local/lib/python3.12/dist-packages/orbax/checkpoint/_src/checkpointers/standard_checkpointer.py](https://localhost:8080/#) in __init__(self, async_options, multiprocessing_options, file_options, checkpoint_metadata_store, temporary_path_class, **kwargs)
     83     """
     84     super().__init__(
---> 85         standard_checkpoint_handler.StandardCheckpointHandler(
     86             multiprocessing_options=multiprocessing_options,
     87             **kwargs,

[/usr/local/lib/python3.12/dist-packages/orbax/checkpoint/_src/handlers/standard_checkpoint_handler.py](https://localhost:8080/#) in __init__(self, save_concurrent_gb, restore_concurrent_gb, multiprocessing_options, pytree_metadata_options)
     95     """
     96     self._supported_types = checkpoint_utils.STANDARD_ARRAY_TYPES
---> 97     self._impl = pytree_checkpoint_handler.PyTreeCheckpointHandler(
     98         save_concurrent_gb=save_concurrent_gb,
     99         restore_concurrent_gb=restore_concurrent_gb,

[/usr/local/lib/python3.12/dist-packages/orbax/checkpoint/_src/handlers/pytree_checkpoint_handler.py](https://localhost:8080/#) in __init__(self, aggregate_filename, save_concurrent_gb, restore_concurrent_gb, use_ocdbt, use_zarr3, multiprocessing_options, type_handler_registry, handler_impl, pytree_metadata_options, array_metadata_validator, enable_pinned_host_transfer)
    536     self._save_concurrent_bytes = _concurrent_bytes(save_concurrent_gb)
    537     self._restore_concurrent_bytes = _concurrent_bytes(restore_concurrent_gb)
--> 538     self._handler_impl = handler_impl or BasePyTreeCheckpointHandler(
    539         save_concurrent_bytes=self._save_concurrent_bytes,
    540         restore_concurrent_bytes=self._restore_concurrent_bytes,

[/usr/local/lib/python3.12/dist-packages/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py](https://localhost:8080/#) in __init__(self, save_concurrent_bytes, restore_concurrent_bytes, use_ocdbt, use_zarr3, multiprocessing_options, type_handler_registry, enable_post_merge_validation, pytree_metadata_options, array_metadata_validator, enable_pinned_host_transfer)
    367 
    368     if enable_pinned_host_transfer is None:
--> 369       enable_pinned_host_transfer = jax.default_backend() == 'gpu'
    370     self._enable_pinned_host_transfer = enable_pinned_host_transfer
    371 

[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in default_backend()
   1013 def default_backend() -> str:
   1014   """Returns the platform name of the default XLA backend."""
-> 1015   return get_backend(None).platform
   1016 
   1017 

[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in get_backend(platform)
    942     platform: None | str | xla_client.Client = None
    943 ) -> xla_client.Client:
--> 944   return _get_backend_uncached(platform)
    945 
    946 

[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in _get_backend_uncached(platform)
    921   platform = (platform or _XLA_BACKEND.value or _PLATFORM_NAME.value or None)
    922 
--> 923   bs = backends()
    924   if platform is not None:
    925     platform = canonicalize_platform(platform)

[/usr/local/lib/python3.12/dist-packages/jax/_src/xla_bridge.py](https://localhost:8080/#) in backends()
    826           else:
    827             err_msg += " (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)"
--> 828           raise RuntimeError(err_msg)
    829 
    830     assert _default_backend is not None

RuntimeError: Unable to initialize backend 'cuda': INVALID_ARGUMENT: Unexpected option name passed to PJRT_Client_Create: abort_collectives_on_failure (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)

I tryed unsuccessfully importing other versions of jax[cuda12] or even jax[cuda13], I got other different errors. And the statement JAX_PLATFORMS=cpuraises the error message "cpu unknown".

[nicolasguillard]

Fixed with this pull request : Transient Gemma-1B Loading Error in Colab GPU Runtime.