Fix the Tensorflow version (#154)

* Remove the transformer engine from setup script since it's not required and it's slowing down the setup.

* Revert the Transformer Engine change. Fix the tensorflow version.

Author: wang2yn84

docs/getting_started/first_run.md

@@ -12,7 +12,11 @@ multiple hosts.
 1. Clone MaxDiffusion in your TPU VM.
 1. Within the root directory of the MaxDiffusion `git` repo, install dependencies by running:
 ```bash
-bash setup.sh MODE=stable
+If you are running on TPU:
+bash setup.sh MODE=stable DEVICE=tpu
+
+If you are running on GPU:
+bash setup.sh MODE=stable DEVICE=gpu
 ```
 
 ## Getting Starting: Multihost development

requirements.txt

@@ -19,7 +19,7 @@ Pillow
 pylint
 pyink
 pytest==8.2.2
-tensorflow>=2.17.0
+tensorflow==2.17.0
 tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git
@@ -29,4 +29,4 @@ tokenizers==0.21.0
 huggingface_hub==0.24.7
 transformers==4.48.1
 einops==0.8.0
-sentencepiece
+sentencepiece

setup.sh

@@ -58,7 +58,7 @@ fi
 # Install JAX and JAXlib based on the specified mode
 if [[ "$MODE" == "stable" || ! -v MODE ]]; then
   # Stable mode
-  if [[ $DEVICE == "tpu" ]]; then 
+  if [[ $DEVICE == "tpu" ]]; then
     echo "Installing stable jax, jaxlib for tpu"
     if [[ -n "$JAX_VERSION" ]]; then
       echo "Installing stable jax, jaxlib, libtpu version ${JAX_VERSION}"
@@ -90,14 +90,14 @@ elif [[ $MODE == "nightly" ]]; then
       # Install Transformer Engine
       export NVTE_FRAMEWORK=jax
       pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@stable
-  elif [[ $DEVICE == "tpu" ]]; then 
+  elif [[ $DEVICE == "tpu" ]]; then
     echo "Installing jax-nightly,jaxlib-nightly"
     # Install jax-nightly
     pip3 install --pre -U jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
     # Install jaxlib-nightly
     pip3 install --pre -U jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
     # Install libtpu-nightly
-    pip3 install --pre -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 
+    pip3 install --pre -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   fi
   echo "Installing nightly tensorboard plugin profile"
   pip3 install tbp-nightly --upgrade
@@ -107,7 +107,7 @@ else
 fi
 
 # Install dependencies from requirements.txt
-pip3 install -U -r requirements.txt
+pip3 install -U -r requirements.txt || echo "Failed to install dependencies in the requirements" >&2
 
 # Install maxdiffusion
-pip3 install -U .
+pip3 install -U . || echo "Failed to install maxdiffusion" >&2