Update pipelining import after change on pytorch

APIs conform to the pytorch rules. This PR should be able to land safely after tonight's nightly pytorch build which includes the above PR. ghstack-source-id: c575bc7835472128c09798544caa38bf1908e5ca Pull Request resolved: #356
pytorch · May 23, 2024 · c87e8bc · c87e8bc
1 parent 638ec48
commit c87e8bc
Showing 1 changed file with 1 addition and 1 deletion.
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,
 )