Fix T4 OOM error visibility and parameter count claim in Lab 5.4

[Yassir-249]

Problem

Lab 5.4 (course_5/gdm_lab_5_4_full_parameter_fine_tuning_of_gemma.ipynb) has two issues that disrupt the lesson the lab is trying to deliver.

1. The OOM error doesn't match what the narrative tells students to expect

The lab text says students "should observe an 'out of memory' (OOM) error" when calling model.fit(...). With the package versions pinned by the project (keras==3.13.2, keras_hub==0.26.0, jax==0.7.2, cuDNN 9.10) the actual error students see is:

XlaRuntimeError: INTERNAL: No reference output found!

Root cause: setting XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 makes JAX preallocate 14.25 GiB and never release compilation buffers from model loading — so by the time .fit() runs, only ~1.7 GiB is free in the JAX pool. The cuDNN autotuner can't get workspace memory and surfaces an opaque INTERNAL error instead of the clean RESOURCE_EXHAUSTED the narrative implies.

Same root cause as #29 (which fixes the analogous issue in Labs 7.5 and 7.6).

2. The "1.3 billion parameters" narrative claim double-counts the token embedding

The "Understanding Gemma's Memory Footprint" cell says:

The Gemma 1B model has approximately 1.3 billion parameters (1B for the main transformer blocks and 300M for the token embeddings).

But Keras reports Total params: 999,885,952 for gemma3_1b. The token embedding is implemented as a ReversibleEmbedding shared between input and output, so its ~302M parameters are already counted inside the backbone's ~999M. Unique parameter count is ~1B, not 1.3B.

Fix

1. Apply the same allocator pattern as #29

In the install/imports cell:

• Move the JAX env-var settings before the import keras line
• Replace XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 with XLA_PYTHON_CLIENT_ALLOCATOR=platform
• Add XLA_FLAGS=--xla_gpu_enable_command_buffer=

After the fix, the same model.fit(...) call produces:

XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 12.52GiB
                 (13438550016B) on device ordinal 0

The 12.52 GiB matches the predicted Adam optimizer state (m + v + gradient at fp32 for ~1B trainable params), giving students a clear, sized OOM that maps directly onto the lab's stated lesson.

2. Correct the parameter count claim

Update the "Understanding Gemma's Memory Footprint" markdown cell from:

The Gemma 1B model has approximately 1.3 billion parameters (1B for the main transformer blocks and 300M for the token embeddings).

to:

The Gemma 1B model has approximately 1 billion parameters in total (around 700M for the main transformer blocks and around 300M for the token embeddings, which are shared between the input and output).

Verification

Tested end-to-end on Colab T4 GPU runtime against live HEAD of main:

	Error
Before fix	XlaRuntimeError: INTERNAL: No reference output found! (or INTERNAL: No valid config found! at lower batch size — both opaque)
After fix	XlaRuntimeError: RESOURCE_EXHAUSTED: Failed to allocate request for 12.52GiB (13438550016B) on device ordinal 0

Related

• fix: pin deps and fix T4 OOM in labs 7.5/7.6  #29 — same allocator pattern applied to Labs 7.5 and 7.6 (still open). This PR follows that same template and is independent.

[Yassir-249]

Resolved on main — the T4 OOM visibility and the parameter-count claim are both addressed now. Closing, thanks!