do NOT review - offloading for other recipes

[felipemello1]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?

Test plan

tune run full_finetune_single_device --config llama3_1/8B_full_single_device \
dataset._component_=torchtune.datasets.text_completion_dataset \
dataset.source=text \
dataset.split="train" \
dataset.data_files="/tmp/fake_long_dataset/train/fake_long_dataset.txt" \
tokenizer.max_seq_len=8192 \
batch_size=4 \
max_steps_per_epoch=10 \
metric_logger=torchtune.training.metric_logging.WandBLogger \
gradient_accumulation_steps=1 \
epochs=1 \
log_peak_memory_stats=True \
optimizer=bitsandbytes.optim.PagedAdamW8bit \
optimizer_in_bwd=True \
compile=True \
enable_activation_checkpointing=True \
enable_activation_offloading=True \
offload_with_streams=True

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 7aa7e49 with merge base 4fbe7b2 ():
💚 Looks good so far! There are no failures yet. 💚

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

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 4.97238% with 172 lines in your changes missing coverage. Please review.

Project coverage is 72.29%. Comparing base (ee343e6) to head (7aa7e49).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
torchtune/training/_activation_offloading.py	5.88%	128 Missing ⚠️
recipes/full_finetune_distributed.py	0.00%	11 Missing ⚠️
recipes/full_finetune_single_device.py	0.00%	11 Missing ⚠️
recipes/lora_finetune_distributed.py	0.00%	11 Missing ⚠️
recipes/lora_finetune_single_device.py	0.00%	11 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1578      +/-   ##
==========================================
- Coverage   73.08%   72.29%   -0.79%     
==========================================
  Files         289      290       +1     
  Lines       14157    14348     +191     
==========================================
+ Hits        10346    10373      +27     
- Misses       3811     3975     +164     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.