Try import ckpt convert (#2476)

Author: xingmingyyj

paddleformers/trainer/utils/ckpt_converter.py

@@ -20,16 +20,48 @@
 
 import paddle
 from paddle.distributed.fleet.utils.log_util import logger
-from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
-    _load_state_dict,
-    get_rank_to_read_files,
-)
-from paddle.distributed.flex_checkpoint.dcp.metadata import (
-    LocalTensorIndex,
-    LocalTensorMetadata,
-    Metadata,
-)
-from paddle.distributed.flex_checkpoint.dcp.utils import flatten_state_dict
+
+try:
+    from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+        _load_state_dict,
+        get_rank_to_read_files,
+    )
+except ModuleNotFoundError:
+    try:
+        from paddle.distributed.checkpoint.load_state_dict import (
+            _load_state_dict,
+            get_rank_to_read_files,
+        )
+    except ModuleNotFoundError:
+        _load_state_dict = None
+        get_rank_to_read_files = None
+
+
+try:
+    from paddle.distributed.flex_checkpoint.dcp.metadata import (
+        LocalTensorIndex,
+        LocalTensorMetadata,
+        Metadata,
+    )
+except ModuleNotFoundError:
+    try:
+        from paddle.distributed.checkpoint.metadata import (
+            LocalTensorIndex,
+            LocalTensorMetadata,
+            Metadata,
+        )
+    except ModuleNotFoundError:
+        LocalTensorIndex = None
+        LocalTensorMetadata = None
+        Metadata = None
+
+try:
+    from paddle.distributed.flex_checkpoint.dcp.utils import flatten_state_dict
+except ModuleNotFoundError:
+    try:
+        from paddle.distributed.checkpoint.utils import flatten_state_dict
+    except ModuleNotFoundError:
+        flatten_state_dict = None
 
 MODEL_WEIGHT_SUFFIX = ".pdparams"
 OPTIMIZER_WEIGHT_SUFFIX = ".pdopt"