refactor parallel module (mindspore-lab#536)

qurikuduo · May 28, 2023 · 4dec44b · 4dec44b
1 parent 69fc70d
commit 4dec44b
Show file tree

Hide file tree

Showing 8 changed files with 287 additions and 322 deletions.
diff --git a/mindnlp/_legacy/ops/__init__.py b/mindnlp/_legacy/ops/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Operators for legacy version"""
+
+from . import parallel
+from .parallel import *
+
+__all__ = []
+__all__.extend(parallel.__all__)
diff --git a/mindnlp/_legacy/ops/parallel.py b/mindnlp/_legacy/ops/parallel.py
@@ -0,0 +1,31 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Parallel Operators"""
+from mindspore import Parameter
+from mindspore.ops import AllGather as msAllGather
+from mindspore.ops.primitive import _run_op
+
+class AllGather(msAllGather):
+    """AllGather op."""
+    def __call__(self, *args):
+        should_elim, output = self.check_elim(*args)
+        for arg in args:
+            if isinstance(arg, Parameter) and arg.has_init:
+                arg.init_data()
+        if should_elim:
+            return output
+        return _run_op(self, self.name, args)
+
+__all__ = ['AllGather']
diff --git a/mindnlp/modules/parallel/__init__.py b/mindnlp/modules/parallel/__init__.py
diff --git a/mindnlp/modules/parallel/mappings.py b/mindnlp/modules/parallel/mappings.py
diff --git a/mindnlp/parallel/__init__.py b/mindnlp/parallel/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2023 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.
+# ============================================================================
+"""MindNLP Parallel modules, which is ported from Megatron."""
+
+from . import layers
+from .layers import *
+
+__all__ = []
+__all__.extend(layers.__all__)