Skip to content

fix: pin deps and fix T4 OOM in labs 7.5/7.6 #29

Open
aymansaeed9 wants to merge 5 commits into
google-deepmind:mainfrom
aymansaeed9:main
Open

fix: pin deps and fix T4 OOM in labs 7.5/7.6 #29
aymansaeed9 wants to merge 5 commits into
google-deepmind:mainfrom
aymansaeed9:main

Conversation

@aymansaeed9

Copy link
Copy Markdown

Problem

  1. kauldron >= 1.4.0 and typeguard >= 4.4 break notebook 7.1 imports and runtime type-checking .
  2. Labs 7.5 and 7.6 OOM on T4 GPUs because the default JAX pre-allocator doesn't free compilation buffers after model loading

Fix

  • Pin kauldron < 1.4.0 and typeguard < 4.4 in pyproject.toml
  • Use XLA_PYTHON_CLIENT_ALLOCATOR=platform instead of XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 in labs 7.5 and 7.6

Supersedes #27. Related: #25

@google-cla

google-cla Bot commented Apr 14, 2026

Copy link
Copy Markdown

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.

upaq pushed a commit that referenced this pull request Apr 17, 2026
…by @aymansaeed9.

- Lab 7.1: Pin kauldron version for compatibility to version 1.3.0; typeguard to 4.4.4.
- Labs 7.5 and 7.6: Override default JAX memory allocation to avoid out of memory errors on T4 GPUs

PiperOrigin-RevId: 901205486
Change-Id: I4490427a6a0b0cee9be6b074c5e741ce6b14ecae
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant