Add train step name and setup.sh install transformer jax from pip

Signed-off-by: kunjan <kunjanp@google.com>

Author: coolkp

setup.sh

@@ -77,8 +77,7 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
         echo "Installing stable jax, jaxlib, libtpu for NVIDIA gpu"
         pip3 install "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
     fi
-    export NVTE_FRAMEWORK=jax
-    pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    pip install "transformer_engine[jax]"
   fi
 
 elif [[ $MODE == "nightly" ]]; then
@@ -88,8 +87,7 @@ elif [[ $MODE == "nightly" ]]; then
       # Install jax-nightly
       pip install -U --pre jax jaxlib jax-cuda12-plugin[with_cuda] jax-cuda12-pjrt -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
       # Install Transformer Engine
-      export NVTE_FRAMEWORK=jax
-      pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+      pip install "transformer_engine[jax]"
   elif [[ $DEVICE == "tpu" ]]; then
     echo "Installing jax-nightly,jaxlib-nightly"
     # Install jax-nightly

src/maxdiffusion/trainers/sdxl_trainer.py

@@ -109,15 +109,19 @@ def load_dataset(self, pipeline, params, train_states):
     p_vae_apply = None
     rng = None
     if config.dataset_type == "tf" and config.cache_latents_text_encoder_outputs:
-      p_encode = jax.jit(
-          partial(
+      text_encoder_partial = partial(
               maxdiffusion_utils.encode_xl,
               text_encoders=[pipeline.text_encoder, pipeline.text_encoder_2],
               text_encoder_params=[train_states["text_encoder_state"].params, train_states["text_encoder_2_state"].params],
           )
+      text_encoder_partial.__name__="Text encoder"
+      p_encode = jax.jit(
+         text_encoder_partial
       )
+      vae_partial = partial(maxdiffusion_utils.vae_apply, vae=pipeline.vae, vae_params=train_states["vae_state"].params)
+      vae_partial.__name__="VAE Partial"
       p_vae_apply = jax.jit(
-          partial(maxdiffusion_utils.vae_apply, vae=pipeline.vae, vae_params=train_states["vae_state"].params)
+         vae_partial
       )
       rng = self.rng
 
@@ -152,8 +156,10 @@ def compile_train_step(self, pipeline, params, train_states, state_shardings, da
 
     self.rng, train_rngs = jax.random.split(self.rng)
     with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      train_step_partial = partial(_train_step, pipeline=pipeline, params=params, config=self.config)
+      train_step_partial.__name__ = "Train Step"
       p_train_step = jax.jit(
-          partial(_train_step, pipeline=pipeline, params=params, config=self.config),
+          train_step_partial,
           in_shardings=(
               state_shardings["unet_state_shardings"],
               state_shardings["vae_state_shardings"],