[TP][2D][EZ] Fix Error in FSDP 2D test (pytorch#107975)

As title, TP dimension should be the second dim, so we need to pass tp_degree to the second rather the first dim of the mesh tensor.

Pull Request resolved: pytorch#107975
Approved by: https://github.com/wz337, https://github.com/awgu

test/distributed/tensor/parallel/test_fsdp_2d_parallel.py

@@ -96,7 +96,7 @@ def init_model(
     # 2-D mesh is [dp, tp]
     twod_mesh = DeviceMesh(
         device_type="cuda",
-        mesh=torch.arange(0, world_size).view(model_parallel_size, -1),
+        mesh=torch.arange(0, world_size).view(-1, model_parallel_size),
     )
 
     fsdp_pg = twod_mesh.get_dim_groups()[0]