[WIP] compile graph breaks

[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)

addressing compile graph breaks

Test plan

Rough numbers:
11b, compiling only decoder

1. without the changes,

28s on the first step
21s on the second step
0.9s any other step

2. with dynamic shape changes

~28s in the first step,
~ 0.9s any other step

3. further adding requires grad (i am not sure if this does anything to the loss)
   h = self.tok_embeddings(tokens)
   h.requires_grad = True

First step ~18s
~ 0.9s any other step

So we save ~21s to 31s of compilation time by avoiding the graph breaks due to shape plus optionally adding requires_grad. This is without compiling the encoder, which NaNs if we do (Brian will look into that).

TODO: Need to check if there is a perf impact when packed=True and there wont be graph breaks.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [DomainsOnly] Jobs fail with GLIBC version not found

❌ 3 New Failures, 6 Cancelled Jobs

As of commit b41114a with merge base d5c54f3 ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.10, stable) (gh)
  tests/torchtune/modules/test_transformer_decoder.py::TestTransformerDecoder::test_kv_cache_batch_size_exceeded
• Recipe Tests / recipe_test (3.11) (gh)
  tests/recipes/test_ppo_full_finetune_single_device.py::TestPPOFullFinetuneSingleDeviceRecipe::test_training_state_on_resume_with_optimizer_in_bwd
• Unit Test / unit_tests (3.9) (gh)
  tests/torchtune/modules/test_transformer_decoder.py::TestTransformerDecoder::test_kv_cache_batch_size_exceeded

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

• GPU tests / gpu_test (3.11, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.9, stable) (gh)
  ##[error]The operation was canceled.
• Recipe Tests / recipe_test (3.10) (gh)
  ##[error]The operation was canceled.
• Recipe Tests / recipe_test (3.9) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.10) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.11) (gh)
  ##[error]The operation was canceled.

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