Skip to content

Commit 75b8f84

Browse files
Microvefacebook-github-bot
authored andcommitted
Integrate PT2 pipeline with FullModelCompilation Config
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
1 parent 558afe2 commit 75b8f84

File tree

1 file changed

+1
-0
lines changed

1 file changed

+1
-0
lines changed

torchrec/distributed/train_pipeline/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
EvalPipelineSparseDist, # noqa
1313
PrefetchTrainPipelineSparseDist, # noqa
1414
StagedTrainPipeline, # noqa
15+
TorchCompileConfig, # noqa
1516
TrainPipeline, # noqa
1617
TrainPipelineBase, # noqa
1718
TrainPipelinePT2, # noqa

0 commit comments

Comments
 (0)