implement activation offloading and opt_in_bwd in knowledge_distillation recipes

[felipemello1]

Current knowledge distillation recipes don't have support for activation offloading and opt_in_bwd.

The implementation should be similar to the one in other recipes, like full_finetuning_distributed.

After enabling it in the recipe, it should also be enabled in the configs related to KD.

PRs with reference implementation:
activation offloading: #1847
opt_in_bwd implementation: #1833

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

after implementing it, run it with the flag on/off and plot the graphs of loss/memory/words per second. The easier way is to add the wandb logger to the config.

to update configs in bulk, you can use the script here: #1954

[AnuravModak]

Hi @felipemello1 is it available for external contributors? if yes kindly assign it to me will look into it. Thanks in advance!

[felipemello1]

@AnuravModak yes! Any issue with "community help wanted" is something that we think would be a great fit for external contributors. Feel free to ask me questions and just submit a PR.

More info on contributing here, but the TLDR:
main/CONTRIBUTING.md

1. Fork torchtune: pytorch/torchtune
   then clone from it: git clone github.com<YOUR_GITHUB_USER>/torchtune.git

2. install dependencies

cd torchtune
conda create -n torchtune python=3.11
conda activate torchtune
pip install --pre --upgrade torch torchvision torchao --index-url download.pytorch.org/whl/nightly/cu124
pip install -e ".[dev]"
pre-commit install

3. create the PR

[rajuptvs]

Hi @felipemello1 , can i give this a go, i have currently implemented the activation offloading and opt_in_bwd for single device knowledge distillation.

Happy to create a PR ( if not being actively worked on already) and get a review of single device kd, before going to the distributed.

Thanks!