Make IR private (#1099)

pytorch · May 2, 2024 · 6b19ff5 · 6b19ff5
1 parent f018a52
commit 6b19ff5
Show file tree

Hide file tree

Showing 43 changed files with 45 additions and 46 deletions.
diff --git a/README.md b/README.md
@@ -109,11 +109,10 @@ mn = MyNetwork(in_dim, layer_dims).to(device)
 
 This network is written as free-form Python code; it has not been modified for any specific parallelism technique.
 
-Let us see our first usage of the `pippy.IR.Pipe` interface:
+Let us see our first usage of the `pippy.Pipe` interface:
 
 ```python
-from pippy import pipeline
-from pippy.IR import annotate_split_points, Pipe, SplitPoint
+from pippy import pipeline, annotate_split_points, Pipe, SplitPoint
 
 annotate_split_points(mn, {'layer0': SplitPoint.END,
                            'layer1': SplitPoint.END})

diff --git a/examples/checkpoint/toy_model.py b/examples/checkpoint/toy_model.py
@@ -8,7 +8,7 @@
 from pippy.compile import compile_stage
 
 from pippy.SaveModule import save_checkpoint
-from pippy.IR import pipe_split
+from pippy import pipe_split
 from torch.utils.data import Dataset, random_split
 
 

diff --git a/examples/cpu_init/gpt2_cpu_init.py b/examples/cpu_init/gpt2_cpu_init.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_albert.py b/examples/huggingface/pippy_albert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_bart.py b/examples/huggingface/pippy_bart.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_bert.py b/examples/huggingface/pippy_bert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_blenderbot.py b/examples/huggingface/pippy_blenderbot.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_camemBert.py b/examples/huggingface/pippy_camemBert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_convBert.py b/examples/huggingface/pippy_convBert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_deberta.py b/examples/huggingface/pippy_deberta.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_debertaV2.py b/examples/huggingface/pippy_debertaV2.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_distilBert.py b/examples/huggingface/pippy_distilBert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_electra.py b/examples/huggingface/pippy_electra.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_fnet.py b/examples/huggingface/pippy_fnet.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_gpt2.py b/examples/huggingface/pippy_gpt2.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_gptNeo.py b/examples/huggingface/pippy_gptNeo.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_layoutLM.py b/examples/huggingface/pippy_layoutLM.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_m2m100.py b/examples/huggingface/pippy_m2m100.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_mbart.py b/examples/huggingface/pippy_mbart.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_megatronBert.py b/examples/huggingface/pippy_megatronBert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_mobileBert.py b/examples/huggingface/pippy_mobileBert.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_mt5.py b/examples/huggingface/pippy_mt5.py
@@ -16,7 +16,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_opt.py b/examples/huggingface/pippy_opt.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_pegasus.py b/examples/huggingface/pippy_pegasus.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_plBart.py b/examples/huggingface/pippy_plBart.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_t5.py b/examples/huggingface/pippy_t5.py
@@ -17,7 +17,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_trOCR.py b/examples/huggingface/pippy_trOCR.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_unet.py b/examples/huggingface/pippy_unet.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/huggingface/pippy_xlnet.py b/examples/huggingface/pippy_xlnet.py
@@ -10,7 +10,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import SplitPoint, annotate_split_points
+from pippy import SplitPoint, annotate_split_points
 from pippy.PipelineSchedule import ScheduleGPipe
 from pippy.PipelineStage import PipelineStage
 

diff --git a/examples/profiling/mlp_profiling.py b/examples/profiling/mlp_profiling.py
@@ -10,7 +10,7 @@
 from torch.profiler import profile, ProfilerActivity
 
 from pippy.compile import compile_stage
-from pippy.IR import pipe_split
+from pippy import pipe_split
 
 
 d_hid = 1024

diff --git a/examples/tp+pp/pippy_tp.py b/examples/tp+pp/pippy_tp.py
@@ -7,7 +7,7 @@
 
 import pippy
 import pippy.fx
-from pippy.IR import pipe_split
+from pippy import pipe_split
 from pippy.compile import compile_stage
 
 import torch.distributed as dist

diff --git a/examples/unrolling/README.md b/examples/unrolling/README.md
@@ -36,7 +36,7 @@ iteration block:
 # Add a split point after each iter_block
 annotate_split_points(
     model,
-    {"iter_block": PipeSplitWrapper.SplitPoint.END},
+    {"iter_block": SplitPoint.END},
 )
 ```
 

diff --git a/examples/unrolling/pippy_unroll.py b/examples/unrolling/pippy_unroll.py
@@ -7,7 +7,7 @@
 import torch.distributed as dist
 
 from pippy import pipeline
-from pippy.IR import annotate_split_points, PipeSplitWrapper, pipeline
+from pippy import annotate_split_points, SplitPoint, pipeline
 from pippy.PipelineStage import PipelineStage
 
 
@@ -57,7 +57,7 @@ def forward(self, x):
 # Add a split point after each iter_block
 annotate_split_points(
     model,
-    {"iter_block": PipeSplitWrapper.SplitPoint.END},
+    {"iter_block": SplitPoint.END},
 )
 
 batch_size = 32

diff --git a/pippy/ModelSplit.py b/pippy/ModelSplit.py
@@ -5,7 +5,7 @@
 import torch
 import torch.fx as fx
 
-from .IR import aten_pipe_split_alias
+from ._IR import aten_pipe_split_alias
 
 
 logger = logging.getLogger(__name__)

diff --git a/pippy/PipelineSchedule.py b/pippy/PipelineSchedule.py
@@ -9,7 +9,7 @@
 import torch.distributed as dist
 from torch.profiler import record_function
 
-from .IR import Pipe
+from ._IR import Pipe
 from .microbatch import merge_chunks, split_args_kwargs_into_chunks
 from .PipelineStage import PipelineStageBase
 

diff --git a/pippy/PipelineStage.py b/pippy/PipelineStage.py
@@ -14,8 +14,8 @@
 
 from ._backward import stage_backward
 from ._debug import map_debug_info
+from ._IR import Pipe
 from ._utils import flatten_args, modify_graph_op_device
-from .IR import Pipe
 
 logger = logging.getLogger(__name__)
 

diff --git a/pippy/IR.py → pippy/_IR.py b/pippy/IR.py → pippy/_IR.py
diff --git a/pippy/__init__.py b/pippy/__init__.py
@@ -1,5 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
-from .IR import (
+from ._IR import (
     annotate_split_points,
     ArgsChunkSpec,
     KwargsChunkSpec,

diff --git a/test/hf_test.py b/test/hf_test.py
@@ -9,7 +9,7 @@
 import torch
 
 import transformers.utils.fx as fx
-from pippy.IR import (
+from pippy._IR import (
     annotate_split_points,
     MultiUseParameterConfig,
     pipeline,

diff --git a/test/local_test_c10d_ddp.py b/test/local_test_c10d_ddp.py
@@ -7,7 +7,7 @@
 
 import torch
 import torch.distributed as dist
-from pippy.IR import pipe_split
+from pippy._IR import pipe_split
 from torch.nn.parallel import DistributedDataParallel
 
 

diff --git a/test/local_test_checkpoint.py b/test/local_test_checkpoint.py
@@ -10,9 +10,9 @@
 
 import torch.distributed as dist
 import torch.optim as optim
-from pippy.compile import compile_stage
 
-from pippy.IR import pipe_split, TrivialLossWrapper
+from pippy._IR import pipe_split, TrivialLossWrapper
+from pippy.compile import compile_stage
 from pippy.utilities.hf_checkpoint import load_checkpoint, save_checkpoint
 
 

diff --git a/test/local_test_null_coalesce_accumulate.py b/test/local_test_null_coalesce_accumulate.py
@@ -7,7 +7,7 @@
 
 import torch
 from pippy import run_pippy
-from pippy.IR import (
+from pippy._IR import (
     _null_coalesce_accumulate,
     pipe_split,
     pipeline,

diff --git a/test/test_ir.py b/test/test_ir.py
@@ -7,7 +7,7 @@
 
 import torch
 
-from pippy.IR import (
+from pippy._IR import (
     _null_coalesce_accumulate,
     annotate_split_points,
     MultiUseParameterConfig,