[Checkpoint] Test file for checkpoint_wrapper

[rohan-varma]

Stack from ghstack (oldest at bottom):

• Fix non-reentrant hooks based checkpointing #78752
• Add apply_activation_checkpointing helper function #78704
• -> [Checkpoint] Test file for checkpoint_wrapper #78702

As we make additional changes and don't couple checkpoint_wrapper with
FSDP, would be good to have standalone test file for it.

Currently, we port the only non FSDP specific unittest that exists, and will add more as development continues.

[facebook-github-bot]

🔗 Helpful links

• 🧪  See artifacts and rendered test results at hud.pytorch.org/pr/78702
• 📄  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 our office hours

✅ No Failures (0 Pending)

As of commit 22f044b (more details on the Dr. CI page):

Expand to see more

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[awgu]

Is the plan to keep test_checkpoint_wrapper.py in test/distributed/fsdp/ in the short-term and relocate it in the long term?

[rohan-varma]

@awgu Yes, test_checkpoint_wrapper.py in fsdp tests tests integration / interaction with FSDP as that is a very common use case.

Although agreed it is not ideal to have 2 test files of the same name due to confusion, that one can probably be named test_checkpoint_wrapper_with_fsdp which I'll do in a follow up.

[rohan-varma]

@pytorchbot merge

[pytorchmergebot]

Merge failed due to Refusing to merge as mandatory check(s) Lint failed for rule superuser
Raised by https://github.com/pytorch/pytorch/actions/runs/2435959380

[rohan-varma]

@pytorchbot merge

[pytorchmergebot]

@pytorchbot successfully started a merge job. Check the current status here

[github-actions]

Hey @rohan-varma.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.