clean up

Author: H-Huang

run_train.sh

@@ -10,11 +10,8 @@ set -ex
 # use envs as local overwrites for convenience
 # e.g.
 # LOG_RANK=0,1 NGPU=4 ./run_train.sh
-# NGPU=${NGPU:-"8"}
-NGPU=${NGPU:-"4"}
-# export LOG_RANK=${LOG_RANK:-0,1,2,3,4,5,6,7}
-# export LOG_RANK=${LOG_RANK:-0,1,2,3}
-export LOG_RANK=${LOG_RANK:-0,1,2,3}
+NGPU=${NGPU:-"8"}
+export LOG_RANK=${LOG_RANK:-0}
 CONFIG_FILE=${CONFIG_FILE:-"./torchtitan/models/llama3/train_configs/debug_model.toml"}
 TRAIN_FILE=${TRAIN_FILE:-"torchtitan.train"}
 

torchtitan/tools/profiling.py

@@ -15,7 +15,7 @@
 from torchtitan.tools.logging import logger
 
 # the number of warmup steps before the active step in each profiling cycle
-WARMUP = 0
+WARMUP = 3
 
 # how much memory allocation/free ops to record in memory snapshots
 MEMORY_SNAPSHOT_MAX_ENTRIES = 100000