run sdpa with dtensor

[tianyu-l]

Stack from ghstack (oldest at bottom):

• -> run sdpa with dtensor #163

There are several caveats:

1. FSDP + TP uses bfloat16 as default param dtype, which seems to create minor numerical discrepancies. I'm still investigating if the discrepancies are caused by collectives or dtensor.
2. Flash attention only supports bfloat16 and float16.
   (1) if dtype is torch.float32, efficient attention will be used, which is not supported by DTensor yet (will add this);
   (2) if dtype is torch.float64, decomposed attention will be used, which is not compatible with DTensor due to the usage of intermediate torch.Tensor mask.
3. 1D TP only parallelism would fail with this PR by default for the reason mentioned in 2.(1).

[drisspg]

Looks good, nice cleanup!

[tianyu-l]

Closing this PR and use #180 instead to work with FSDP2 changes.