Add a HOP to bypass tracing of a wrapper function while tracing the wrapped function #153487

soulitzer · 2025-05-13T19:57:34Z

Stack from ghstack (oldest at bottom):

Usage:

from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper

# Your ordinary function wrapper
def my_hop_fn_impl(fn, *args, k=1, **kwargs):
    def wrapper(*args, **kwargs):
        out = fn(*args, **kwargs)
        if isinstance(out, tuple):
            return (out[0] + k,)
        return out + k

    return wrapper

# Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph
def my_hop_fn(fn, *args, k=1, **kwargs):
    return dynamo_bypassing_wrapper(
        functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs
    )

Notes:

The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn.
The wrapper_fn passed to dynamo_bypassing_wrapper should have signature Callable -> Callable

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

[ghstack-poisoned]

pytorch-bot · 2025-05-13T19:57:38Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/153487

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 194cdce with merge base 7c9d94e ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

inductor-rocm / rocm-py3.10-inductor / test (inductor, 1, 2, linux.rocm.gpu.2) (gh) (similar failure)
distributed/fsdp/test_fsdp_tp_integration.py::TestTPFSDPIntegration::test_fsdp_tp_extension_grad

This comment was automatically generated by Dr. CI and updates every 15 minutes.

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

Usage: ```python from torch._higher_order_ops.wrap import wrap_generic # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return wrap_generic( fn, *args, wrapper_fn=functools.partial(my_hop_fn_impl, k=k), **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `wrap_generic` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

torch/_dynamo/variables/higher_order_ops.py

ydwu4 · 2025-05-20T23:26:28Z

torch/_dynamo/variables/higher_order_ops.py

+            raise RuntimeError(
+                f"WrapGenericHigherOrderVariable: Unsupported function {type(func_var)}"
+            )
+        gmod_kwargs = {


Using the wrapper_fn as kwargs seems pretty fragile like it might collide with inner kwargs. Maybe we could just pass the wrapper fn as an explicit input? Something like:

DynamoBypassingWrapper(wrapper_fn, inner_fn, *inner_args, **inner_kwargs)

The proxy node we created in dynamo could make wrapper_fn a string key that can be used to access gmod.meta.

Good point, updated

Usage: ```python from torch._higher_order_ops.wrap import wrap_generic # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return wrap_generic( fn, *args, wrapper_fn=functools.partial(my_hop_fn_impl, k=k), **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `wrap_generic` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ydwu4 · 2025-05-21T16:33:46Z

torch/_higher_order_ops/wrap.py

+        import torch._dynamo  # noqa: F401
+        from torch._dynamo import disable
+
+        is_compiling = (


Is this accurate? An alternative is:

How about we do something like this: ```python def __call__(sefl, wrapper_fn_or_key, inner_fn, *inner_args, **inner_kwargs): is_compiling = isinstance(wrapper_fn_or_key, str) if is_compiling: assert isinstance(inner_fn, torch.fx.GraphModule) wrapper_fn = inner_fn.meta[wrapper_fn_or_key] eles: wrapper_fn = wrapper_fn_or_key

Thanks this looks cleaner

ydwu4 · 2025-05-21T16:34:31Z

torch/_higher_order_ops/wrap.py

@@ -87,6 +87,41 @@ def wrapper():
 wrap_with_autocast = WrapWithAutocast()




Maybe some doc?

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ydwu4 · 2025-05-21T18:11:25Z

torch/_higher_order_ops/wrap.py

@@ -87,6 +87,43 @@ def wrapper():
 wrap_with_autocast = WrapWithAutocast()


+# This HOP allows you to bypass dynamo tracing of the wrapper function while
+# still tracing the inner function.
+# The wrapper function should receive a single callable argument, and return


nit: something like "two callables, the first callable being the wrapper_fn which takes a single argument inner_fn. The second callable is inner_fn and the rest of args and kwags are for the inner fn."

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer · 2025-05-21T18:25:52Z

@pytorchbot merge

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer · 2025-05-21T20:04:50Z

@pytorchbot merge

pytorchmergebot · 2025-05-21T20:07:07Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-05-21T21:05:50Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / cuda12.4-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer · 2025-05-21T22:40:42Z

@pytorchbot merge

pytorchmergebot · 2025-05-21T22:42:48Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-05-21T23:25:22Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / test (default, 3, 3, macos-m1-stable)

Details for Dev Infra team

Raised by workflow job

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer · 2025-05-21T23:33:43Z

@pytorchbot merge

pytorchmergebot · 2025-05-21T23:36:14Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-05-21T23:57:28Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

Lint / lintrunner-noclang / linux-job

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

…acing the wrapped function" Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer · 2025-05-22T00:13:50Z

@pytorchbot merge

pytorchmergebot · 2025-05-22T00:16:01Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…rapped function (#153487) Usage: ```python from torch._higher_order_ops.wrap import dynamo_bypassing_wrapper # Your ordinary function wrapper def my_hop_fn_impl(fn, *args, k=1, **kwargs): def wrapper(*args, **kwargs): out = fn(*args, **kwargs) if isinstance(out, tuple): return (out[0] + k,) return out + k return wrapper # Calling `my_hop_fn` instead of the impl directly captures a HOP into the dynamo graph def my_hop_fn(fn, *args, k=1, **kwargs): return dynamo_bypassing_wrapper( functools.partial(my_hop_fn_impl, k=k), fn, *args, **kwargs ) ``` Notes: - The dynamo captured graph now stashes arbitrary callable objects (the wrapper_fn) - this is equivalent to what SAC does today with policy_fn. - The `wrapper_fn` passed to `dynamo_bypassing_wrapper ` should have signature `Callable -> Callable` Pull Request resolved: #153487 Approved by: https://github.com/ydwu4

Allow HOP-ifying out-of-tree functions in compile

59b2377

[ghstack-poisoned]

pytorch-bot bot added ciflow/inductor module: dynamo labels May 13, 2025

This was referenced May 13, 2025

Make test_create_graph_and_full_backward_hook_cycle more robust to unrelated warnings #153407

Open

Add AC_TRACER Infra TorchDispatchMode key #152158

Open

Add dynamo config to HOP-ify context managers #152159

Open

Update on "Allow HOP-ifying out-of-tree functions in compile"

88f91a2

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer requested a review from zou3519 as a code owner May 13, 2025 20:42

soulitzer added 2 commits May 13, 2025 17:03

Update on "Allow HOP-ifying out-of-tree functions in compile"

7be544f

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

Update on "Allow HOP-ifying out-of-tree functions in compile"

e41e5a2

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer added the topic: not user facing topic category label May 13, 2025

pytorch deleted a comment from github-actions bot May 13, 2025

Update on "Allow HOP-ifying out-of-tree functions in compile"

f169d99

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

soulitzer requested a review from ydwu4 May 16, 2025 16:33

soulitzer added 3 commits May 19, 2025 16:34

Update on "Allow HOP-ifying out-of-tree functions in compile"

0b408ec

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

Update on "Allow HOP-ifying out-of-tree functions in compile"

bc2a83c

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

ydwu4 reviewed May 20, 2025

View reviewed changes

soulitzer added 2 commits May 21, 2025 11:57

soulitzer requested a review from ydwu4 May 21, 2025 16:10

soulitzer changed the title ~~Allow HOP-ifying out-of-tree functions in compile~~ Add a HOP to bypass tracing of a wrapper function while tracing the wrapped function May 21, 2025

ydwu4 reviewed May 21, 2025

View reviewed changes

soulitzer requested a review from ydwu4 May 21, 2025 17:26

ydwu4 approved these changes May 21, 2025

View reviewed changes

pytorchmergebot removed the merging label May 21, 2025

pytorchmergebot added the merging label May 21, 2025

pytorchmergebot removed the merging label May 21, 2025

soulitzer requested a review from tugsbayasgalan as a code owner May 21, 2025 22:14

pytorchmergebot added the merging label May 21, 2025

pytorchmergebot removed the merging label May 21, 2025

pytorchmergebot added the merging label May 21, 2025

pytorchmergebot removed the merging label May 21, 2025

pytorchmergebot added the merging label May 22, 2025

pytorchmergebot closed this in f2af30f May 22, 2025

pytorchmergebot added Merged and removed merging labels May 22, 2025

github-actions bot deleted the gh/soulitzer/364/head branch June 23, 2025 02:20

		@@ -87,6 +87,41 @@ def wrapper():
		wrap_with_autocast = WrapWithAutocast()

Add a HOP to bypass tracing of a wrapper function while tracing the wrapped function #153487

Add a HOP to bypass tracing of a wrapper function while tracing the wrapped function #153487

Uh oh!

Conversation

soulitzer commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/153487

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Uh oh!

ydwu4 May 20, 2025

Choose a reason for hiding this comment

Uh oh!

soulitzer May 21, 2025

Choose a reason for hiding this comment

Uh oh!

ydwu4 May 21, 2025

Choose a reason for hiding this comment

Uh oh!

soulitzer May 21, 2025

Choose a reason for hiding this comment

Uh oh!

ydwu4 May 21, 2025

Choose a reason for hiding this comment

Uh oh!

ydwu4 May 21, 2025

Choose a reason for hiding this comment

Uh oh!

soulitzer commented May 21, 2025

Uh oh!

soulitzer commented May 21, 2025

Uh oh!

pytorchmergebot commented May 21, 2025

Merge started

Uh oh!

pytorchmergebot commented May 21, 2025

Merge failed

Uh oh!

soulitzer commented May 21, 2025

Uh oh!

pytorchmergebot commented May 21, 2025

Merge started

Uh oh!

pytorchmergebot commented May 21, 2025

Merge failed

Uh oh!

soulitzer commented May 21, 2025

Uh oh!

pytorchmergebot commented May 21, 2025

Merge started

Uh oh!

pytorchmergebot commented May 21, 2025

Merge failed

Uh oh!

soulitzer commented May 22, 2025

Uh oh!

pytorchmergebot commented May 22, 2025

Merge started

Uh oh!

Uh oh!

soulitzer commented May 13, 2025 •

edited

Loading

pytorch-bot bot commented May 13, 2025 •

edited

Loading