Skip to content

Make tensorflow optional in requirements #5086

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aireenmei wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make tensorflow optional in requirements #5086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion docs/guides/data_input_pipeline/data_input_tfds.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,18 @@
# TFDS pipeline
# TFDS pipeline (Deprecated)

```{warning}
The TFDS input pipeline (`dataset_type=tfds`) is **deprecated**. We recommend migrating to the [Grain pipeline](data_input_grain.md) with `dataset_type=grain` and `grain_file_type=tfrecord`. You can keep the same TFRecord dataset paths.
```

````{note}
TensorFlow and TensorFlow Datasets (TFDS) are optional dependencies in MaxText. If you need to use the legacy TFDS pipeline, install the optional dependencies by running:

```bash
install_tpu_pre_train_extra_deps --with-tf
# or for GPU:
# install_cuda12_pre_train_extra_deps --with-tf
```
````

1. Download the Allenai C4 dataset in TFRecord format to a Cloud Storage bucket. For information about cost, see [this discussion](https://github.com/allenai/allennlp/discussions/5056)

Expand Down
8 changes: 8 additions & 0 deletions docs/install_maxtext.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ This is the easiest way to get started with the latest stable version.
```bash
uv pip install maxtext[tpu]=={{version}} --resolution=lowest
install_tpu_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_tpu_pre_train_extra_deps --with-tf
```

- **Option 2:** Install `maxtext[cuda12]`, used for pre-training and decoding
Expand All @@ -61,6 +63,8 @@ This is the easiest way to get started with the latest stable version.
```bash
uv pip install maxtext[cuda12]=={{version}} --resolution=lowest
install_cuda12_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_cuda12_pre_train_extra_deps --with-tf
```

- **Option 3:** Install `maxtext[tpu-post-train]`, used for post-training on
Expand Down Expand Up @@ -130,13 +134,17 @@ environment to avoid dependency conflicts.
```bash
uv pip install -e .[tpu] --resolution=lowest
install_tpu_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_tpu_pre_train_extra_deps --with-tf
```

- **Option 2:** Install `.[cuda12]`

```bash
uv pip install -e .[cuda12] --resolution=lowest
install_cuda12_pre_train_extra_deps
# Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
# install_cuda12_pre_train_extra_deps --with-tf
```

- **Option 3:** Install `.[tpu-post-train]`
Expand Down
8 changes: 6 additions & 2 deletions src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ ENV ENV_JAX_VERSION=$JAX_VERSION
ARG DEVICE
ENV ENV_DEVICE=$DEVICE

# TODO: remove default once separate TF and TF-free nightly image workflows are established
ARG TF=true
ENV ENV_TF=$TF

ARG PACKAGE_DIR
ENV PACKAGE_DIR=$PACKAGE_DIR

Expand All @@ -56,10 +60,10 @@ COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/

# Install dependencies - these steps are cached unless the copied files change
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE} TF=${ENV_TF}"
RUN --mount=type=cache,target=/root/.cache/uv \
export UV_LINK_MODE=copy && \
bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE} TF=${ENV_TF}

# Now copy the remaining code (source files that may change frequently)
COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
Expand Down
8 changes: 6 additions & 2 deletions src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,10 @@ ENV ENV_LIBTPU_VERSION=$LIBTPU_VERSION
ARG DEVICE
ENV ENV_DEVICE=$DEVICE

# TODO: remove default once separate TF and TF-free nightly image workflows are established
ARG TF=true
ENV ENV_TF=$TF

ARG PACKAGE_DIR
ENV PACKAGE_DIR=$PACKAGE_DIR

Expand All @@ -65,10 +69,10 @@ COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/
COPY libtpu.so* /root/custom_libtpu/

# Install dependencies - these steps are cached unless the copied files change
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE} TF=${ENV_TF}"
RUN --mount=type=cache,target=/root/.cache/uv \
export UV_LINK_MODE=copy && \
bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE} TF=${ENV_TF}

# Now copy the remaining code (source files that may change frequently)
COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
Expand Down
2 changes: 0 additions & 2 deletions src/dependencies/extra_deps/cuda12_overrides.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,5 @@ datasets>=4.8.5
fsspec==2026.2.0
gcsfs==2026.2.0
orbax-checkpoint>=0.12.1
tensorflow==2.20.0
tensorflow-text==2.20.1
markdown>=3.10.2
xprof<=2.23.0
1 change: 1 addition & 0 deletions src/dependencies/extra_deps/jetstream_github_deps.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
1 change: 0 additions & 1 deletion src/dependencies/extra_deps/pre_train_github_deps.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1 @@
google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
4 changes: 4 additions & 0 deletions src/dependencies/extra_deps/tf_requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
tensorflow==2.20.0
tensorflow-datasets
tensorflow-text==2.20.1
Comment thread
aireenmei marked this conversation as resolved.
seqio
2 changes: 0 additions & 2 deletions src/dependencies/extra_deps/tpu_overrides.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,5 @@ datasets>=4.8.5
fsspec==2026.2.0
gcsfs==2026.2.0
orbax-checkpoint>=0.12.4
tensorflow==2.20.0
tensorflow-text==2.20.1
markdown>=3.10.2
xprof<=2.23.0
4 changes: 0 additions & 4 deletions src/dependencies/requirements/base_requirements/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,8 @@ pylint
pytest
pytype
sentencepiece
seqio
tensorboard-plugin-profile
tensorboardx
tensorflow-datasets
tensorflow-text
tensorflow
tiktoken
tokamax!=0.1.0
transformers>=5.8.0
Expand Down
31 changes: 8 additions & 23 deletions src/dependencies/requirements/generated_requirements/cuda12-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ chex>=0.1.92
click>=8.5.0
cloud-accelerator-diagnostics>=0.1.1
cloudpickle>=3.1.2
clu>=0.0.12
colorama>=0.4.6
contourpy>=1.3.3
coverage>=7.16.0
Expand All @@ -40,11 +39,8 @@ defusedxml>=0.7.1
dill>=0.4.1
distlib>=0.4.3
distro>=1.9.0
dm-tree>=0.1.10
docstring-parser>=0.18.0
drjax>=0.2.1
editdistance>=0.8.1
einops>=0.8.2
einshape>=1.0
etils>=1.14.0
evaluate>=0.2.0
Expand All @@ -60,11 +56,11 @@ fsspec>=2026.2.0
gast>=0.7.0
gcsfs>=2026.2.0
google-api-core>=2.34.0
google-api-python-client>=2.199.0
google-api-python-client>=2.200.0
google-auth>=2.57.0
google-auth-httplib2>=0.4.2
google-auth-oauthlib>=1.4.1
google-cloud-aiplatform>=2.0.1
google-cloud-aiplatform>=2.1.0
google-cloud-appengine-logging>=1.10.0
google-cloud-audit-log>=0.6.2
google-cloud-bigquery>=3.44.0
Expand All @@ -76,7 +72,7 @@ google-cloud-resource-manager>=1.18.0
google-cloud-storage>=3.13.1
google-cloud-storage-control>=1.14.0
google-crc32c>=1.8.0
google-genai>=2.20.0
google-genai>=2.21.0
google-pasta>=0.2.0
google-resumable-media>=2.10.2
googleapis-common-protos>=1.75.2
Expand Down Expand Up @@ -115,7 +111,7 @@ kiwisolver>=1.5.0
latex2sympy2-extended>=1.11.0
libclang>=18.1.1
libcst>=1.9.0
lxml>=6.1.2
lxml>=6.1.3
markdown>=3.10.2
markdown-it-py>=4.2.0
markupsafe>=3.0.3
Expand Down Expand Up @@ -149,7 +145,7 @@ nvidia-cuda-nvcc-cu12>=12.9.86 ; sys_platform == 'linux'
nvidia-cuda-nvrtc-cu12>=12.9.86 ; sys_platform == 'linux'
nvidia-cuda-runtime-cu12>=12.9.79 ; sys_platform == 'linux'
nvidia-cudnn-cu12>=9.24.0.43 ; sys_platform == 'linux'
nvidia-cudnn-frontend>=1.27.0
nvidia-cudnn-frontend>=1.28.0
nvidia-cufft-cu12>=11.4.1.4 ; sys_platform == 'linux'
nvidia-cusolver-cu12>=11.7.5.82 ; sys_platform == 'linux'
nvidia-cusparse-cu12>=12.5.10.65 ; sys_platform == 'linux'
Expand All @@ -171,12 +167,11 @@ parameterized>=0.9.0
pathspec>=1.1.1
pathwaysutils>=0.1.11
pillow>=12.3.0
platformdirs>=4.11.5
platformdirs>=4.11.7
pluggy>=1.6.0
portpicker>=1.6.0
pre-commit>=4.6.2
prometheus-client>=0.26.0
promise>=2.3
propcache>=0.5.2
proto-plus>=1.28.4
protobuf>=6.33.6
Expand All @@ -191,7 +186,6 @@ pydantic>=2.13.5
pydantic-core>=2.46.5
pydot>=4.0.1
pyelftools>=0.33
pyglove>=0.4.5
pygments>=2.21.0
pyink>=26.5.1
pylint>=4.0.8
Expand All @@ -206,7 +200,7 @@ pytokens>=0.4.1
pytype>=2024.10.11
pyyaml>=6.0.3
qwix>=0.1.8
regex>=2026.8.31
regex>=2026.9.3
requests>=2.34.2
requests-oauthlib>=2.0.0
responses>=0.18.0
Expand All @@ -215,11 +209,9 @@ safetensors>=0.8.0
scipy>=1.18.0
scipy-stubs>=1.18.0.1
sentencepiece>=0.2.2
seqio>=0.0.20
setuptools>=84.0.0
shellingham>=1.5.4
shortuuid>=1.0.13
simple-parsing>=0.1.9
simplejson>=4.1.2
six>=1.17.0
sniffio>=1.3.1
Expand All @@ -228,14 +220,8 @@ starlette>=1.6.0
sympy>=1.12
tabulate>=0.10.0
tenacity>=9.1.4
tensorboard>=2.20.0
tensorboard-data-server>=0.7.0
tensorboard-plugin-profile>=2.13.0
tensorboardx>=2.6.5
tensorflow>=2.20.0
tensorflow-datasets>=4.9.10
tensorflow-metadata>=1.21.0
tensorflow-text>=2.20.1
tensorstore>=0.1.85
termcolor>=3.3.0
tiktoken>=0.14.0
Expand All @@ -259,12 +245,11 @@ uritemplate>=4.2.0
urllib3>=2.7.0
uvicorn>=0.52.4
uvloop>=0.22.1
virtualenv>=21.7.7
virtualenv>=21.7.8
wadler-lindig>=0.1.7
websockets>=16.1.1
werkzeug>=3.1.8
wheel>=0.48.0
wrapt>=2.4.0
xprof>=2.20.0
xxhash>=4.0.1
yarl>=1.24.5
Expand Down
23 changes: 4 additions & 19 deletions src/dependencies/requirements/generated_requirements/decoupled-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ cheroot>=11.1.2
chex>=0.1.92
click>=8.5.0
cloudpickle>=3.1.2
clu>=0.0.12
colorama>=0.4.6
contourpy>=1.3.3
coverage>=7.16.0
Expand All @@ -43,11 +42,8 @@ defusedxml>=0.7.1
dill>=0.4.1
distlib>=0.4.3
distro>=1.9.0
dm-tree>=0.1.10
docstring-parser>=0.18.0
drjax>=0.2.1
editdistance>=0.8.1
einops>=0.8.2
einshape>=1.0
etils>=1.14.0
evaluate>=0.2.0
Expand Down Expand Up @@ -94,7 +90,7 @@ kiwisolver>=1.5.0
latex2sympy2-extended>=1.11.0
libclang>=18.1.1
libcst>=1.9.0
lxml>=6.1.2
lxml>=6.1.3
markdown>=3.10.2
markdown-it-py>=4.2.0
markupsafe>=3.0.3
Expand Down Expand Up @@ -134,12 +130,11 @@ parameterized>=0.9.0
pathspec>=1.1.1
pathwaysutils>=0.1.11
pillow>=12.3.0
platformdirs>=4.11.5
platformdirs>=4.11.7
pluggy>=1.6.0
portpicker>=1.6.0
pre-commit>=4.6.2
prometheus-client>=0.26.0
promise>=2.3
propcache>=0.5.2
protobuf>=6.33.6
psutil>=7.2.2
Expand All @@ -153,7 +148,6 @@ pydantic>=2.13.5
pydantic-core>=2.46.5
pydot>=4.0.1
pyelftools>=0.33
pyglove>=0.4.5
pygments>=2.21.0
pyink>=26.5.1
pylint>=4.0.8
Expand All @@ -168,7 +162,7 @@ pytokens>=0.4.1
pytype>=2024.10.11
pyyaml>=6.0.3
qwix>=0.1.8
regex>=2026.8.31
regex>=2026.9.3
requests>=2.34.2
requests-oauthlib>=2.0.0
responses>=0.18.0
Expand All @@ -177,11 +171,9 @@ safetensors>=0.8.0
scipy>=1.18.0
scipy-stubs>=1.18.0.1
sentencepiece>=0.2.2
seqio>=0.0.20
setuptools>=84.0.0
shellingham>=1.5.4
shortuuid>=1.0.13
simple-parsing>=0.1.9
simplejson>=4.1.2
six>=1.17.0
sniffio>=1.3.1
Expand All @@ -190,14 +182,8 @@ starlette>=1.6.0
sympy>=1.12
tabulate>=0.10.0
tenacity>=9.1.4
tensorboard>=2.20.0
tensorboard-data-server>=0.7.0
tensorboard-plugin-profile>=2.13.0
tensorboardx>=2.6.5
tensorflow>=2.20.0
tensorflow-datasets>=4.9.10
tensorflow-metadata>=1.21.0
tensorflow-text>=2.20.1
tensorstore>=0.1.85
termcolor>=3.3.0
tiktoken>=0.14.0
Expand All @@ -218,12 +204,11 @@ uritemplate>=4.2.0
urllib3>=2.7.0
uvicorn>=0.52.4
uvloop>=0.22.1
virtualenv>=21.7.7
virtualenv>=21.7.8
wadler-lindig>=0.1.7
websockets>=16.1.1
werkzeug>=3.1.8
wheel>=0.48.0
wrapt>=2.4.0
xxhash>=4.0.1
yarl>=1.24.5
zipp>=4.1.0
Expand Down
Loading
Loading