Improve FSDP support for low-bit optimizers

[gau-nernst]

• Use DTensor.from_local(run_check=False) to wrap quantized optim state (instead of swapping _local_tensor)
• Make block_size a fixed attribute, calculated inside __init__ (instead of dynamically calculate every time)
• Implement all_gather_into_tensor and wait_tensor to support DTensor.full_tensor()

[pytorch-bot]

🔗 Helpful Links

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

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

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

[drisspg]

curious q: Is there some description of the codes/ scales tensor and their relation to each other?

I can see the pattern that codes has .5x (4bit) and 1x (8bit) the bsize * scale numels
But does this assert square blocks?
I think some description here would be helpful

[gau-nernst]

I will add some description. Basically for 8-bit and FP8, codes has the same shape as the "outer shape", while for 4-bit, since there is bit-packing, I find that it's easier to let codes be a flattened 1D buffer and keep track of the shape manually.
To get the scale, the float tensor is actually flattened first and reshape to (-1, block_size). This is done to relax the requirement that the last dimension must be divisible by block_size -> now we only need numel (total size) to be divisible by block_size. This is especially needed when block size is large (8-bit optim uses block_size=2048 as done in bnb). Since optim update is element-wise, we don't really need to care if the original tensor is 1D, 2D, or n-D (well, maybe there is some structure in n-D tensor that flattening it might not be so wise). I believe the original implementation in bnb does this as well.
-> scale is always a 1D tensor, with size=original_tensor.numel() // block_size

[gau-nernst]

@drisspg Added some docs. Lmk if it is still unclear.

[awgu]

The DTensor-related changes look good to me. I wonder, did you ever have to run .full_tensor() in the training path, or was it only used outside e.g. for debugging?

[gau-nernst]

I don't use FSDP at work or personal projects because I don't have access to multi-GPU machines, so can't really answer your question 😅. Only added FSDP support for low bit optimizers due to request from people.

At least in torchtune, I saw that .full_tensor() is used to retrieve full optim state dict before saving checkpoint https://github.com/pytorch/torchtune/blob/0057fe7cf83e14f0b62538a8d4d20719e0a88639/torchtune/utils/_distributed.py#L437 (though it might be unnecessary, or even less efficient, compared to saving each shard separately, provided we resume training with the same setup).

[jerryzh168]

btw the link is not valid, can you remove . in the end?

[gau-nernst]

Done

[jerryzh168]

same for this one

[msaroufim]

Nice! We should look to get the low bit optimizers out of prototype soon!