[WIP] Minimize all reduces when accumulating gradients

[joecummings]

Co-authored-by: Nathan Azrak nathan.azrak@gmail.com

Context

See #2415

Changelog

Test plan

UX

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Cancelled Jobs

As of commit 32fbfa7 with merge base d3ab3b7 ():

NEW FAILURE - The following job has failed:

• GPU tests / gpu_test (3.11, stable) (gh)
  tests/recipes/test_full_finetune_distributed.py::TestFullFinetuneDistributedRecipe::test_training_state_on_resume_from_distributed_checkpoint_single_rank[llama3/8B_full-llama3-tune-4-1-True]

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.10, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.9, stable) (gh)
  ##[error]The operation was canceled.

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