Move distributed checkpointing to a subdirectory of experimental

pytorch · Sep 29, 2023 · c060b6e · c060b6e
1 parent 63a9020
commit c060b6e
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/test/spmd/test_xla_distributed_checkpoint.py b/test/spmd/test_xla_distributed_checkpoint.py
@@ -15,9 +15,10 @@
     create_default_local_save_plan,
     create_default_global_save_plan,
 )
-from torch_xla.experimental.distributed_checkpoint import SPMDLoadPlanner, SPMDSavePlanner
-from torch_xla.experimental._distributed_checkpoint_helpers import (
-    _sharded_cpu_state_dict, _CpuShards, _is_sharded_tensor)
+from torch_xla.experimental.checkpoint import SPMDLoadPlanner, SPMDSavePlanner
+from torch_xla.experimental.checkpoint._helpers import (_sharded_cpu_state_dict,
+                                                        _CpuShards,
+                                                        _is_sharded_tensor)
 
 
 class DistributedCheckpointTestBase(test_xla_sharding_base.XlaShardingTest):

diff --git a/torch_xla/experimental/checkpoint/__init__.py b/torch_xla/experimental/checkpoint/__init__.py
@@ -0,0 +1,6 @@
+from .planners import SPMDSavePlanner, SPMDLoadPlanner
+
+__all__ = [
+    "SPMDSavePlanner",
+    "SPMDLoadPlanner",
+]
diff --git a/...mental/_distributed_checkpoint_helpers.py → ...h_xla/experimental/checkpoint/_helpers.py b/...mental/_distributed_checkpoint_helpers.py → ...h_xla/experimental/checkpoint/_helpers.py
diff --git a/...la/experimental/distributed_checkpoint.py → ...h_xla/experimental/checkpoint/planners.py b/...la/experimental/distributed_checkpoint.py → ...h_xla/experimental/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.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):
   """