[WIP] 2D parallelism for training

[joecummings]

No description provided.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 2 Cancelled Jobs

As of commit cfd31cd with merge base d4465c8 ():

NEW FAILURES - The following jobs have failed:

• Build Docs / build_docs (3.11) (gh)
• GPU tests / gpu_test (3.11, stable) (gh)
  tests/recipes/test_full_finetune_distributed.py::TestFullFinetuneDistributedRecipe::test_training_state_on_resume_from_distributed_checkpoint_multi_rank[llama3/8B_full-llama3-tune-4-1-True]
• Lint / lint (3.10) (gh)
  Process completed with exit code 1.

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

• GPU tests / gpu_test (3.10, stable) (gh)
  tests/recipes/test_full_finetune_distributed.py::TestFullFinetuneDistributedRecipe::test_training_state_on_resume_from_distributed_checkpoint_multi_rank[llama3/8B_full-llama3-tune-4-1-True]
• 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.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 6.66667% with 28 lines in your changes missing coverage. Please review.

Project coverage is 23.64%. Comparing base (46d8153) to head (cfd31cd).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
recipes/full_finetune_distributed.py	0.00%	15 Missing ⚠️
torchtune/training/_distributed.py	13.33%	13 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2318      +/-   ##
==========================================
- Coverage   23.87%   23.64%   -0.24%     
==========================================
  Files         359      361       +2     
  Lines       21260    21488     +228     
==========================================
+ Hits         5076     5080       +4     
- Misses      16184    16408     +224     

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

[SalmanMohammadi]

FYI the default parallelize plan we have didn't work when I was using it with the smaller Llama3 models with tied embeddings.