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

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"""Utilities for :meth:`DTensor.from_local` with ``run_check=True`` (PyTorch parity).""" 

16from __future__ import annotations 

17 

18from typing import Optional, Sequence, Tuple 

19 

20from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

21from hyper_parallel.core.dtensor.placement_types import Placement 

22from hyper_parallel.platform import get_platform 

23 

24platform = get_platform() 

25Tensor = platform.Tensor 

26 

27 

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

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_broadcast( 

39 tensor: Tensor, 

40 mesh: DeviceMesh, 

41 mesh_dim: int, 

42 *, 

43 group_src: int = 0, 

44) -> Tensor: 

45 """Broadcast *tensor* along one mesh dimension.""" 

46 _ensure_mesh_process_groups(mesh) 

47 group = mesh.get_group(mesh_dim) 

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

49 tensor = tensor.contiguous() 

50 rank_list = mesh.get_rank_list_along_axis(mesh_dim) 

51 src = rank_list[group_src] 

52 platform.broadcast(tensor, src, group=group) 

53 return tensor 

54 

55 

56def _tensor_meta(local_tensor: Tensor, *, check_shape_stride: bool) -> dict: 

57 meta = { 

58 "dtype": str(local_tensor.dtype), 

59 "requires_grad": bool(getattr(local_tensor, "requires_grad", False)), 

60 } 

61 if check_shape_stride: 

62 meta["shape"] = tuple(local_tensor.shape) 

63 if hasattr(local_tensor, "stride"): 

64 meta["stride"] = tuple(local_tensor.stride()) 

65 return meta 

66 

67 

68def check_tensor_meta( 

69 local_tensor: Tensor, 

70 group, 

71 group_size: int, 

72 *, 

73 check_shape_stride: bool, 

74) -> None: 

75 """Gather tensor metadata across *group* and verify consistency.""" 

76 local_meta = _tensor_meta(local_tensor, check_shape_stride=check_shape_stride) 

77 gathered = [None] * group_size 

78 platform.all_gather_object(gathered, local_meta, group=group) 

79 if not all(meta == local_meta for meta in gathered if meta is not None): 

80 raise ValueError( 

81 "Inconsistent tensor metadata across ranks in from_local(run_check=True): " 

82 f"local={local_meta}, gathered={gathered}" 

83 ) 

84 

85 

86def _mesh_check_group(device_mesh: DeviceMesh): 

87 """Return the process group and size covering all ranks in *device_mesh*.""" 

88 _ensure_mesh_process_groups(device_mesh) 

89 if device_mesh.ndim == 1: 

90 return device_mesh.get_group(0), device_mesh.size(0) 

91 flat_mesh = device_mesh.flatten() 

92 _ensure_mesh_process_groups(flat_mesh) 

93 return flat_mesh.get_group(0), flat_mesh.size(0) 

94 

95 

96def run_from_local_checks( 

97 local_tensor: Tensor, 

98 device_mesh: DeviceMesh, 

99 resolved_placements: Sequence[Placement], 

100 *, 

101 shape: Optional[Tuple[int, ...]] = None, 

102 stride: Optional[Tuple[int, ...]] = None, 

103) -> None: 

104 """Validate local shards and align replicate placements before wrapping as DTensor.""" 

105 check_shape_stride = shape is None and stride is None 

106 group, group_size = _mesh_check_group(device_mesh) 

107 check_tensor_meta( 

108 local_tensor, 

109 group, 

110 group_size, 

111 check_shape_stride=check_shape_stride, 

112 ) 

113 for mesh_dim, placement in enumerate(resolved_placements): 

114 if placement.is_replicate(): 

115 mesh_broadcast(local_tensor, device_mesh, mesh_dim)