source env.sh in train.sh

training/trillium/Mixtral-8x7B-Pytorch/XPK/train.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+source env.sh
+
 # Extract the number after '-' in TPU_TYPE
 TPU_NUM=$(echo "$TPU_TYPE" | grep -oP '(?<=-)\d+')
 
@@ -21,7 +23,7 @@ export PROFILE_DURATION_MS=100000
 export PROFILE_LOGDIR=${PROFILE_LOG_DIR}
 export XLA_PERSISTENT_CACHE_PATH=/app/xla_cache/
 export TPU_LIBRARY_PATH=/workspace/_libtpu.so
-export NUM_SLICE=${NUM_SLICE}
+export NUM_TPU_SLICE=${NUM_SLICE}
 export LIBTPU_INIT_ARGS="--xla_tpu_enable_flash_attention=false --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=81920"
 
 huggingface-cli login --token=${HF_TOKEN}