[Traceable FSDP][Compiled Autograd] Add queue_callback() support

[yf225]

Adds support for Variable._execution_engine.queue_callback(), which is used in FSDP2.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

❌ 16 New Failures, 4 Unrelated Failures

As of commit 81aa586 with merge base a0429c0 ():

NEW FAILURES - The following jobs have failed:

• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamic_aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  hf_BigBird
• Lint / lintrunner-noclang / linux-job (gh)
  >>> Lint for torch/_dynamo/variables/misc.py:
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-focal-cuda12.4-py3.10-gcc9 / build (gh)
  /var/lib/jenkins/workspace/aten/src/ATen/cuda/CUDASparseDescriptors.h:119:68: error: ‘cusparseStatus_t cusparseCreateBsrsm2Info(bsrsm2Info**)’ is deprecated: The routine will be removed in the next major release [-Werror=deprecated-declarations]
• pull / linux-focal-py3.11-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_autograd.py::TestAutograd::test_reentrant_with_callbacks_both_depths
• pull / linux-focal-py3.12-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_autograd.py::TestAutograd::test_reentrant_with_callbacks_both_depths
• pull / linux-focal-py3.8-clang10 / test (default, 1, 3, linux.2xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-focal-py3.8-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_autograd.py::TestAutograd::test_reentrant_with_callbacks_both_depths
• pull / linux-jammy-py3.10-clang15-asan / test (default, 1, 6, linux.4xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / linux-jammy-py3.8-gcc11 / test (default, 1, 3, linux.2xlarge) (gh)
  inductor/test_distributed_patterns.py::DistributedPatternTests::test_module_backward_hooks_aot
• pull / win-vs2019-cpu-py3 / build (gh)
  sccache: error: couldn't connect to server

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
  hf_BigBird

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• pull / linux-focal-cuda12.4-py3.10-gcc9-sm86 / build (gh) (#127104)
  /var/lib/jenkins/workspace/aten/src/ATen/cuda/CUDASparseDescriptors.h:119:68: error: ‘cusparseStatus_t cusparseCreateBsrsm2Info(bsrsm2Info**)’ is deprecated: The routine will be removed in the next major release [-Werror=deprecated-declarations]

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

[xmfan]

do we need these final callbacks in the graph? or just need them to execute at the end. In the second case, we can execute them on the c++ side just like eager does

[yf225]

do we need these final callbacks in the graph? or just need them to execute at the end. In the second case, we can execute them on the c++ side just like eager does

yes we need Dynamo to trace through those final callbacks, so they need to be in the graph

[xmfan]

I think we can avoid the globals on OutputGraph if we construct the CompiledAutogradEngine within the graph

Would be good to double check:

• Graph break between queue_callback and exec_final_callbacks
• Memory check for tensors involved the callback