Update

[ghstack-poisoned]

create_seed_checkpoint.sh

@@ -5,29 +5,32 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+#
+# create_seed_checkpoint.sh
+#
+# Run this script to create a seed checkpoint used to initialize a model from step-0.
+# Seed checkpoints are used to initialize pipeline-parallel models since the model initializer
+# functions don't cleanly run on chunked model parts after meta-initialization.
+#
+# Use the same model config to generate your seed checkpoint as you use for training.
+# e.g.
+# CONFIG=<path to model_config> ./create_seed_checkpoint.sh
+
 set -ex
 
-# libUV is a scalable backend for TCPStore which is used in processGroup
-# rendezvous. This is the recommended backend for distributed training.
 export USE_LIBUV=1
 TRAINER_DIR=${1:-/home/$USER/local/torchtitan}
-
-# use envs as local overrides for convenience
-# e.g.
-# LOG_RANK=0,1 NGPU=4 ./run_llama_train.sh
-
 NGPU=1
 LOG_RANK=0
-
-
 CONFIG_FILE=${CONFIG_FILE:-"./train_configs/debug_model.toml"}
 
 seed_checkpoint="--checkpoint.enable_checkpoint --checkpoint.create_seed_checkpoint"
+force_1d="--training.data_parallel_degree 1 --training.tensor_parallel_degree 1 --training.pipeline_parallel_degree 1"
 overrides=""
 if [ $# -ne 0 ]; then
     overrides="$*"
 fi
 
 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} $seed_checkpoint $overrides
+train.py --job.config_file ${CONFIG_FILE} $seed_checkpoint $force_1d $overrides

train.py

@@ -405,10 +405,10 @@ def loss_fn(pred, labels):
 
     metric_logger.close()
     logger.info("Training completed")
-    destroy_process_group()
 
 
 if __name__ == "__main__":
     config = JobConfig()
     config.parse_args()
     main(config)
+    destroy_process_group()