Copybara import of gpu-recipes:

  - 09f6608a72d36044cfcfd4ff45b623ec6eab8a09 Update image and remove xla flag

GitOrigin-RevId: 09f6608a72d36044cfcfd4ff45b623ec6eab8a09

Author: Copybara

training/a4high/llama-3.1-405b/maxtext-pretraining-gke/README.md

@@ -113,7 +113,7 @@ default settings, run the following command from your client:
 cd $RECIPE_ROOT
 helm install -f values.yaml \
     --set-file maxtext_config=$REPO_ROOT/src/frameworks/a4high/maxtext-configs/llama-3.1-405b-256gpus-a4h-fp8.yaml \
-    --set workload.image=us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150313  \
+    --set workload.image=us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317  \
     --set workload.run_name=$USER-llama-3-1-405b-maxtext-fp8 \
     --set workload.gpus=256 \
     --set queue=$KUEUE_NAME \
@@ -132,7 +132,7 @@ helm install -f values.yaml \
 cd $RECIPE_ROOT
 helm install -f values.yaml \
     --set-file maxtext_config=$REPO_ROOT/src/frameworks/a4high/maxtext-configs/llama-3.1-405b-256gpus-a4h-fp8.yaml \
-    --set workload.image=us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150313  \
+    --set workload.image=us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317  \
     --set workload.run_name=$USER-llama-3-1-405b-maxtext-fp8 \
     --set workload.gpus=256 \
     --set queue=$KUEUE_NAME \

training/a4high/llama-3.1-405b/maxtext-pretraining-gke/values.yaml

@@ -41,7 +41,6 @@ xlaFlags: >-
   --xla_gpu_enable_latency_hiding_scheduler=true
   --xla_gpu_enable_triton_gemm=false
   --xla_gpu_enable_command_buffer=FUSION,CUSTOM_CALL
-  --xla_gpu_enable_highest_priority_async_stream=true
   --xla_gpu_all_reduce_combine_threshold_bytes=17179869184
   --xla_gpu_all_gather_combine_threshold_bytes=17179869184
   --xla_gpu_reduce_scatter_combine_threshold_bytes=17179869184