Skip to content

Commit 047760b

Browse files
TroyGardenmeta-codesync[bot]
authored andcommitted
add train pipeline class in logging (#3486)
Summary: Pull Request resolved: #3486 # context * add TrainPipeline class logging to help debugging which train pipeline is used in the training job * this will produce the following log * the logging is done once when the train pipeline class is initiated (per rank) ``` INFO:torchrec.distributed.train_pipeline.train_pipelines:TrainPipeline class: <class 'torchrec.distributed.train_pipeline.train_pipelines.TrainPipelineSparseDist'> ``` Reviewed By: nipung90 Differential Revision: D85577739 fbshipit-source-id: 710a84cdb611cbaa2d255e40fb8464ee3826dac3
1 parent 3a6cf2e commit 047760b

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

torchrec/distributed/train_pipeline/train_pipelines.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,8 @@ def __init__(self) -> None:
110110
# pipeline state such as in foward, in backward etc, used in training recover scenarios
111111
self._state: PipelineState = PipelineState.IDLE
112112

113+
logger.info(f"TrainPipeline class: {type(self)}")
114+
113115
def sync_embeddings(
114116
self,
115117
model: torch.nn.Module,

0 commit comments

Comments
 (0)