diff --git a/examples/legacy/llama/pretrain_llama_7b_ptd.sh b/examples/legacy/llama/pretrain_llama_7b_ptd.sh
index 49096363..f549b629 100644
--- a/examples/legacy/llama/pretrain_llama_7b_ptd.sh
+++ b/examples/legacy/llama/pretrain_llama_7b_ptd.sh
@@ -2,6 +2,7 @@
 
 export CUDA_DEVICE_MAX_CONNECTIONS=1
 export NPU_ASD_ENABLE=0
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
 
 GPUS_PER_NODE=8
 MASTER_ADDR=localhost