[Auto Parallel] Improve the interface and the underlying mechanisms

python/paddle/distributed/__init__.py

@@ -43,10 +43,6 @@
 
 from .auto_parallel import shard_op  # noqa: F401
 from .auto_parallel import shard_tensor  # noqa: F401
-from .auto_parallel import set_shard_mask  # noqa: F401
-from .auto_parallel import set_offload_device  # noqa: F401
-from .auto_parallel import set_pipeline_stage  # noqa: F401
-from .auto_parallel import ProcessMesh  # noqa: F401
 
 from .fleet import BoxPSDataset  # noqa: F401
 

python/paddle/distributed/auto_parallel/__init__.py

@@ -14,10 +14,11 @@
 
 from .interface import shard_tensor  # noqa: F401
 from .interface import shard_op  # noqa: F401
-from .interface import set_shard_mask  # noqa: F401
-from .interface import set_offload_device  # noqa: F401
-from .interface import set_pipeline_stage  # noqa: F401
-from .interface import ProcessMesh  # noqa: F401
+from .process_mesh import ProcessMesh
+# from .interface import set_shard_mask  # noqa: F401
+# from .interface import set_offload_device  # noqa: F401
+# from .interface import set_pipeline_stage  # noqa: F401
+# from .interface import ProcessMesh  # noqa: F401
 from .completion import complete_annotation  # noqa: F401
 from .completion import complete_backward_annotation  # noqa: F401
 from .reshard import reshard  # noqa: F401