Move distributed checkpointing to a subdirectory of experimental (#5656)

pytorch · Oct 10, 2023 · e7e7c5e · e7e7c5e
1 parent 6dccf0a
commit e7e7c5e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/test/spmd/test_xla_distributed_checkpoint.py b/test/spmd/test_xla_distributed_checkpoint.py
@@ -16,7 +16,7 @@
     create_default_global_save_plan,
 )
 from torch_xla.experimental.distributed_checkpoint import SPMDLoadPlanner, SPMDSavePlanner
-from torch_xla.experimental._distributed_checkpoint_helpers import (
+from torch_xla.experimental.distributed_checkpoint._helpers import (
     _sharded_cpu_state_dict, _CpuShards, _is_sharded_tensor)
 
 

diff --git a/torch_xla/experimental/distributed_checkpoint/__init__.py b/torch_xla/experimental/distributed_checkpoint/__init__.py
@@ -0,0 +1,6 @@
+from .planners import SPMDSavePlanner, SPMDLoadPlanner
+
+__all__ = [
+    "SPMDSavePlanner",
+    "SPMDLoadPlanner",
+]
diff --git a/...mental/_distributed_checkpoint_helpers.py → ...mental/distributed_checkpoint/_helpers.py b/...mental/_distributed_checkpoint_helpers.py → ...mental/distributed_checkpoint/_helpers.py
diff --git a/...la/experimental/distributed_checkpoint.py → ...mental/distributed_checkpoint/planners.py b/...la/experimental/distributed_checkpoint.py → ...mental/distributed_checkpoint/planners.py
@@ -35,16 +35,11 @@
 from torch.distributed.checkpoint.utils import find_state_dict_object
 from torch.utils._pytree import tree_map
 from torch_xla.experimental.xla_sharding import XLAShardedTensor, XLAShard
-from torch_xla.experimental._distributed_checkpoint_helpers import (
+from torch_xla.experimental.distributed_checkpoint._helpers import (
     FLATTEN_MAPPING, flatten_state_dict, dedup_tensors, _is_sharded_tensor,
     set_element, narrow_tensor_by_index, _unwrap_xla_sharded_tensor, _CpuShards)
 from typing import Any, Dict, List, Tuple, Union
 
-__all__ = [
-    "SPMDSavePlanner",
-    "SPMDLoadPlanner",
-]
-
 
 class SPMDSavePlanner(SavePlanner):
   """