import fbgemm sparse_ops before compilation (pytorch#2144)

Ivan Kobzarev · facebook-github-bot · commit a117ae21af3d · 2024-06-21T11:35:00.000-07:00
Summary: Pull Request resolved: pytorch#2144 Importing fbgemm.sparse_ops that are needed for torchrec compilation before model.compile call. The main concern is torch.package tests. Reviewed By: gnahzg, Microve Differential Revision: D58830081 fbshipit-source-id: 3497f996c881d9c05814e779e68153b7956dc316
diff --git a/torchrec/distributed/train_pipeline/train_pipelines.py b/torchrec/distributed/train_pipeline/train_pipelines.py
@@ -243,6 +243,10 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
                 torch._dynamo.config.force_unspec_int_unbacked_size_like_on_torchrec_kjt = (
                     True
                 )
+
+                # Importing only before compilation to not slow-done train_pipelines import
+                torch.ops.import_module("fbgemm_gpu.sparse_ops")
+
                 self._model.compile(
                     fullgraph=cc.fullgraph, dynamic=cc.dynamic, backend=cc.backend
                 )

Original file line number	Diff line number	Diff line change
`@@ -243,6 +243,10 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:`
`243`	`243`	`torch._dynamo.config.force_unspec_int_unbacked_size_like_on_torchrec_kjt = (`
`244`	`244`	`True`
`245`	`245`	`)`
	`246`	`+`
	`247`	`+ # Importing only before compilation to not slow-done train_pipelines import`
	`248`	`+ torch.ops.import_module("fbgemm_gpu.sparse_ops")`
	`249`	`+`
`246`	`250`	`self._model.compile(`
`247`	`251`	`fullgraph=cc.fullgraph, dynamic=cc.dynamic, backend=cc.backend`
`248`	`252`	`)`