mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/ops_generate/resources/resource_loader.py

@@ -0,0 +1,36 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Module of base class for resource loader."""
+
+from abc import ABC, abstractmethod
+from typing import Dict
+
+from .resource_list import ResourceType
+
+
+class ResourceLoader(ABC):
+    """
+    Abstract class for resource loader.
+    """
+    @abstractmethod
+    def load(self) -> Dict[ResourceType, object]:
+        """
+        Load resource.
+
+        Returns:
+            Dict[ResourceType, object]: The resource type and resource object map.
+        """
+        raise NotImplementedError

mindspore/ops_generate/resources/resource_manager.py

@@ -0,0 +1,64 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Module managing resource."""
+
+from common.op_proto import OpProtoLoader, DeprecatedOpProtoLoader, FuncOpProtoLoader
+from api.op_api_proto import OpApiProtoLoader
+
+from .resource_loader import ResourceLoader
+from .resource_list import ResourceType
+from .yaml_loader import OpDocYamlLoader, TensorMethodDocYamlLoader, MintFuncDocYamlLoader
+
+
+class ResourceManager():
+    """
+    ResourceManager is a class for managing resources.
+    """
+
+    def __init__(self):
+        self.resource_map = {}
+
+    def register_resource(self, loader: ResourceLoader) -> None:
+        """
+        Register resource.
+        """
+        self.resource_map.update(loader.load())
+
+    def get_resource(self, type: ResourceType) -> object:
+        """
+        Get resource by type.
+        """
+        if type not in self.resource_map:
+            raise ValueError(f"Resource '{type.name}' not registered")
+        return self.resource_map[type]
+
+
+def prepare_resources() -> ResourceManager:
+    """
+    Load needed resources.
+    """
+    resource_mgr = ResourceManager()
+    resource_mgr.register_resource(OpProtoLoader())
+    resource_mgr.register_resource(DeprecatedOpProtoLoader())
+    resource_mgr.register_resource(FuncOpProtoLoader())
+    resource_mgr.register_resource(OpDocYamlLoader())
+    resource_mgr.register_resource(TensorMethodDocYamlLoader())
+    resource_mgr.register_resource(MintFuncDocYamlLoader())
+    resource_mgr.register_resource(OpApiProtoLoader(
+        resource_mgr.get_resource(ResourceType.OP_PROTO),
+        resource_mgr.get_resource(ResourceType.DEPRECATED_OP_PROTO),
+        resource_mgr.get_resource(ResourceType.FUNC_OP_PROTO)))
+    return resource_mgr

mindspore/ops_generate/resources/yaml_loader.py

@@ -0,0 +1,88 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Module loading ops yaml."""
+
+import os
+from typing import Sequence, Union
+
+from common.gen_utils import safe_load_yaml_from_dir
+import common.gen_constants as K
+
+from .resource_loader import ResourceLoader
+from .resource_list import ResourceType
+
+
+class YamlLoader(ResourceLoader):
+    """
+    YamlLoader is a utility class for loading yaml files.
+    """
+    def __init__(self, resouce_type: ResourceType, yaml_path: Union[Sequence[str], str]):
+        """
+        Initialize YamlLoader.
+
+        Args:
+            resouce_type (ResourceType): The type of the resource.
+            yaml_path (Union[Sequence[str], str]): The path to the yaml file or directory.
+        """
+        self.type = resouce_type
+        if isinstance(yaml_path, str):
+            self.yaml_path = [yaml_path]
+        else:
+            self.yaml_path = yaml_path
+
+    def load(self) -> dict:
+        """
+        Load yaml files.
+
+        Returns:
+            tuple[int, object]: The resource id and the yaml dict.
+        """
+        for yaml_path in self.yaml_path:
+            if not os.path.isdir(yaml_path):
+                raise ValueError(f"yaml path '{yaml_path}' not found")
+
+        yaml_dict = {}
+        for yaml_path in self.yaml_path:
+            yaml_dict.update(safe_load_yaml_from_dir(yaml_path))
+
+        return {self.type: yaml_dict}
+
+
+class OpDocYamlLoader(YamlLoader):
+    """
+    OpDocYamlLoader is a class for loading op primitive doc yaml files.
+    """
+    def __init__(self):
+        op_doc_yaml_path = os.path.join(K.WORK_DIR, K.MS_OP_DEF_YAML_PATH, "doc")
+        super().__init__(ResourceType.OP_DOC_YAML, op_doc_yaml_path)
+
+
+class TensorMethodDocYamlLoader(YamlLoader):
+    """
+    TensorMethodDocYamlLoader is a class for loading tensor method doc yaml files.
+    """
+    def __init__(self):
+        tensor_method_doc_yaml_path = os.path.join(K.WORK_DIR, K.MS_TENSOR_METHOD_DOC_YAML_PATH)
+        super().__init__(ResourceType.TENSOR_METHOD_DOC_YAML, tensor_method_doc_yaml_path)
+
+
+class MintFuncDocYamlLoader(YamlLoader):
+    """
+    MintFuncDocYamlLoader is a class for loading mint func doc yaml files.
+    """
+    def __init__(self):
+        mint_func_doc_yaml_path = os.path.join(K.WORK_DIR, K.MS_MINT_FUNC_DOC_YAML_PATH)
+        super().__init__(ResourceType.MINT_FUNC_DOC_YAML, mint_func_doc_yaml_path)

mindspore/ops_generate/tensor_py_cc_generator.py

@@ -0,0 +1,122 @@
+# Copyright 2025 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""
+Generates mindspore/ccsrc/pybind_api/ir/tensor_py.cc which includes the CPython Tensor APIs.
+"""
+
+import os
+import common.gen_constants as K
+from common.gen_utils import save_file
+import common.template as template
+from common.template import Template
+from common.base_generator import BaseGenerator
+from pyboost import pyboost_utils
+
+class TensorPyCppGenerator(BaseGenerator):
+    """
+    This class is responsible for generating mindspore/ccsrc/pybind_api/ir/tensor_register/
+    auto_generate/tensor_py_gen.cc
+    """
+    def __init__(self):
+        self.TENSOR_PY_CC_TEMPLATE = template.TENSOR_PY_CC_TEMPLATE
+        self.TENSOR_PY_H_TEMPLATE = template.TENSOR_PY_H_TEMPLATE
+        self.cpy_wrapper_template = Template("  DEFINE_TENSOR_METHOD_CPYWRAPPER(${pascal_api_name}) \\")
+        self.tensor_api_def_template = Template(
+            '{"${snake_api_name}"'
+            ', (PyCFunction)TensorMethod${pascal_api_name}_CPyWrapper, METH_VARARGS | METH_KEYWORDS},'
+        )
+        self.stubtensor_api_def_template = Template(
+            'py::cpp_function TensorMethod${snake_api_name}_wrapper(\n'
+            '  [](const py::object& self, const py::args& args, const py::kwargs& kwargs) {\n'
+            '      return  TensorMethod${pascal_api_name}(self, args, kwargs);\n'
+            '  },\n'
+            '  py::is_method(stubTensorClass)\n'
+            ');\n'
+            'stubTensorClass.attr("${snake_api_name}") = TensorMethod${snake_api_name}_wrapper;'
+        )
+
+    def generate(self, work_path, tensor_method_protos, alias_func_mapping):
+        """
+        Generates the content for the helper file and saves it to the specified path.
+
+        Args:
+            work_path (str): The directory where the generated file will be saved.
+            tensor_method_protos (dict): A dict mapping from Tensor func API names to their proto lists.
+            alias_func_mapping (dict): A dictionary mapping function name to its alias function names.
+
+        Returns:
+            None
+        """
+        wrapper_defs = []
+        tensor_api_defs = []
+        stubtensor_api_defs = []
+        for api_name, _ in tensor_method_protos.items():
+            pascal_api_name = pyboost_utils.format_func_api_name(api_name)
+            snake_api_name = api_name
+            wrapper_defs.append(self.cpy_wrapper_template.replace(pascal_api_name=pascal_api_name))
+            tensor_api_defs.append(
+                self.tensor_api_def_template.replace(
+                    snake_api_name=snake_api_name,
+                    pascal_api_name=pascal_api_name
+                )
+            )
+            stubtensor_api_defs.append(
+                self.stubtensor_api_def_template.replace(
+                    snake_api_name=snake_api_name,
+                    pascal_api_name=pascal_api_name
+                )
+            )
+            if api_name in alias_func_mapping:
+                alias_api_names = alias_func_mapping[api_name]
+                for alias_api_name in alias_api_names:
+                    snake_api_name = alias_api_name
+                    tensor_api_defs.append(
+                        self.tensor_api_def_template.replace(
+                            snake_api_name=snake_api_name,
+                            pascal_api_name=pascal_api_name
+                        )
+                    )
+                    stubtensor_api_defs.append(
+                        self.stubtensor_api_def_template.replace(
+                            snake_api_name=snake_api_name,
+                            pascal_api_name=pascal_api_name
+                        )
+                    )
+
+        # delete the ' \' for the last wrapper macro definition
+        wrapper_defs[-1] = wrapper_defs[-1][:-2]
+
+        file_str = self.TENSOR_PY_CC_TEMPLATE.replace(
+            tensor_api_defs=tensor_api_defs,
+            stubtensor_api_defs=stubtensor_api_defs
+        )
+        save_file(
+            os.path.join(work_path, K.TENSOR_PY_CC_PATH),
+            "tensor_py_gen.cc",
+            file_str
+        )
+
+        file_str = self.TENSOR_PY_H_TEMPLATE.replace(CPyWrapper_defs=wrapper_defs)
+        save_file(
+            os.path.join(work_path, K.TENSOR_PY_CC_PATH),
+            "tensor_py_gen.h",
+            file_str
+        )
+
+def _format_api_name(api_name):
+    has_suffix = api_name.endswith("_")
+    parts = api_name.strip("_").split("_")
+    formatted_api_name = "".join(part.capitalize() for part in parts)
+    return formatted_api_name + '_' if has_suffix else formatted_api_name

mindspore/parallel/__init__.py

@@ -19,7 +19,9 @@ from mindspore.parallel.algo_parameter_config import get_algo_parameters, reset_
     set_algo_parameters
 from mindspore.parallel.checkpoint_transform import rank_list_for_transform, transform_checkpoint_by_rank, \
     transform_checkpoints, merge_pipeline_strategys, sync_pipeline_shared_parameters, \
-    load_segmented_checkpoints, set_op_strategy_config
+    load_segmented_checkpoints, set_op_strategy_config, load_distributed_checkpoint, \
+    merge_sliced_parameter, restore_group_info_list, build_searched_strategy
+from mindspore.parallel.checkpoint_convert import rank_list_for_convert, convert_checkpoint_by_rank, convert_checkpoints
 from mindspore.parallel.parameter_broadcast import parameter_broadcast
 from mindspore.parallel.shard import shard, Layout
 from mindspore.parallel.transform_safetensors import unified_safetensors
@@ -27,4 +29,6 @@ from mindspore.parallel.transform_safetensors import unified_safetensors
 __all__ = ["set_algo_parameters", "reset_algo_parameters", "get_algo_parameters", "rank_list_for_transform",
            "transform_checkpoint_by_rank", "transform_checkpoints", "merge_pipeline_strategys", "shard",
            "sync_pipeline_shared_parameters", "Layout", "parameter_broadcast", "load_segmented_checkpoints",
-           "unified_safetensors", "set_op_strategy_config"]
+           "unified_safetensors", "load_distributed_checkpoint", "merge_sliced_parameter", "restore_group_info_list",
+           "build_searched_strategy", "set_op_strategy_config", "rank_list_for_convert",
+           "convert_checkpoint_by_rank", "convert_checkpoints"]

mindspore/parallel/_auto_parallel_context.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2023 Huawei Technologies Co., Ltd
+# Copyright 2020-2025 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@ from mindspore import context
 import mindspore.log as logger
 from mindspore.parallel._dp_allreduce_fusion import _set_fusion_strategy_by_idx, _set_fusion_strategy_by_size
 from mindspore.parallel._ps_context import _is_role_pserver
+from mindspore.parallel.shard import Layout
 from mindspore._c_expression import AutoParallelContext
 from mindspore._checkparam import args_type_check
 from mindspore import _checkparam as Validator
@@ -63,6 +64,7 @@ class _ParallelOptimizerConfig:
     GRADIENT_ACCUMULATION_SHARD = "gradient_accumulation_shard"
     PARALLEL_OPTIMIZER_THRESHOLD = "parallel_optimizer_threshold"
     OPTIMIZER_WEIGHT_SHARD_SIZE = "optimizer_weight_shard_size"
+    OPTIMIZER_LEVEL = "optimizer_level"
 
 
 class _PipelineConfig:
@@ -77,6 +79,8 @@ class _PipelineScheduler:
     PIPELINE_1F1B = "1f1b"
     PIPELINE_GPIPE = "gpipe"
     PIPELINE_SEQPIPE = "seqpipe"
+    PIPELINE_SEQVPP = "seqvpp"
+    PIPELINE_SEQSMARTVPP = "seqsmartvpp"
 
 
 class _AutoParallelContext:
@@ -100,6 +104,7 @@ class _AutoParallelContext:
     def __init__(self):
         self._context_handle = AutoParallelContext.get_instance()
         self._dataset_strategy_using_str = True
+        self._dataset_layout = None
 
     def check_context_handle(self):
         """
@@ -441,6 +446,9 @@ class _AutoParallelContext:
             raise ValueError("The context configuration parameter 'parallel_mode' only support 'stand_alone', "
                              "'data_parallel', 'hybrid_parallel', 'semi_auto_parallel' and 'auto_parallel', "
                              "but got the value : {}.".format(parallel_mode))
+        if run_mode == context.ParallelMode.DATA_PARALLEL and self.get_enable_parallel_optimizer():
+            logger.warning("'enable_parallel_optimizer' is not suggested in 'data_parallel' mode, "
+                           "consider using 'semi_auto_parallel' or 'auto_parallel' mode.")
 
     def get_parallel_mode(self):
         """Get parallel mode."""
@@ -585,6 +593,9 @@ class _AutoParallelContext:
         if not isinstance(dataset_strategy, tuple):
             raise TypeError("For 'set_auto_parallel_context', the argument 'dataset_strategy' "
                             "must be str or tuple type, but got the type : {}.".format(type(dataset_strategy)))
+        if dataset_strategy and isinstance(dataset_strategy[0], Layout):
+            self._set_dataset_strategy_layout(dataset_strategy)
+            return
         for ele in dataset_strategy:
             if not isinstance(ele, tuple):
                 raise TypeError("For 'set_auto_parallel_context', the element of argument "
@@ -599,8 +610,36 @@ class _AutoParallelContext:
         self._dataset_strategy_using_str = False
         self._context_handle.set_dataset_strategy(dataset_strategy)
 
+    def _set_dataset_strategy_layout(self, dataset_strategy):
+        """set dataset layout to c++ by using pybind."""
+        dataset_devmat = []
+        dataset_tensormap = []
+        dataset_alias_name = []
+        self._dataset_layout = dataset_strategy
+        for ele in dataset_strategy:
+            if not isinstance(ele, Layout):
+                raise TypeError(f"All the dataset_strategy elements should be Layout, but got {type(ele)}")
+            layout_to_dict = ele.to_dict()
+            dataset_devmat.append(layout_to_dict["device_matrix"])
+            dataset_alias_name.append(layout_to_dict["alias_name"])
+            if layout_to_dict["interleaved_parallel"]:
+                raise ValueError("For dataset_strategy, layout does not support interleaved_parallel")
+            tensor_map = []
+            for value in layout_to_dict["tensor_map"]:
+                if isinstance(value, tuple):
+                    tensor_map.append(value)
+                elif isinstance(value, int):
+                    tensor_map.append((value,))
+                else:
+                    raise TypeError(f"value in tensor map must be tuple or int, but got {type(value)}")
+            dataset_tensormap.append(tuple(tensor_map))
+        self._context_handle.set_dataset_layout(dataset_devmat, dataset_tensormap, dataset_alias_name)
+
+
     def get_dataset_strategy(self):
         """Get dataset sharding strategy."""
+        if self._dataset_layout is not None:
+            return self._dataset_layout
         self.check_context_handle()
         if self._dataset_strategy_using_str:
             if self._context_handle.get_full_batch():
@@ -888,6 +927,9 @@ class _AutoParallelContext:
                             "the argument 'enable_parallel_optimizer' must be bool, but got the type : {}."
                             .format(type(enable_parallel_optimizer)))
         self._context_handle.set_enable_parallel_optimizer(enable_parallel_optimizer)
+        if enable_parallel_optimizer and self.get_parallel_mode() == context.ParallelMode.DATA_PARALLEL:
+            logger.warning("'enable_parallel_optimizer' is not suggested in 'data_parallel' mode, "
+                           "consider using 'semi_auto_parallel' or 'auto_parallel' mode.")
 
     def set_force_fp32_communication(self, force_fp32_communication):
         """
@@ -918,7 +960,7 @@ class _AutoParallelContext:
 
             - pipeline_interleave(bool): Setting true enable interleave scheduler for pipeline parallelism. This
                                          scheduler requires more memory but less bubble.
-            - pipeline_scheduler(string): There are two choices, "1f1b" and "gpipe". default is "1f1b"
+            - pipeline_scheduler(str): There are two choices, "1f1b" and "gpipe". default is "1f1b"
 
               - 1f1b: It requires less memory and bubble ratio, for it run backward pass when corresponding forward pass
                       finished.
@@ -954,7 +996,9 @@ class _AutoParallelContext:
 
         Validator.check_string(pipeline_config[pp_scheduler], [_PipelineScheduler.PIPELINE_1F1B,
                                                                _PipelineScheduler.PIPELINE_GPIPE,
-                                                               _PipelineScheduler.PIPELINE_SEQPIPE])
+                                                               _PipelineScheduler.PIPELINE_SEQPIPE,
+                                                               _PipelineScheduler.PIPELINE_SEQVPP,
+                                                               _PipelineScheduler.PIPELINE_SEQSMARTVPP])
         if not pipeline_config[pp_interleave] and pipeline_config[pp_scheduler] != _PipelineScheduler.PIPELINE_1F1B:
             raise ValueError(f"When pipeline_interleave is False, {pp_scheduler} is not supported")
 
@@ -994,19 +1038,21 @@ class _AutoParallelContext:
                                                  shape[n] \* size(dtype). Non-negative. Unit: KB. Default: 64.
             - optimizer_weight_shard_size(int): Set the optimizer weight shard group size if you want to specific the
                                                 maximum group size across devices when the parallel optimizer is
-                                                enabled. The numerical range can be (0, device_num]. Default value
-                                                is -1, which means the optimizer weight shard group size will
-                                                the data parallel group of each parameter. Default -1.
-
+                                                enabled. The numerical range can be (0, device_num] or -1. If pipeline
+                                                parallelism is enabled, the numerical range is (0, device_num/stage]
+                                                or -1. Default value is -1, which means the optimizer weight shard
+                                                group size will be equal to the data parallel group of each parameter.
         """
         self.check_context_handle()
         grad_shard_name = _ParallelOptimizerConfig.GRADIENT_ACCUMULATION_SHARD
         threshold_name = _ParallelOptimizerConfig.PARALLEL_OPTIMIZER_THRESHOLD
         optimizer_weight_shard_size_name = _ParallelOptimizerConfig.OPTIMIZER_WEIGHT_SHARD_SIZE
+        optimizer_level_name = _ParallelOptimizerConfig.OPTIMIZER_LEVEL
 
         for config_name in parallel_optimizer_config:
             unknown_config = []
-            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name]:
+            if config_name not in [grad_shard_name, threshold_name, optimizer_weight_shard_size_name,
+                                   optimizer_level_name]:
                 unknown_config.append(config_name)
 
             if unknown_config:
@@ -1017,6 +1063,10 @@ class _AutoParallelContext:
                 parallel_optimizer_config[grad_shard_name], grad_shard_name, grad_shard_name)
             self._context_handle.set_grad_accumulation_shard(
                 parallel_optimizer_config[grad_shard_name])
+            if optimizer_level_name in parallel_optimizer_config \
+                    and parallel_optimizer_config[optimizer_level_name] != "level2":
+                raise ValueError(f"The optimizer_level is set as {parallel_optimizer_config[optimizer_level_name]}, "
+                                 "thus cannot set grad_accumulation_shard as True.")
 
         if threshold_name in parallel_optimizer_config:
             Validator.check_non_negative_int(
@@ -1026,8 +1076,23 @@ class _AutoParallelContext:
 
         if optimizer_weight_shard_size_name in parallel_optimizer_config:
             value = parallel_optimizer_config[optimizer_weight_shard_size_name]
-            Validator.check_positive_int(value)
-            self.set_optimizer_weight_shard_size(value)
+            if value != -1:
+                Validator.check_positive_int(value, prim_name="optimizer_weight_shard_size")
+                self.set_optimizer_weight_shard_size(value)
+
+        if optimizer_level_name in parallel_optimizer_config:
+            optimizer_level = parallel_optimizer_config[optimizer_level_name]
+            if optimizer_level not in ["level1", "level2", "level3"]:
+                raise ValueError("Optimizer level should in ['level1', 'level2', 'level3'], but got {}"
+                                 .format(optimizer_level))
+
+            if self._context_handle.get_grad_accumulation_shard() and optimizer_level != "level2":
+                raise ValueError("The grad_accumulation shard is set, thus cannot set optimizer_level != 'level2'")
+            if optimizer_level == "level2":
+                self._context_handle.set_grad_accumulation_shard(True)
+            if optimizer_level == "level3":
+                self._context_handle.set_zero3(True)
+                self._context_handle.set_grad_accumulation_shard(False)
 
     def get_grad_accumulation_shard(self):
         """Get grad accumulation shard."""
@@ -1136,6 +1201,7 @@ class _AutoParallelContext:
         self.check_context_handle()
         self._context_handle.reset()
         _ParallelFusionConfig.reset()
+        self._dataset_layout = None
 
     def _check_and_default_group(self, group):
         """Validate the given group, if group is empty, returns a default fusion group"""
@@ -1245,6 +1311,36 @@ class _AutoParallelContext:
             self.set_enable_all_gather_fusion(openstate)
             self.set_enable_reduce_scatter_fusion(openstate)
 
+    def set_auto_parallel_new_interface(self, auto_parallel_new_interface):
+        """
+        Set AutoParallel(cell) new interface flag.
+
+        Args:
+            auto_parallel_new_interface (bool): Mark whether to use the new interface.
+        """
+        self.check_context_handle()
+        self._context_handle.set_auto_parallel_new_interface(auto_parallel_new_interface)
+
+    def get_auto_parallel_new_interface(self):
+        """Get auto_parallel_new_interface."""
+        self.check_context_handle()
+        return self._context_handle.get_auto_parallel_new_interface()
+
+    def set_init_param_in_compile(self, init_param_in_compile):
+        """
+        Set flag marking whether to init parameters in compiling process.
+
+        Args:
+            init_param_in_compile (bool): Mark whether to init parameters in compiling process.
+        """
+        self.check_context_handle()
+        self._context_handle.set_init_param_in_compile(init_param_in_compile)
+
+    def get_init_param_in_compile(self):
+        """Get init_param_in_compile."""
+        self.check_context_handle()
+        return self._context_handle.get_init_param_in_compile()
+
 _AUTO_PARALLEL_CONTEXT = None
 
 
@@ -1295,7 +1391,10 @@ _set_auto_parallel_context_func_map = {
     "comm_fusion": auto_parallel_context().set_comm_fusion,
     "dump_local_norm": auto_parallel_context().set_dump_local_norm,
     "dump_local_norm_path": auto_parallel_context().set_dump_local_norm_path,
-    "dump_device_local_norm": auto_parallel_context().set_dump_device_local_norm}
+    "dump_device_local_norm": auto_parallel_context().set_dump_device_local_norm,
+    "auto_parallel_new_interface": auto_parallel_context().set_auto_parallel_new_interface,
+    "init_param_in_compile": auto_parallel_context().set_init_param_in_compile}
+
 
 _get_auto_parallel_context_func_map = {
     "device_num": auto_parallel_context().get_device_num,
@@ -1330,7 +1429,9 @@ _get_auto_parallel_context_func_map = {
     "full_batch_is_set": auto_parallel_context().get_full_batch_is_set,
     "dump_local_norm": auto_parallel_context().get_dump_local_norm,
     "dump_local_norm_path": auto_parallel_context().get_dump_local_norm_path,
-    "dump_device_local_norm": auto_parallel_context().get_dump_device_local_norm}
+    "dump_device_local_norm": auto_parallel_context().get_dump_device_local_norm,
+    "auto_parallel_new_interface": auto_parallel_context().get_auto_parallel_new_interface,
+    "init_param_in_compile": auto_parallel_context().get_init_param_in_compile}
 
 
 @args_type_check(device_num=int, global_rank=int, gradients_mean=bool, gradient_fp32_sync=bool,
@@ -1472,6 +1573,33 @@ def _get_auto_parallel_context(attr_key):
     return get_func()
 
 
+def _get_all_auto_parallel_context():
+    """get auto parallel context before reset"""
+    _auto_paralell_context_value_map = {}
+    _pipeline_config = {}
+    for key, value in _get_auto_parallel_context_func_map.items():
+        if key == "pipeline_interleave":
+            _pipeline_config[key] = value()
+        elif key == "pipeline_scheduler":
+            _pipeline_config[key] = value()
+        else:
+            _auto_paralell_context_value_map[key] = value()
+    return _auto_paralell_context_value_map, _pipeline_config
+
+
+def _recover_auto_parallel_context(context_value_map, pp_config):
+    """set auto parallel context after transformation"""
+    # set the same auto parallel context after transform
+    from mindspore.context import reset_auto_parallel_context
+    reset_auto_parallel_context()
+    for key, value in context_value_map.items():
+        # list is empty or full_batch_is_set is not needed to set
+        if (isinstance(value, list) and not value) or (key == "full_batch_is_set"):
+            continue
+        _set_auto_parallel_context_func_map[key](value)
+    _set_auto_parallel_context_func_map["pipeline_config"](pp_config)
+
+
 def _reset_auto_parallel_context():
     """
     Reset auto parallel context attributes to the default values: