Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / _collective_utils.py: 36%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Mesh-scoped collectives for :func:`distribute_tensor` (PyTorch DTensor parity).""" 

16from __future__ import annotations 

17 

18from typing import Sequence 

19 

20from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

21from hyper_parallel.platform import get_platform 

22 

23platform = get_platform() 

24Tensor = platform.Tensor 

25 

26 

27def _ensure_mesh_process_groups(mesh: DeviceMesh) -> None: 

28 """Lazily create per-axis process groups when mesh was built with ``init_backend=False``.""" 

29 if hasattr(mesh, "_dim_group_names") and mesh._dim_group_names is not None: 

30 return 

31 mesh._dim_group_names = DeviceMesh._init_process_groups( # pylint: disable=protected-access 

32 mesh._mesh_shape, 

33 mesh.mesh_dim_names, 

34 mesh._rank_list, 

35 ) 

36 

37 

38def mesh_scatter( 

39 output: Tensor, 

40 scatter_list: Sequence[Tensor], 

41 mesh: DeviceMesh, 

42 mesh_dim: int, 

43 *, 

44 group_src: int = 0, 

45) -> Tensor: 

46 """Scatter tensor chunks along one mesh dimension (PyTorch ``mesh_scatter`` parity).""" 

47 _ensure_mesh_process_groups(mesh) 

48 group = mesh.get_group(mesh_dim) 

49 contiguous_list = [ 

50 chunk.contiguous() if hasattr(chunk, "is_contiguous") and not chunk.is_contiguous() else chunk 

51 for chunk in scatter_list 

52 ] 

53 if platform.get_group_rank(group) == group_src: 

54 platform.scatter(output, list(contiguous_list), group=group, group_src=group_src) 

55 else: 

56 platform.scatter(output, None, group=group, group_src=group_src) 

57 return output 

58 

59 

60def mesh_broadcast( 

61 tensor: Tensor, 

62 mesh: DeviceMesh, 

63 mesh_dim: int, 

64 *, 

65 group_src: int = 0, 

66) -> Tensor: 

67 """Broadcast a tensor along one mesh dimension (PyTorch ``mesh_broadcast`` parity).""" 

68 _ensure_mesh_process_groups(mesh) 

69 group = mesh.get_group(mesh_dim) 

70 if hasattr(tensor, "is_contiguous") and not tensor.is_contiguous(): 

71 tensor = tensor.contiguous() 

72 platform.broadcast(tensor, group=group, group_src=group_src) 

73 return tensor