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, List 

26 

27import torch.distributed._tensor as torch_dt 

28 

29logging.basicConfig(level=logging.INFO) 

30logger = logging.getLogger(__name__) 

31 

32# Global backend flag 

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

34 

35 

36class _NeverMatch: 

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

38 __slots__ = () 

39 

40 

41# Lazy-export cache 

42_LAZY_CACHE: dict = {} 

43 

44 

45def _invalidate_lazy_cache() -> None: 

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

47 _LAZY_CACHE.clear() 

48 

49 

50def detect_dtensor_backend( 

51 adamw_params: List[Any], 

52 muon_params: List[Any], 

53) -> str: 

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

55 global _DTENSOR_BACKEND # pylint: disable=global-statement 

56 

57 sample_param = _extract_first_param(muon_params) 

58 

59 if sample_param is None: 

60 sample_param = _extract_first_param(adamw_params) 

61 

62 if sample_param is None: 

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

64 _DTENSOR_BACKEND = "hyper" 

65 _invalidate_lazy_cache() 

66 return _DTENSOR_BACKEND 

67 

68 param_cls_module = type(sample_param).__module__ 

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

70 _DTENSOR_BACKEND = "torch" 

71 else: 

72 _DTENSOR_BACKEND = "hyper" 

73 

74 logger.info("Detected DTensor backend: '%s'.", _DTENSOR_BACKEND) 

75 _invalidate_lazy_cache() 

76 return _DTENSOR_BACKEND 

77 

78 

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

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

81 for group in param_groups: 

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

83 for p in params: 

84 return p 

85 

86 for p in param_groups: 

87 return p 

88 

89 return None 

90 

91 

92# Accessor functions 

93def get_dtensor_cls(): 

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

95 if _DTENSOR_BACKEND == "torch": 

96 return torch_dt.DTensor 

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

98 return DTensor 

99 

100 

101def get_device_mesh_cls(): 

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

103 if _DTENSOR_BACKEND == "torch": 

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

105 return DeviceMesh 

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

107 return DeviceMesh 

108 

109 

110def get_shard_cls(): 

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

112 if _DTENSOR_BACKEND == "torch": 

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

114 return Shard 

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

116 return Shard 

117 

118 

119def get_replicate_cls(): 

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

121 if _DTENSOR_BACKEND == "torch": 

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

123 return Replicate 

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

125 return Replicate 

126 

127 

128def get_strided_shard_cls(): 

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

130 if _DTENSOR_BACKEND == "torch": 

131 return _NeverMatch 

132 

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

134 return StridedShard 

135 

136 

137# DTensor union type resolver 

138def _import_hyper_dtensor(): 

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

140 try: 

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

142 return DTensor 

143 except ImportError: 

144 return torch_dt.DTensor 

145 

146 

147def _resolve_dtensor_union(): 

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

149 return torch_dt.DTensor | _import_hyper_dtensor() 

150 

151 

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

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

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

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

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

157 

158 

159# lazy exports 

160_LAZY_RESOLVERS = { 

161 "DTensor": _resolve_dtensor_union, 

162 "DeviceMesh": get_device_mesh_cls, 

163 "Shard": get_shard_cls, 

164 "Replicate": get_replicate_cls, 

165 "StridedShard": get_strided_shard_cls, 

166} 

167 

168 

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

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

171 resolver = _LAZY_RESOLVERS.get(name) 

172 if resolver is not None: 

173 value = _LAZY_CACHE.get(name) 

174 if value is None: 

175 value = resolver() 

176 _LAZY_CACHE[name] = value 

177 return value 

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

179 

180 

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

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

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