Integrate PT2 pipeline with FullModelCompilation Config (pytorch#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

@@ -12,6 +12,7 @@
     EvalPipelineSparseDist,  # noqa
     PrefetchTrainPipelineSparseDist,  # noqa
     StagedTrainPipeline,  # noqa
+    TorchCompileConfig,  # noqa
     TrainPipeline,  # noqa
     TrainPipelineBase,  # noqa
     TrainPipelinePT2,  # noqa