Implement activation offloading and opt_in_bwd in knowledge_distillation recipes

[AnuravModak]

Towards #1959 .

So this is the first PR for implementing the changes in
KD recipes:
https://github.com/pytorch/torchtune/blob/main/recipes/knowledge_distillation_single_device.py
https://github.com/pytorch/torchtune/blob/main/recipes/knowledge_distillation_distributed.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[AnuravModak]

Hi @felipemello1 , As part of this implementation i have made changes only in knowledge_distillation_single_device.py for activation offloading. And enabled testing parameter in tests/recipes/test_knowledge_distillation_single_device.py. Let me know if its in the right direction, if Yes will conclude the remaining implementation as well!

[joecummings]

This looks great, thanks for doing this!! As part of testing can you output memory usage between activation offloading enabled and disabled as part of the way to confirm that this is doing what we expect it to?

[AnuravModak]

Hi @joecummings ,

Added an import for wandb.
Initialized WandB in the beginning of the training loop if self.use_wandb is set to True.
Logged relevant metrics (log_dict) to WandB if self.use_wandb is True.

Conditional Logging:
The logging mechanism is conditioned by the self.use_wandb flag, enabling or disabling WandB integration.

Profiling and Memory:
Added memory tracking and profiling, which can be used for monitoring resource usage during training.

Let me know if it works!

[ebsmothers]

Hi @AnuravModak thanks for the PR! A couple suggestions:

1. Regarding @joecummings's previous suggestion to add WandB logs, you should just be able to add WandB as a metric logger from the command line. So you will not need to make any changes to the code to do this. Once you have the results you can paste them in the PR summary. You can see e.g. the test plan in DoRA fixes #2139 for an example of this.
2. I think some extra files under .idea directory snuck in, can you remove them?
3. Why are the activation offloading changes in knowledge_distillation_single_device (specifically L461-L475) commented out?