[low-bit optim] Fix edge cases for FSDP2 integration

[gau-nernst]

This PR fixes 2 issues that came up in torchtune when fine-tuning Llama3.2-vision

1. Sometimes there is a strange torch.compile() error with DTensor when there is .grad field. pytorch/torchtune#1978 (comment)

This is actually an old issue #652 (see #652 (comment) for more details). However, it does not always happen: our CI test passes (for PyTorch 2.6, but PyTorch 2.5 has this issue), finetune Llama text (not multimodal) in torchtune has no issues, but finetuning Llama3.2-vision faces the error -> It's not clear why and how this happens. The error message seems to indicate that torch.compile() tries to do dynamic-shape, even though we are explicitly using dynamic=False

The solution is to call .detach() on param, which shares the same weight storage, but now it doesn't have .grad anymore. Thanks to this, low-bit optim + FSDP2 also work for PyTorch 2.5 CI now (previously it didn't).

I can't add a test for this, since I don't know how/when this happens.

2. Wrong DTensor creation when there is uneven sharding (i.e. 1st dim is not divisible by world size)

Usually we don't have uneven shards for LLMs, thus this error didn't surface. However, for ViT, it might be possible due to pos_embed: in some implementation pos_embed includes CLS token, hence the first dim is num_visual_tokens + 1.

The fix is simple: pass shape (and stride) to DTensor.from_local(). An appropriate test has also been added.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 38bf355 with merge base 6ff3904 ():
💚 Looks good so far! There are no failures yet. 💚

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

[gau-nernst]

@vkuzo Do you have time to take a look at this PR? Thank you 🙏