Use TE 2.6.0

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -41,7 +41,7 @@ jobs:
       - name: Print dependencies
         run: |
           pip uninstall -y transformer-engine transformer-engine-jax
-          pip install -U transformer-engine[jax]==2.5.0
+          pip install -U transformer-engine[jax]==2.6.0
           pip freeze
 
       - name: Run MaxText Training
@@ -81,20 +81,20 @@ jobs:
   #     - name: Run MaxDiffusion Training
   #       run: |
   #         # This command is adapted from your DAG for a single-slice configuration.
-  #         NVTE_FUSED_ATTN=1 pip install . && \
-  #           python -m src.maxdiffusion.train_sdxl src/maxdiffusion/configs/base_xl.yml \
-  #           hardware=gpu \
-  #           train_new_unet=true \
-  #           train_text_encoder=false \
-  #           cache_latents_text_encoder_outputs=true \
-  #           per_device_batch_size=1 \
-  #           attention=dot_product \
-  #           activations_dtype=bfloat16 \
-  #           weights_dtype=bfloat16 \
-  #           max_train_steps=200 \
-  #           enable_profiler=True \
-  #           run_name=1slice-VGpuVersion.XPK_H100_a3-maxdiffusion-jax-stable-stack-2025-09-26-04-12-02 \
-  #           output_dir=gs://rbierneni-multipod-dev/${{ github.run_id }}
+          # NVTE_FUSED_ATTN=1 pip install . && \
+          #   python -m src.maxdiffusion.train_sdxl src/maxdiffusion/configs/base_xl.yml \
+          #   hardware=gpu \
+          #   train_new_unet=true \
+          #   train_text_encoder=false \
+          #   cache_latents_text_encoder_outputs=true \
+          #   per_device_batch_size=1 \
+          #   attention=dot_product \
+          #   activations_dtype=bfloat16 \
+          #   weights_dtype=bfloat16 \
+          #   max_train_steps=200 \
+          #   enable_profiler=True \
+          #   run_name=1slice-VGpuVersion.XPK_H100_a3-maxdiffusion-jax-stable-stack-2025-09-26-04-12-02 \
+          #   output_dir=gs://rbierneni-multipod-dev/${{ github.run_id }}
 
 # jobs:
 #   build: