From be464c16eb2b5ac76353d08dbf0945b26638637b Mon Sep 17 00:00:00 2001 From: Shuai Yang Date: Tue, 8 Oct 2024 14:13:23 -0700 Subject: [PATCH] Integrate PT2 pipeline with FullModelCompilation Config (#2475) Summary: In current FMC setting, we start compilation from the first iteration. However, when we do sparse arch compilation, torchrec needs a few iterations of dry runs before compilation. So we need to postpone the compilation to happen in the torchrec pipeline. Despite the difference, we aim to provide an uniform interface for users. Users can set `FullModelCompilation Config` as usual and if users also use `training.pipeline_type=pt2`, this diff can handle the postpone mechanism automatically. Differential Revision: D63500331 --- torchrec/distributed/train_pipeline/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchrec/distributed/train_pipeline/__init__.py b/torchrec/distributed/train_pipeline/__init__.py index e49f7ff98..d7b38d2b0 100644 --- a/torchrec/distributed/train_pipeline/__init__.py +++ b/torchrec/distributed/train_pipeline/__init__.py @@ -12,6 +12,7 @@ EvalPipelineSparseDist, # noqa PrefetchTrainPipelineSparseDist, # noqa StagedTrainPipeline, # noqa + TorchCompileConfig, # noqa TrainPipeline, # noqa TrainPipelineBase, # noqa TrainPipelinePT2, # noqa