Revert "[FSDP2] Move to public torch.distributed.fsdp (pytorch#141868)"

This reverts commit 45583a5.

Reverted pytorch#141868 on behalf of https://github.com/atalman due to failing internally ([comment](pytorch#141868 (comment)))

Author: pytorchmergebot

docs/source/distributed.fsdp.fully_shard.rst

@@ -79,7 +79,6 @@ Features described in this documentation are classified by release status:
    torch.distributed.algorithms.join <distributed.algorithms.join>
    torch.distributed.elastic <distributed.elastic>
    torch.distributed.fsdp <fsdp>
-   torch.distributed.fsdp.fully_shard <distributed.fsdp.fully_shard>
    torch.distributed.tensor.parallel <distributed.tensor.parallel>
    torch.distributed.optim <distributed.optim>
    torch.distributed.pipelining <distributed.pipelining>

docs/source/index.rst

@@ -10,7 +10,7 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from torch.distributed.fsdp import fully_shard
+from torch.distributed._composable.fsdp import fully_shard
 from torch.nn.parallel.scatter_gather import _is_namedtuple
 from torch.testing._internal.common_cuda import TEST_CUDA
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu

test/distributed/_composable/fsdp/test_fully_shard_autograd.py

@@ -7,8 +7,8 @@
 import torch
 import torch.nn as nn
 from torch.distributed._composable import replicate
+from torch.distributed._composable.fsdp import fully_shard
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
-from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest, MLPStack

test/distributed/_composable/fsdp/test_fully_shard_clip_grad_norm_.py

@@ -11,30 +11,30 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.distributed._composable import checkpoint, replicate
-from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
-from torch.distributed.fsdp import (
+from torch.distributed._composable.fsdp import (
     FSDPModule,
     fully_shard,
     MixedPrecisionPolicy,
     OffloadPolicy,
 )
-from torch.distributed.fsdp._fully_shard._fsdp_collectives import (
+from torch.distributed._composable.fsdp._fsdp_collectives import (
     _div_if_needed,
     _get_gradient_divide_factors,
     foreach_all_gather,
     foreach_all_gather_copy_out,
     foreach_reduce,
 )
-from torch.distributed.fsdp._fully_shard._fsdp_common import FSDPMeshInfo, TrainingState
-from torch.distributed.fsdp._fully_shard._fsdp_init import (
+from torch.distributed._composable.fsdp._fsdp_common import FSDPMeshInfo, TrainingState
+from torch.distributed._composable.fsdp._fsdp_init import (
     _get_post_forward_mesh_info,
     _init_default_fully_shard_mesh,
 )
-from torch.distributed.fsdp._fully_shard._fsdp_param import ShardedState
-from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup
-from torch.distributed.tensor import DTensor
+from torch.distributed._composable.fsdp._fsdp_param import ShardedState
+from torch.distributed._composable.fsdp._fsdp_param_group import FSDPParamGroup
+from torch.distributed._tensor import DTensor
+from torch.distributed._tensor.experimental import implicit_replication
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.tensor.debug import CommDebugMode
-from torch.distributed.tensor.experimental import implicit_replication
 from torch.testing._internal.common_cuda import TEST_CUDA
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import (

test/distributed/_composable/fsdp/test_fully_shard_comm.py

@@ -12,19 +12,17 @@
 
 import torch
 import torch._dynamo.testing
+import torch.distributed._composable.fsdp._fsdp_param
 import torch.nn.functional as F
 from torch import nn
 from torch._dynamo.utils import counters
 from torch._inductor import comms
 from torch._inductor.utils import is_fallback_op, run_and_get_code
+from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed._composable.fsdp._fsdp_common import TrainingState
+from torch.distributed._composable.fsdp._fsdp_param_group import FSDPParamGroup
 from torch.distributed._tensor import init_device_mesh
-from torch.distributed.fsdp import (
-    fully_shard,
-    FullyShardedDataParallel as FSDP,
-    ShardingStrategy,
-)
-from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
-from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
 from torch.testing import FileCheck
 from torch.testing._internal.common_distributed import (
     at_least_x_gpu,
@@ -85,7 +83,7 @@ def _test_disable_compiling_hooks(
     ):
         torch._dynamo.reset()
         trace_rules_check_count = 0
-        HOOKS_FILE_NAME = "torch/distributed/fsdp/_fully_shard/_fsdp_state.py"
+        HOOKS_FILE_NAME = "torch/distributed/_composable/fsdp/_fsdp_state.py"
         HOOK_WRAPPER_NAME = "fsdp_hook_wrapper"
 
         def patched_trace_rules_check(*args, **kwargs):

test/distributed/_composable/fsdp/test_fully_shard_compile.py

@@ -13,8 +13,8 @@
 import torch.nn as nn
 import torch.utils._pytree as pytree
 from torch.autograd.grad_mode import _unsafe_preserve_version_counter
+from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
-from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
 from torch.testing._internal.common_cuda import TEST_CUDA
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import (

test/distributed/_composable/fsdp/test_fully_shard_extensions.py

@@ -10,8 +10,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.distributed._composable import checkpoint, replicate
-from torch.distributed.fsdp import fully_shard
-from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
+from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed._composable.fsdp._fsdp_param_group import (
     RegisterPostBackwardFunction,
 )
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu

test/distributed/_composable/fsdp/test_fully_shard_frozen.py

@@ -4,8 +4,8 @@
 import torch
 import torch.nn as nn
 from torch.amp.grad_scaler import GradScaler, OptState
+from torch.distributed._composable.fsdp import fully_shard
 from torch.distributed._tensor import init_device_mesh
-from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
     parallelize_module,

test/distributed/_composable/fsdp/test_fully_shard_grad_scaler.py

@@ -9,6 +9,13 @@
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed._composable import replicate
+from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed._composable.fsdp._fsdp_init import (
+    _get_managed_modules,
+    _get_managed_states,
+)
+from torch.distributed._composable.fsdp._fsdp_param import ParamModuleInfo
+from torch.distributed._composable.fsdp._fsdp_param_group import _get_param_module_infos
 from torch.distributed._tensor import (
     DeviceMesh,
     distribute_tensor,
@@ -17,15 +24,6 @@
     Shard,
 )
 from torch.distributed.device_mesh import init_device_mesh
-from torch.distributed.fsdp import fully_shard
-from torch.distributed.fsdp._fully_shard._fsdp_init import (
-    _get_managed_modules,
-    _get_managed_states,
-)
-from torch.distributed.fsdp._fully_shard._fsdp_param import ParamModuleInfo
-from torch.distributed.fsdp._fully_shard._fsdp_param_group import (
-    _get_param_module_infos,
-)
 from torch.distributed.fsdp._init_utils import (
     _init_inter_node_process_group,
     _init_intra_node_process_group,
@@ -1158,26 +1156,5 @@ def shard_placement_fn(param: nn.Parameter) -> Optional[Shard]:
             fully_shard(model, shard_placement_fn=shard_placement_fn)
 
 
-# TODO: Remove this test class once we remove the old import path:
-# torch/distributed/_composable/fsdp
-class TestFullyShardOldImport(FSDPTestMultiThread):
-    @property
-    def world_size(self) -> int:
-        return 2
-
-    @unittest.skipIf(not TEST_CUDA, "no cuda")
-    def test_old_import_training(self):
-        from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy
-
-        model = nn.Sequential(nn.Linear(16, 16), nn.Linear(16, 16))
-        mp_policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16)
-        fully_shard(model[0], mp_policy=mp_policy)
-        fully_shard(model[1], mp_policy=mp_policy)
-        fully_shard(model, mp_policy=mp_policy)
-
-        inp = torch.randn((8, 16), device="cuda")
-        model(inp).sum().backward()
-
-
 if __name__ == "__main__":
     run_tests()