Skip to content

[DRAFT] Evaluate feasability of using FunctionalTensor for Example Value #155606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[DRAFT] Evaluate feasability of using FunctionalTensor for Example Value #155606

wants to merge 1 commit into from
+863 −95

Conversation

Lucaskabela
Copy link
Contributor

@Lucaskabela Lucaskabela commented Jun 10, 2025 •
edited by pytorch-bot bot
Loading

WIP - gathering initial signals

To replace one view and mutation metadata pass

Key test:

These 4 cases we were most concerned with can be ran with

python test.py

python test/functorch/test_aotdispatch.py

Results in

----------------------------------------------------------------------
Ran 2131 tests in 2267.300s

FAILED (failures=53, errors=49, skipped=450, expected failures=60)

Sample failure:

    python test/functorch/test_aotdispatch.py TestAOTExport.test_aot_export_with_torch_cond

Like

    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/convert_frame.py", line 784, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/convert_frame.py", line 823, in _compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/bytecode_transformation.py", line 1424, in transform_code_object
    transformations(instructions, code_options)
  File "/home/lucaskabela/pytorch/torch/_dynamo/convert_frame.py", line 266, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/convert_frame.py", line 744, in transform
    tracer.run()
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 3488, in run
    super().run()
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 1359, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 1263, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 831, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 2221, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/lucaskabela/pytorch/torch/_dynamo/symbolic_convert.py", line 1189, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/variables/lazy.py", line 201, in realize_and_forward
    return getattr(self.realize(), name)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 75, in graph_break_as_hard_error
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 1086, in call_function
    (true_r, true_treespec, true_graph, true_lifted_freevars) = speculate_branch(
                                                                ^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 1061, in speculate_branch
    ) = speculate_subgraph(
        ^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 823, in speculate_subgraph
    aliasing_info = subtracer.has_aliasing()
                    ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_dynamo/output_graph.py", line 3065, in has_aliasing
    example_value = _collect_fake_inputs([node])[0]
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lucaskabela/pytorch/torch/_higher_order_ops/utils.py", line 372, in _collect_fake_inputs
    assert isinstance(val, FakeTensor)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @jerryzh168 @gujinghui @PenghuiCheng @jianyuh @min-jean-cho @yanbing-j @Guobing-Chen @Xia-Weiwen @snadampal @mcarilli @ptrblck @leslie-fang-intel @EikanWang @voznesenskym @penguinwu @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov

@Lucaskabela Lucaskabela requested a review from bdhirsh as a code owner June 10, 2025 20:55
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jun 10, 2025 •
edited
Loading

🔗 Helpful Links

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

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

❌ 72 New Failures, 1 Pending, 3 Unrelated Failures

As of commit 9cee5f7 with merge base 9e5f4a8 (image):

NEW FAILURES - The following jobs have failed:

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

UNSTABLE - The following jobs are marked as unstable, possibly due to flakiness on trunk:

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

@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 2 times, most recently from 853dce0 to a86b126 Compare June 10, 2025 23:42
@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 3 times, most recently from fa32d14 to 50d5189 Compare June 12, 2025 20:46
@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 2 times, most recently from 776779f to 42df2fe Compare June 16, 2025 23:26
@Lucaskabela Lucaskabela requested a review from mruberry as a code owner June 16, 2025 23:26
@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch from 42df2fe to 885557f Compare June 16, 2025 23:32
@Lucaskabela Lucaskabela marked this pull request as draft June 16, 2025 23:34
Lucaskabela
Lucaskabela commented Jun 17, 2025
View reviewed changes
torch/_dynamo/utils.py
@@ -3245,7 +3246,7 @@ def get_fake_value(node, tx, allow_non_graph_fake=False):
)

try:
with tx.fake_mode, enable_python_dispatcher():
with tx.functional_mode, tx.fake_mode, enable_python_dispatcher():
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

run_node has issues here for self.bn - namely, several lines above we perform a deep copy on the nnmodule for BatchNorm (

pytorch/torch/_dynamo/utils.py

Line 3234 in 12b0213

nnmodule = deepcopy_to_fake_tensor(nnmodule, tx.fake_mode)
)

There are seemingly two conflicting things here:

  1. the example_value in node extracted here needs to be functional to record potential mutations when running the node, and a functional tensor can only be acted on by other functional tensors
  2. the deep copy linked above does not support functional tensors for the Parameter's (such as running_mean) as the torch__dispatch__ does not support Parameter types

So the seeming solution would be to enable deep copy on Parameter type for functional tensors, but since these are parameters aren't they inheritly stateful?

@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 6 times, most recently from 585c463 to 638a21c Compare June 24, 2025 17:48
@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 2 times, most recently from c6fa6cf to a9806f5 Compare June 26, 2025 22:09
@pytorch-bot pytorch-bot bot added ciflow/h100-symm-mem ciflow/linux-aarch64 linux aarch64 CI workflow ciflow/mps Run MPS tests (subset of trunk) module: amp (automated mixed precision) autocast module: cpu CPU specific problem (e.g., perf, algorithm) module: mkldnn Related to Intel IDEEP or oneDNN (a.k.a. mkldnn) integration NNC oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: quantization release notes category release notes: distributed (checkpoint) release notes: inductor (aoti) labels Jul 7, 2025
@linux-foundation-easycla Linux Foundation: EasyCLA
Copy link

linux-foundation-easycla bot commented Jul 7, 2025 •
edited
Loading

CLA Signed

The committers listed above are authorized under a signed CLA.

  • ✅ login: Lucaskabela / name: Lucas Kabela (9cee5f7)

@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch 5 times, most recently from b92f98d to 229a574 Compare July 7, 2025 23:45
@Lucaskabela Lucaskabela force-pushed the lucaskabela/functional_in_dynamo branch from 229a574 to 9cee5f7 Compare July 8, 2025 17:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@bdhirsh bdhirsh Awaiting requested review from bdhirsh bdhirsh is a code owner

@bobrenjc93 bobrenjc93 Awaiting requested review from bobrenjc93 bobrenjc93 is a code owner

@laithsakka laithsakka Awaiting requested review from laithsakka laithsakka is a code owner

@zou3519 zou3519 Awaiting requested review from zou3519 zou3519 is a code owner

@anijain2305 anijain2305 Awaiting requested review from anijain2305

@mruberry mruberry Awaiting requested review from mruberry

@ezyang ezyang Awaiting requested review from ezyang ezyang will be requested when the pull request is marked ready for review ezyang is a code owner

@Chillee Chillee Awaiting requested review from Chillee Chillee will be requested when the pull request is marked ready for review Chillee is a code owner

Assignees
No one assigned
Labels
ciflow/h100-symm-mem ciflow/inductor ciflow/linux-aarch64 linux aarch64 CI workflow ciflow/mps Run MPS tests (subset of trunk) module: amp (automated mixed precision) autocast module: cpu CPU specific problem (e.g., perf, algorithm) module: dynamo module: inductor module: mkldnn Related to Intel IDEEP or oneDNN (a.k.a. mkldnn) integration NNC oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: AO frontend release notes: distributed (checkpoint) release notes: inductor (aoti) release notes: quantization release notes category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@Lucaskabela