Add backwards compatibility for config variable names

Author: denys-fridman

moe_pretraining/nemo/run_deepseek_v3_671b.sh

@@ -19,17 +19,17 @@ cd $(dirname $0)
 
 set -e
 
-# Vars without defaults
+# Vars without defaults (support both old and new variable names for backwards compatibility)
 #     Slurm settings
 : "${ACCOUNT:?ACCOUNT not set}"
 : "${PARTITION:?PARTITION not set}"
 
 #     Job settings
-: "${LOG_DIR:?LOG_DIR not set}"
+LOG_DIR="${LOG_DIR:-${JOB_DIR:?LOG_DIR or JOB_DIR not set}}"
 : "${IMAGE:?IMAGE not set}"
 
 #     Dataset settings
-: "${DATA_DIR:?DATA_DIR not set}"
+DATA_DIR="${DATA_DIR:-${DATADIR:-${PREPROCESSED_PATH:?DATA_DIR, DATADIR, or PREPROCESSED_PATH not set}}}"
 : "${MODEL_CKPT:?MODEL_CKPT not set}"
 
 # Vars with defaults
@@ -63,7 +63,7 @@ set -e
 #     Training settings
 : "${MAX_STEPS:=1000}"
 : "${WARMUP_STEPS:=0}"
-: "${LR:="2e-4"}"
+LR="${LR:-${MAX_LR:-"2e-4"}}"
 : "${MIN_LR:="5e-6"}"
 : "${SEED:=1234}"