Remove composable API's fully_shard from torchrec test (#2549)

Summary:
Pull Request resolved: #2549

The fully_shard name is now used by FSDP2 (torch.distributed._composable.fsdp.fully_shard) and the Composable API's fully_shard (torch.distributed._composable.fully_shard) is being deprecated. Therefore, we want to remove torch.distributed._composable.fully_shard from torchrec as well.
Deprecation message from PyTorch:
https://github.com/pytorch/pytorch/blob/main/torch/distributed/_composable/fully_shard.py#L41-L48

Reviewed By: fegin, iamzainhuda

Differential Revision: D65702643

fbshipit-source-id: 22494e5a6a201a32c86f571c7604383a2e3b3a02

Author: wz337

torchrec/distributed/composable/tests/test_fsdp.py

@@ -14,7 +14,8 @@
 
 import torch
 from torch import nn
-from torch.distributed._composable import fully_shard
+
+# from torch.distributed._composable import fully_shard
 from torch.distributed._shard.sharded_tensor import ShardedTensor
 from torch.distributed._tensor import DTensor
 
@@ -24,6 +25,7 @@
     load_state_dict,
     save_state_dict,
 )
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel
 from torch.distributed.fsdp.wrap import ModuleWrapPolicy
 from torch.distributed.optim import (
@@ -91,15 +93,15 @@ def _run(  # noqa
                 device=ctx.device,
                 plan=row_wise(),
             )
-            m.dense = fully_shard(
+            m.dense = FSDP(  # pyre-ignore
                 m.dense,
+                auto_wrap_policy=ModuleWrapPolicy({nn.Linear}),
                 device_id=ctx.device.index,
-                policy=ModuleWrapPolicy({nn.Linear}),
             )
-            m.over = fully_shard(
+            m.over = FSDP(
                 m.over,
+                auto_wrap_policy=ModuleWrapPolicy({nn.Linear}),
                 device_id=ctx.device.index,
-                policy=ModuleWrapPolicy({nn.Linear}),
             )
 
             dense_opt = KeyedOptimizerWrapper(