[compiled autograd] move inputs to cuda with non_blocking=True

[xmfan]

Stack from ghstack (oldest at bottom):

• [compiled autograd] flatten runtime inputs with fast path #129116
• -> [compiled autograd] move inputs to cuda with non_blocking=True #129181
• [compiled autograd] treat input params as static #128987
• [compiled autograd] use in_compiled_autograd_region instead of compiled_autograd_enabled_count #128982
• [aot autograd] collect static parameter metadata when graphs fallback to inference #128905
• [compiled autograd] update benchmarks to use cli flags for fullgraph/dynamic #127960

non_blocking=True requires first pinning, which shouldn't be a problem given that they are cpu scalars

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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! (8 Unrelated Failures)

As of commit b999cb8 with merge base 9d06e37 ():

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

• pull / linux-jammy-py3.8-gcc11 / test (distributed, 1, 2, linux.2xlarge) (gh) (similar failure)
  Process completed with exit code 1.

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

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (aot_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128903)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128902)
  hf_BigBird
• inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128901)
  hf_BigBird
• inductor / rocm6.1-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2, unstable) (gh) (#128871)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (aot_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128929)
  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, unstable) (gh) (#128931)
  hf_BigBird
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / test (dynamo_eager_torchbench, 1, 2, linux.g5.4xlarge.nvidia.gpu, unstable) (gh) (#128932)
  hf_BigBird

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

[eellison]

looks good, one question

[eellison]

Should we put a numel() limit on this ? The general guidance for pin memory is to avoid overallocating.

[xmfan]

i guess it doesn't matter rn since every input in runtime_inputs_to_move would have numel=1