Add support for name kwarg in mark_dynamic

[bobrenjc93]

Stack from ghstack (oldest at bottom):

• [wip] consolidate on fullgraph name #163421
• -> Add support for name kwarg in mark_dynamic #163246

Ergonomic improvement to allow sharing symbols without having to do the complex torch._check paradigm as described by @anijain2305 in his recent UED:

Different symbols for KV cached tensors - One property in my case was that the KV cache for different attention blocks had the same seq length, but there is no API to enforce that. The only way is to add torch._check but torch.compile must trace those functions to instruct the dynamic shape infra. This required me to change the model code.
Changing model code is not the best experience. Lets see how transformers maintainers react to my PR. Maybe an API with fullgraph=True is a better bet here.

cc @ezyang @EikanWang @jgong5 @wenzhe-nrv @voznesenskym @penguinwu @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

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

✅ No Failures

As of commit f7c3ddf with merge base 39450e7 ():
💚 Looks good so far! There are no failures yet. 💚

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

[ezyang]

I think doing it with strings is unwise, as in a large enough codebase it can be difficult to avoid collisions, which will cause extremely strange errors. The only use case for strings is if there is a single global configuration spot for dynamic that applies everywhere, but mark_dynamic works both for fullgraph and graph break cases, and also it can propagate unpredictably as its data flow. Easy fix: explicitly allocate Dim symbols (similar to how export does it) and then use those to dedupe by object identity.

Speaking of which, why don't we use export's Dim directly? cc @avikchaudhuri

[bobrenjc93]

I think doing it with strings is unwise, as in a large enough codebase it can be difficult to avoid collisions, which will cause extremely strange errors.

Fair point. I assumed this would be such a rare power-user case that it wouldn’t matter in practice, but you’re right that using object identity is cleaner. I did think about the Dims API, though it seemed like a heavy lift for existing models with mark_dynamic infrastructure (for example, ads PT2 wrappers). Since this is a power-user feature anyway, that cost may not be a big deal. I’ll close this PR for now and see if I can get the Dims API working for compile in a more ergonomic way. Worst case, there may be a middle-ground approach where we create dim-like objects and thread them through the mark_dynamic calls.