refactor attn layers

[vchiley]

This PR refactors all attn_impl to use one cls which calls diff attn fns (flash, triton, torch) internally.
This means the state dict of one attn_impl is guaranteed to be the same as the other (this is tested now).

Note: this PR uses unpacked attn variants (ie q, k, v are chunked out from qkv), since kv_caching will concatenated tokens to kv and therefore q will be diff seq len from k and v.

Test also now compare all variants to each other with clip, qkln, and alibi (when possible). This is like 55 really quick tests.

@dskhudia lmk if this works for the inference work

@dskhudia also noted that we'll want to move to torch.nn.functional.scaled_dot_product_attention. maybe now, maybe later.

@samhavens re ift. Note: the flash path will unpad the input and do the attn calculation unpadded; the other impl should do calc on the entire pad input but it should be only up to the longest seq in the batch (not the max seq len of the model)

[bmosaicml]

Great work Vitaliy this looks clean and ready for caching

[vchiley]

Running a 125m (with clip 6) using flash, triton, torch in wandb proj: attn_refactor. MFU is about where it should be. Training curves look identical.

wandb proj: attn_refactor also has a 7b run with triton to verify MFU is where it should be (using 16 GPUs) (its at > 43% MFU; checks out)
Also adding 7B with flash to see diff (not training to convergence, just for a 100ish steps). Flash version gets 41% MFU.
Note 7B run uses activation_checkpointing_reentrant: false and limit_all_gathers: true.

[dskhudia]

For inference, ideally we should have export test for onnx/torchscript but can be added in a followup PR.

[abhi-mosaic]

This looks awesome @vchiley ! I added a few comments but they're minor.

[dskhudia]

Looks great. Much cleaner now.

also noted that we'll want to move to torch.nn.functional.scaled_dot_product_attention.

^ can be done later once we move to 2.0.