Offloading support for multiple attention layouts

[sanandaraj5597]

Description

CPU offloading currently only supports the sbhd_sbhd_sbhd layout, but we have multiple other layouts for pre-training, fine-tuning of LLM's.

This PR adds support offloading for all attention layouts.

[pggPL]

I think it will not work for offloaded layers, because .to() by default preserves memory format. I think it needs to be changed to .to(device="cpu", memory_format=torch.contiguous_format).

[pggPL]

Otherwise it looks good, have you tested it somehow?

[sanandaraj5597]

I think it will not work for offloaded layers, because .to() by default preserves memory format. I think it needs to be changed to .to(device="cpu", memory_format=torch.contiguous_format).

It will work right because the cpu copy is created in a contiguous fashion. When you do .to(), the cpu copy will be moved to GPU which is also contiguous. And that's the reason why we breakdown all the formats (sbh3d/th3d/...) to contiguous (sbhd_sbhd_sbhd/thd_thd_thd/...) in attention.

have you tested it somehow?

Yes I've run E2E tests on top of this for pre-training and fine-tuning.

[pggPL]

https://docs.pytorch.org/docs/stable/generated/torch.Tensor.to.html It preserves memory formats. So it will split sbh3d into 3 separate tensors (non interleaved), none of them being contiguous - all of them will have the same stride as previously.