expose option to collect sizes as dynamic (#141153) (#2594)

xmfan · facebook-github-bot · commit 677cb0d60e11 · 2024-11-25T18:31:17.000-08:00
Summary: Pull Request resolved: #2594 This is to address recompiles from eager nodes that saved dynamic activations X-link: pytorch/pytorch#141153 Approved by: https://github.com/jansel ghstack dependencies: #141152 Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/db4e8a1d8a888bc3c407c13ba227f3f859349be3 Reviewed By: izaitsevfb Differential Revision: D66384254 Pulled By: xmfan fbshipit-source-id: 89d439353219252c9f5be0b44cf67957092860e3
diff --git a/torchrec/distributed/train_pipeline/train_pipelines.py b/torchrec/distributed/train_pipeline/train_pipelines.py
@@ -1624,7 +1624,7 @@ def get_compiled_autograd_ctx(
             self.initialized = True
             return contextlib.nullcontext()
 
-        return torch._dynamo.compiled_autograd.enable(
+        return torch._dynamo.compiled_autograd._enable(
             # pyre-ignore
             torch.compile(**self.compiled_autograd_options)
         )

Original file line number	Diff line number	Diff line change
`@@ -1624,7 +1624,7 @@ def get_compiled_autograd_ctx(`
`1624`	`1624`	`self.initialized = True`
`1625`	`1625`	`return contextlib.nullcontext()`
`1626`	`1626`
`1627`		`- return torch._dynamo.compiled_autograd.enable(`
	`1627`	`+ return torch._dynamo.compiled_autograd._enable(`
`1628`	`1628`	`# pyre-ignore`
`1629`	`1629`	`torch.compile(**self.compiled_autograd_options)`
`1630`	`1630`	`)`