├── .coverage ├── .flake8 ├── .github └── workflows │ ├── model_tests.yaml │ └── wait_for_ssh_to_drain.sh ├── .gitignore ├── .gitmodules ├── ARCHITECTURE.md ├── CITATION ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── REFERENCE.md ├── binaries ├── build.py └── upload.py ├── check.sh ├── docs ├── .gitignore ├── Makefile ├── README.md ├── requirements.txt └── source │ ├── _static │ └── css │ │ └── custom.css │ ├── _templates │ ├── autosummary │ │ └── class.rst │ ├── classtemplate.rst │ └── layout.html │ ├── conf.py │ ├── docutils.conf │ └── index.rst ├── examples ├── basic │ ├── example.py │ ├── example_manual_stage.py │ └── example_train.py ├── checkpoint │ └── toy_model.py ├── cpu_init │ ├── README.md │ └── gpt2_cpu_init.py ├── huggingface │ ├── hf_utils.py │ ├── pippy_bert.py │ ├── pippy_blenderbot.py │ ├── pippy_camemBert.py │ ├── pippy_convBert.py │ ├── pippy_deberta.py │ ├── pippy_debertaV2.py │ ├── pippy_distilBert.py │ ├── pippy_electra.py │ ├── pippy_fnet.py │ ├── pippy_gpt2.py │ ├── pippy_gptNeo.py │ ├── pippy_layoutLM.py │ ├── pippy_mbart.py │ ├── pippy_megatronBert.py │ ├── pippy_mobileBert.py │ ├── pippy_opt.py │ ├── pippy_trOCR.py │ ├── pippy_unet.py │ └── pippy_xlnet.py ├── llama │ ├── README.md │ └── pippy_llama.py ├── mixture_of_experts │ └── dist_moe.py ├── profiling │ └── mlp_profiling.py ├── tp+pp │ └── pippy_tp.py └── unrolling │ ├── README.md │ └── pippy_unroll.py ├── format.sh ├── pippy ├── ManualPipelineStage.py ├── ModelSplit.py ├── PipelineSchedule.py ├── _IR.py ├── _PipelineStage.py ├── __init__.py ├── _backward.py ├── _debug.py ├── _unflatten.py ├── _utils.py ├── graphsplit.py ├── microbatch.py └── utilities │ ├── __init__.py │ └── hf_checkpoint.py ├── pyproject.toml ├── requirements.txt ├── setup.py ├── test ├── __init__.py ├── hf_test.py ├── multinode_trainer.slurm ├── test_autosplit.py ├── test_bwd.py ├── test_chunkspec.py ├── test_composability.py ├── test_cpu_init.py ├── test_fwd.py ├── test_grad.py ├── test_interleave.py ├── test_ir.py ├── test_microbatch.py ├── test_optim.py ├── test_pipe.py ├── test_pipe_bwd.py ├── test_pipeline_schedule.py ├── test_pipeline_schedule_e2e.py ├── test_pipeline_stage.py ├── test_skip_conn.py ├── test_stage_backward.py ├── test_transformer.py └── test_unflatten.py └── version.txt /.coverage: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.coverage -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.flake8 -------------------------------------------------------------------------------- /.github/workflows/model_tests.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.github/workflows/model_tests.yaml -------------------------------------------------------------------------------- /.github/workflows/wait_for_ssh_to_drain.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.github/workflows/wait_for_ssh_to_drain.sh -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/.gitmodules -------------------------------------------------------------------------------- /ARCHITECTURE.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/ARCHITECTURE.md -------------------------------------------------------------------------------- /CITATION: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/CITATION -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/CODE_OF_CONDUCT.md -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/README.md -------------------------------------------------------------------------------- /REFERENCE.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/REFERENCE.md -------------------------------------------------------------------------------- /binaries/build.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/binaries/build.py -------------------------------------------------------------------------------- /binaries/upload.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/binaries/upload.py -------------------------------------------------------------------------------- /check.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/check.sh -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/.gitignore -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/requirements.txt -------------------------------------------------------------------------------- /docs/source/_static/css/custom.css: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/_static/css/custom.css -------------------------------------------------------------------------------- /docs/source/_templates/autosummary/class.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/_templates/autosummary/class.rst -------------------------------------------------------------------------------- /docs/source/_templates/classtemplate.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/_templates/classtemplate.rst -------------------------------------------------------------------------------- /docs/source/_templates/layout.html: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/_templates/layout.html -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/conf.py -------------------------------------------------------------------------------- /docs/source/docutils.conf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/docutils.conf -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/docs/source/index.rst -------------------------------------------------------------------------------- /examples/basic/example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/basic/example.py -------------------------------------------------------------------------------- /examples/basic/example_manual_stage.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/basic/example_manual_stage.py -------------------------------------------------------------------------------- /examples/basic/example_train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/basic/example_train.py -------------------------------------------------------------------------------- /examples/checkpoint/toy_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/checkpoint/toy_model.py -------------------------------------------------------------------------------- /examples/cpu_init/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/cpu_init/README.md -------------------------------------------------------------------------------- /examples/cpu_init/gpt2_cpu_init.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/cpu_init/gpt2_cpu_init.py -------------------------------------------------------------------------------- /examples/huggingface/hf_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/hf_utils.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_bert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_bert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_blenderbot.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_blenderbot.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_camemBert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_camemBert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_convBert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_convBert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_deberta.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_deberta.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_debertaV2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_debertaV2.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_distilBert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_distilBert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_electra.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_electra.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_fnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_fnet.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_gpt2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_gpt2.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_gptNeo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_gptNeo.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_layoutLM.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_layoutLM.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_mbart.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_mbart.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_megatronBert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_megatronBert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_mobileBert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_mobileBert.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_opt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_opt.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_trOCR.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_trOCR.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_unet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_unet.py -------------------------------------------------------------------------------- /examples/huggingface/pippy_xlnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/huggingface/pippy_xlnet.py -------------------------------------------------------------------------------- /examples/llama/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/llama/README.md -------------------------------------------------------------------------------- /examples/llama/pippy_llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/llama/pippy_llama.py -------------------------------------------------------------------------------- /examples/mixture_of_experts/dist_moe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/mixture_of_experts/dist_moe.py -------------------------------------------------------------------------------- /examples/profiling/mlp_profiling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/profiling/mlp_profiling.py -------------------------------------------------------------------------------- /examples/tp+pp/pippy_tp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/tp+pp/pippy_tp.py -------------------------------------------------------------------------------- /examples/unrolling/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/unrolling/README.md -------------------------------------------------------------------------------- /examples/unrolling/pippy_unroll.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/examples/unrolling/pippy_unroll.py -------------------------------------------------------------------------------- /format.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/format.sh -------------------------------------------------------------------------------- /pippy/ManualPipelineStage.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/ManualPipelineStage.py -------------------------------------------------------------------------------- /pippy/ModelSplit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/ModelSplit.py -------------------------------------------------------------------------------- /pippy/PipelineSchedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/PipelineSchedule.py -------------------------------------------------------------------------------- /pippy/_IR.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_IR.py -------------------------------------------------------------------------------- /pippy/_PipelineStage.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_PipelineStage.py -------------------------------------------------------------------------------- /pippy/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/__init__.py -------------------------------------------------------------------------------- /pippy/_backward.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_backward.py -------------------------------------------------------------------------------- /pippy/_debug.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_debug.py -------------------------------------------------------------------------------- /pippy/_unflatten.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_unflatten.py -------------------------------------------------------------------------------- /pippy/_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/_utils.py -------------------------------------------------------------------------------- /pippy/graphsplit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/graphsplit.py -------------------------------------------------------------------------------- /pippy/microbatch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/microbatch.py -------------------------------------------------------------------------------- /pippy/utilities/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /pippy/utilities/hf_checkpoint.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pippy/utilities/hf_checkpoint.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/pyproject.toml -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch >= 2.3.0.dev 2 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/setup.py -------------------------------------------------------------------------------- /test/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | -------------------------------------------------------------------------------- /test/hf_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/hf_test.py -------------------------------------------------------------------------------- /test/multinode_trainer.slurm: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/multinode_trainer.slurm -------------------------------------------------------------------------------- /test/test_autosplit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_autosplit.py -------------------------------------------------------------------------------- /test/test_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_bwd.py -------------------------------------------------------------------------------- /test/test_chunkspec.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_chunkspec.py -------------------------------------------------------------------------------- /test/test_composability.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_composability.py -------------------------------------------------------------------------------- /test/test_cpu_init.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_cpu_init.py -------------------------------------------------------------------------------- /test/test_fwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_fwd.py -------------------------------------------------------------------------------- /test/test_grad.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_grad.py -------------------------------------------------------------------------------- /test/test_interleave.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_interleave.py -------------------------------------------------------------------------------- /test/test_ir.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_ir.py -------------------------------------------------------------------------------- /test/test_microbatch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_microbatch.py -------------------------------------------------------------------------------- /test/test_optim.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_optim.py -------------------------------------------------------------------------------- /test/test_pipe.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_pipe.py -------------------------------------------------------------------------------- /test/test_pipe_bwd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_pipe_bwd.py -------------------------------------------------------------------------------- /test/test_pipeline_schedule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_pipeline_schedule.py -------------------------------------------------------------------------------- /test/test_pipeline_schedule_e2e.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_pipeline_schedule_e2e.py -------------------------------------------------------------------------------- /test/test_pipeline_stage.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_pipeline_stage.py -------------------------------------------------------------------------------- /test/test_skip_conn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_skip_conn.py -------------------------------------------------------------------------------- /test/test_stage_backward.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_stage_backward.py -------------------------------------------------------------------------------- /test/test_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_transformer.py -------------------------------------------------------------------------------- /test/test_unflatten.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/pytorch/PiPPy/HEAD/test/test_unflatten.py -------------------------------------------------------------------------------- /version.txt: -------------------------------------------------------------------------------- 1 | 0.2.0 2 | --------------------------------------------------------------------------------