[DTensor] enable SimpleFSDP's composability with Tensor Parallel (pytorch#152286)

This PR adds support for SimpleFSDP's composability with Tensor Parallel + torch.compile.

`_StridedShard` is used in SimpleFSDP/FSDP2 to support correct distributed checkpointing when FSDP+TP is applied. Previously, `_StridedShard` is not guarded by torch.compile. This PR adds `_StridedShard` as an additional placement type to be guarded by torch.compile.

Pull Request resolved: pytorch#152286
Approved by: https://github.com/bdhirsh

Author: ruisizhang123

test/distributed/tensor/test_dtensor_compile.py

@@ -29,6 +29,7 @@
     PrepareModuleOutput,
     RowwiseParallel,
 )
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import get_devtype
 from torch.testing._internal.common_utils import (
@@ -196,8 +197,10 @@ def fn(x):
             return a
 
         compiled_fn = torch.compile(backend="aot_eager", fullgraph=True)(fn)
-
-        for x in [Shard(0), Replicate(), Partial()]:
+        split_factors = [2, 3, 4]
+        for x in [Shard(0), Replicate(), Partial()] + [
+            _StridedShard(0, split_factor=s) for s in split_factors
+        ]:
             opt_fn = fn(x)
             compiled_out = compiled_fn(x)
             self.assertEqual(opt_fn, compiled_out)

torch/_dynamo/guards.py

@@ -1702,6 +1702,7 @@ def EQUALS_MATCH(self, guard: Guard):
         if torch.distributed.is_available():
             from torch.distributed.device_mesh import DeviceMesh
             from torch.distributed.tensor.placement_types import (
+                _StridedShard,
                 Partial,
                 Replicate,
                 Shard,
@@ -1712,6 +1713,7 @@ def EQUALS_MATCH(self, guard: Guard):
                 Replicate,
                 Partial,
                 DeviceMesh,
+                _StridedShard,
             )
 
         from torch.export.dynamic_shapes import _IntWrapper