test tree_util changes

.github/container/Dockerfile.axlearn

@@ -1,6 +1,6 @@
 # syntax=docker/dockerfile:1-labs
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
-ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#main
+ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#sbosisio/tree_util
 ARG SRC_PATH_AXLEARN=/opt/axlearn
 
 ###############################################################################

.github/eks-workflow-files/axlearn/axlearn-fuji-model.yml

@@ -25,6 +25,8 @@ spec:
                       mkdir -p ${TRAINER_DIR}
                       OUTPUT_LOG_FILE=${TRAINER_DIR}/output.log
 
+                      export XLA_FLAGS="--xla_gpu_enable_command_buffer="
+
                       python3 /usr/local/bin/fuji-train-perf.py \
                           --module=text.gpt.c4_trainer \
                           --config=${CONFIG} \