Update

[ghstack-poisoned]
pytorch · May 23, 2024 · 44ea24c · 44ea24c
2 parents 9884f42 + 17b7bec
commit 44ea24c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/multinode_trainer.slurm b/multinode_trainer.slurm
@@ -54,7 +54,7 @@ export NCCL_BUFFSIZE=2097152
 #export TORCH_DIST_INIT_BARRIER=1
 export FI_EFA_SET_CUDA_SYNC_MEMOPS=0
 #export USE_LIBUV=1
-CONFIG_FILE=${CONFIG_FILE:-"./train_configs/llama_13b.toml"}
+CONFIG_FILE=${CONFIG_FILE:-"./train_configs/llama2_13b.toml"}
 
 dcgmi profile --pause
 # adjust sbatch --ntasks and sbatch --nodes above and --nnodes below

diff --git a/torchtitan/parallelisms/parallelize_llama.py b/torchtitan/parallelisms/parallelize_llama.py
@@ -19,7 +19,7 @@
     CheckpointImpl,
 )
 from torch.distributed.pipelining import pipeline, SplitPoint
-from torch.distributed.pipelining._PipelineStage import (
+from torch.distributed.pipelining.PipelineStage import (
     _PipelineStage,
     ManualPipelineStage,
 )