Moved tunix to setup.sh to only install when DEVICE=tpu - #2394
Merged
Conversation
Contributor
|
🤖 Hi @parambole, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Collaborator
Rohan-Bierneni
marked this pull request as ready for review
September 24, 2025 23:52
Rohan-Bierneni
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
vipannalla and
yangyuwei
as code owners
September 24, 2025 23:52
Contributor
|
🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Rohan-Bierneni
force-pushed
the
rbierneni-move-tunix-installation
branch
2 times, most recently
from
September 25, 2025 18:26
d116300 to
87b8edc
Compare
Contributor
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
📋 Review Summary
This pull request correctly moves the tunix dependency from the general requirements to be installed only in TPU environments. This resolves an issue with GPU environment setup. The changes in the Dockerfile and setup.sh are consistent and well-implemented.
🔍 General Feedback
- The conditional installation of
tunixis a good solution to handle environment-specific dependencies. - The code is clean and the changes are easy to understand.
Removed tunix from requirements files Install tunix if device=tpu Mark tunix-based tests as tpu-only Ignore sft hooks test for gpu tests
Rohan-Bierneni
force-pushed
the
rbierneni-move-tunix-installation
branch
from
September 25, 2025 20:27
e0ef316 to
38f38dd
Compare
parambole
approved these changes
Sep 25, 2025
bvandermoon
approved these changes
Sep 25, 2025
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
6 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.
Description
Tunix has a requirement on jax[tpu], which is causing errors and bugs in setting up maxtext environment for GPUs since it is in requirements.txt.
When using the JAX AI Image flow, since jax is installed as part of the image and tunix is installed when using JAX AI Image as a base image, there is a weird bug where libtpu is being installed in the final maxtext GPU image. This is probably due to jax and tunix not being installed in the same step, which would normally throw an error but is hidden in this 2-stage flow.
When using setup.sh, since both jax[cuda12] and tunix are being installed in the same step, pip correctly recognizes that tunix depends on jax[tpu] and will throw an error here.
To fix this, we will move tuni from requirements.txt to setup.sh and install the pinned commit when MODE=stable and the unpinned repo when MODE=nightly. We already do something similar for transformer_engine which is a gpu-only dependency.
Note: when tunix supports GPUs, we can move it back to requirements.txt
Note2: Added a small change for migration to jax-ai-image from jax-stable-stack in the entrypoint for JAX AI Images.
Tests
Gpu image built in the unit test run does not contain libtpu: https://screenshot.googleplex.com/AxYbecHHimTdRx8
Tpu image has tunix installed: https://screenshot.googleplex.com/3JM78mv6ZJR76ju
PR is working as intended
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.