wehub-resource-sync
146 lines
4.8 KiB
Python
146 lines
4.8 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
|
|
import numpy as np
|
|
|
|
from paddle.framework import core
|
|
|
|
|
|
class ProcessMesh(core.ProcessMesh):
|
|
r"""
|
|
The class `ProcessMesh` describes the topology of logical processes.
|
|
|
|
Args:
|
|
mesh (list|numpy.array): an N-dimensional array describes the topology
|
|
of logical processes.
|
|
dim_names (list, optional): the i-th element of this list gives the name of the
|
|
i-th dimension.
|
|
|
|
Returns:
|
|
None
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> import paddle.distributed as dist
|
|
>>> paddle.enable_static()
|
|
|
|
>>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]])
|
|
>>> assert mesh.shape == [2, 3]
|
|
>>> assert mesh.process_ids == [2, 4, 5, 0, 1, 3]
|
|
|
|
"""
|
|
|
|
def __init__(self, mesh, dim_names=None):
|
|
if not isinstance(mesh, list) and not isinstance(mesh, np.ndarray):
|
|
raise ValueError(
|
|
'The mesh must be an instance of list or np.ndarray.'
|
|
)
|
|
if isinstance(mesh, list):
|
|
mesh = np.array(mesh)
|
|
|
|
self._mesh = mesh
|
|
|
|
self._shape = list(self._mesh.shape)
|
|
|
|
self._process_ids = self._mesh.flatten().tolist()
|
|
if not all(isinstance(p, int) for p in self._process_ids):
|
|
raise ValueError("All elements of the mesh must be integer")
|
|
|
|
if min(self._process_ids) < 0:
|
|
raise ValueError('All elements of the mesh must be >= 0.')
|
|
|
|
unique_process_ids = set(self._process_ids)
|
|
if len(unique_process_ids) != len(self._process_ids):
|
|
raise ValueError('All elements of the mesh must be unique.')
|
|
|
|
if dim_names is not None:
|
|
assert len(dim_names) == len(self._shape), (
|
|
"The length of dims_names must be same as the shape of the mesh."
|
|
)
|
|
self._dim_names = dim_names
|
|
else:
|
|
self._dim_names = ["d" + str(i) for i in range(len(self._shape))]
|
|
|
|
# Follow the requirement for using pybind11
|
|
core.ProcessMesh.__init__(
|
|
self, self._shape, self._process_ids, self._dim_names
|
|
)
|
|
|
|
@property
|
|
def mesh(self):
|
|
return self._mesh
|
|
|
|
|
|
def compute_compatible_process_mesh(process_meshes):
|
|
"""Compute the compatible process mesh given a list of process meshes."""
|
|
if not process_meshes:
|
|
return None
|
|
|
|
def _compute_compatible_of_two_process_meshes(pm1, pm2):
|
|
if pm1 is None:
|
|
return True, pm2
|
|
if pm2 is None:
|
|
return True, pm1
|
|
if pm1 == pm2:
|
|
return True, pm1
|
|
if pm1.process_ids == pm2.process_ids:
|
|
if len(pm1.shape) >= len(pm2.shape):
|
|
return True, pm1
|
|
else:
|
|
return True, pm2
|
|
process_set1 = set(pm1.process_ids)
|
|
process_set2 = set(pm2.process_ids)
|
|
if process_set1.issubset(process_set2):
|
|
return True, pm2
|
|
if process_set2.issubset(process_set1):
|
|
return True, pm1
|
|
return False, None
|
|
|
|
compatible_result = None
|
|
for process_mesh in process_meshes:
|
|
(
|
|
compatible,
|
|
compatible_result,
|
|
) = _compute_compatible_of_two_process_meshes(
|
|
compatible_result, process_mesh
|
|
)
|
|
if not compatible:
|
|
return None
|
|
if compatible_result.empty():
|
|
return None
|
|
if isinstance(compatible_result, core.ProcessMesh):
|
|
mesh = np.array(compatible_result.process_ids).reshape(
|
|
compatible_result.shape
|
|
)
|
|
return ProcessMesh(mesh, compatible_result.dim_names)
|
|
elif isinstance(compatible_result, ProcessMesh):
|
|
return ProcessMesh(compatible_result.mesh, compatible_result.dim_names)
|
|
else:
|
|
raise ValueError("Unrecognized ProcessMesh.")
|
|
|
|
|
|
def merge_process_mesh(process_meshes):
|
|
"""Merge a list of process meshes."""
|
|
merged_process_mesh = None
|
|
merged_process_ids = set()
|
|
for process_mesh in process_meshes:
|
|
if process_mesh is not None:
|
|
process_ids = set(process_mesh.process_ids)
|
|
merged_process_ids = merged_process_ids.union(process_ids)
|
|
if len(merged_process_ids) != 0:
|
|
merged_process_mesh = ProcessMesh(list(merged_process_ids))
|
|
return merged_process_mesh
|