Use expandable segments in run_llama_train.sh

ghstack-source-id: e8adf491c3a7d2da4a8674bb78b18e793b75a934 Pull Request resolved: pytorch#643
mori360 · Oct 22, 2024 · a5e57e4 · a5e57e4
1 parent e211236
commit a5e57e4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/run_llama_train.sh b/run_llama_train.sh
@@ -19,6 +19,7 @@ if [ $# -ne 0 ]; then
     overrides="$*"
 fi
 
+PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True" \
 torchrun --nproc_per_node=${NGPU} --rdzv_backend c10d --rdzv_endpoint="localhost:0" \
 --local-ranks-filter ${LOG_RANK} --role rank --tee 3 \
 train.py --job.config_file ${CONFIG_FILE} $overrides