[FSDP] Add record_function for explicit prefetching

[awgu]

Stack from ghstack (oldest at bottom):

• -> [FSDP] Add record_function for explicit prefetching #105985
• Revert "Simplify handle indexing (#105006)" #105984

Example:

This can be particularly help when with_stack=False, in which case it is harder to tell the prefetch.

Differential Revision: D47816318

[pytorch-bot]

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• 6-7 hour queue times caused by GPU instances not being available

✅ 1 Unrelated Failure

As of commit 6fd6f82:

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

• linux-focal-rocm5.6-py3.8 / test (default, 2, 3, linux.rocm.gpu, unstable) (gh)

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

[awgu]

@pytorchbot rebase -s

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict gh/awgu/430/orig returned non-zero exit code 1

Rebasing (1/2)
Auto-merging torch/distributed/fsdp/_exec_order_utils.py
CONFLICT (content): Merge conflict in torch/distributed/fsdp/_exec_order_utils.py
Auto-merging torch/distributed/fsdp/flat_param.py
CONFLICT (content): Merge conflict in torch/distributed/fsdp/flat_param.py
error: could not apply 55bd5ac83a7... Revert "Simplify handle indexing (#105006)"
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
Could not apply 55bd5ac83a7... Revert "Simplify handle indexing (#105006)"

Raised by https://github.com/pytorch/pytorch/actions/runs/5662928478

[fegin]

LGTM

[awgu]

Unrelated failure:
linux-focal-rocm5.6-py3.8 / test (default, 2, 3, linux.rocm.gpu, unstable)

[awgu]

@pytorchbot merge -f "unrelated failure"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

[awgu]

@awgu has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.