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).""" 

16# pylint: disable=C9006,C9007 

17from __future__ import annotations 

18 

19from typing import Optional, Sequence, Tuple 

20 

21import torch 

22import torch.distributed as dist 

23 

24from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

25from hyper_parallel.core.dtensor.placement_types import Placement 

26 

27Tensor = torch.Tensor 

28 

29 

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

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

32 return 

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

34 mesh._mesh_shape, 

35 mesh.mesh_dim_names, 

36 mesh._rank_list, 

37 ) 

38 

39 

40def mesh_broadcast( 

41 tensor: Tensor, 

42 mesh: DeviceMesh, 

43 mesh_dim: int, 

44 *, 

45 group_src: int = 0, 

46) -> Tensor: 

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

48 _ensure_mesh_process_groups(mesh) 

49 group = mesh.get_group(mesh_dim) 

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

51 tensor = tensor.contiguous() 

52 rank_list = mesh.get_rank_list_along_axis(mesh_dim) 

53 src = rank_list[group_src] 

54 dist.broadcast(tensor, src, group=group) 

55 return tensor 

56 

57 

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

59 meta = { 

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

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

62 } 

63 if check_shape_stride: 

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

65 if hasattr(local_tensor, "stride"): 

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

67 return meta 

68 

69 

70def check_tensor_meta( 

71 local_tensor: Tensor, 

72 group, 

73 group_size: int, 

74 *, 

75 check_shape_stride: bool, 

76) -> None: 

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

78 local_meta = _tensor_meta(local_tensor, check_shape_stride=check_shape_stride) 

79 gathered = [None] * group_size 

80 dist.all_gather_object(gathered, local_meta, group=group) 

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

82 raise ValueError( 

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

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

85 ) 

86 

87 

88def _mesh_check_group(device_mesh: DeviceMesh): 

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

90 _ensure_mesh_process_groups(device_mesh) 

91 if device_mesh.ndim == 1: 

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

93 flat_mesh = device_mesh.flatten() 

94 _ensure_mesh_process_groups(flat_mesh) 

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

96 

97 

98def run_from_local_checks( 

99 local_tensor: Tensor, 

100 device_mesh: DeviceMesh, 

101 resolved_placements: Sequence[Placement], 

102 *, 

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

104) -> None: 

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

106 check_shape_stride = shape is None 

107 group, group_size = _mesh_check_group(device_mesh) 

108 check_tensor_meta( 

109 local_tensor, 

110 group, 

111 group_size, 

112 check_shape_stride=check_shape_stride, 

113 ) 

114 for mesh_dim, placement in enumerate(resolved_placements): 

115 if placement.is_replicate(): 

116 mesh_broadcast(local_tensor, device_mesh, mesh_dim)