Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / dtensor_compat.py: 0%

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

16"""DTensor backend compatibility layer for the optimizer module. 

17 

18Provides lazy exports (PEP 562) for DTensor, DeviceMesh, Shard, Replicate,  

19and StridedShard based on the detected backend ('torch' or 'hyper'). 

20""" 

21 

22from __future__ import annotations 

23 

24import logging 

25from typing import Any, Dict, List 

26 

27import torch.distributed._tensor as torch_dt 

28 

29logger = logging.getLogger(__name__) 

30 

31# Global backend flag 

32_DTENSOR_BACKEND: str = "hyper" # "hyper" or "torch" 

33 

34 

35class _NeverMatch: 

36 """Safe fallback class that always returns False for ``isinstance()``.""" 

37 __slots__ = () 

38 

39 

40# Lazy-export cache 

41_LAZY_CACHE: Dict[str, Any] = {} 

42 

43 

44def _invalidate_lazy_cache() -> None: 

45 """Clear the lazy-export cache to rebuild on next access.""" 

46 _LAZY_CACHE.clear() 

47 

48 

49def detect_dtensor_backend( 

50 adamw_params: List[Any], 

51 muon_params: List[Any], 

52) -> str: 

53 """Detect and set the DTensor backend ('torch' or 'hyper') from parameter lists.""" 

54 global _DTENSOR_BACKEND # pylint: disable=global-statement 

55 

56 sample_param = _extract_first_param(muon_params) 

57 

58 if sample_param is None: 

59 sample_param = _extract_first_param(adamw_params) 

60 

61 if sample_param is None: 

62 logger.info_rank0("No parameters found for backend detection; defaulting to 'hyper'.") 

63 _DTENSOR_BACKEND = "hyper" 

64 _invalidate_lazy_cache() 

65 return _DTENSOR_BACKEND 

66 

67 param_cls_module = type(sample_param).__module__ 

68 if param_cls_module.startswith("torch.distributed"): 

69 _DTENSOR_BACKEND = "torch" 

70 else: 

71 _DTENSOR_BACKEND = "hyper" 

72 

73 logger.info_rank0("Detected DTensor backend: '%s'.", _DTENSOR_BACKEND) 

74 _invalidate_lazy_cache() 

75 return _DTENSOR_BACKEND 

76 

77 

78def _extract_first_param(param_groups: List[Any]) -> Any: 

79 """Return the first parameter from a list of param groups, or None.""" 

80 for group in param_groups: 

81 params = group.get("params", []) if isinstance(group, dict) else [] 

82 for p in params: 

83 return p 

84 

85 for p in param_groups: 

86 return p 

87 

88 return None 

89 

90 

91# Accessor functions 

92def get_dtensor_cls(): 

93 """Return the DTensor class for the active backend.""" 

94 if _DTENSOR_BACKEND == "torch": 

95 return torch_dt.DTensor 

96 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel 

97 return DTensor 

98 

99 

100def get_device_mesh_cls(): 

101 """Return the DeviceMesh class for the active backend.""" 

102 if _DTENSOR_BACKEND == "torch": 

103 from torch.distributed.device_mesh import DeviceMesh # pylint: disable=import-outside-toplevel 

104 return DeviceMesh 

105 from hyper_parallel.core.dtensor.device_mesh import DeviceMesh # pylint: disable=import-outside-toplevel 

106 return DeviceMesh 

107 

108 

109def get_shard_cls(): 

110 """Return the Shard placement class for the active backend.""" 

111 if _DTENSOR_BACKEND == "torch": 

112 from torch.distributed._tensor.placement_types import Shard # pylint: disable=import-outside-toplevel 

113 return Shard 

114 from hyper_parallel.core.dtensor.placement_types import Shard # pylint: disable=import-outside-toplevel 

115 return Shard 

116 

117 

118def get_replicate_cls(): 

119 """Return the Replicate placement class for the active backend.""" 

120 if _DTENSOR_BACKEND == "torch": 

121 from torch.distributed._tensor.placement_types import Replicate # pylint: disable=import-outside-toplevel 

122 return Replicate 

123 from hyper_parallel.core.dtensor.placement_types import Replicate # pylint: disable=import-outside-toplevel 

124 return Replicate 

125 

126 

127def get_strided_shard_cls(): 

128 """Return the StridedShard placement class. Returns _NEVER_MATCH for 'torch'.""" 

129 if _DTENSOR_BACKEND == "torch": 

130 return _NeverMatch 

131 

132 from hyper_parallel.core.dtensor.placement_types import StridedShard # pylint: disable=import-outside-toplevel 

133 return StridedShard 

134 

135 

136# DTensor union type resolver 

137def _import_hyper_dtensor(): 

138 """Import hyper DTensor class; return torch DTensor as fallback.""" 

139 try: 

140 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=import-outside-toplevel 

141 return DTensor 

142 except ImportError: 

143 return torch_dt.DTensor 

144 

145 

146def _resolve_dtensor_union(): 

147 """Build ``torch_dt.DTensor | hyper_dt.DTensor`` on demand.""" 

148 return torch_dt.DTensor | _import_hyper_dtensor() 

149 

150 

151def to_local_if_dtensor(tensor: Any) -> Any: 

152 """Return the local shard if `tensor` is a DTensor, otherwise return as-is.""" 

153 # Use resolver directly for internal module lookups instead of lazy-loaded DTensor 

154 dtensor_type = _LAZY_CACHE.get("DTensor") or _resolve_dtensor_union() 

155 return tensor.to_local() if isinstance(tensor, dtensor_type) else tensor 

156 

157 

158# lazy exports 

159_LAZY_RESOLVERS = { 

160 "DTensor": _resolve_dtensor_union, 

161 "DeviceMesh": get_device_mesh_cls, 

162 "Shard": get_shard_cls, 

163 "Replicate": get_replicate_cls, 

164 "StridedShard": get_strided_shard_cls, 

165} 

166 

167 

168def __getattr__(name): # type: ignore[no-untyped-def] # pylint: disable=invalid-name 

169 """Resolve module attributes on first access.""" 

170 resolver = _LAZY_RESOLVERS.get(name) 

171 if resolver is not None: 

172 value = _LAZY_CACHE.get(name) 

173 if value is None: 

174 value = resolver() 

175 _LAZY_CACHE[name] = value 

176 return value 

177 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

178 

179 

180def __dir__(): # type: ignore[no-untyped-def] # pylint: disable=invalid-name 

181 """Include lazy-exported names in dir() for IDE autocomplete.""" 

182 return list(globals().keys()) + list(_LAZY_RESOLVERS.keys())