Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / dfunction.py: 88%

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"""Custom distributed autograd function base class.""" 

16from __future__ import annotations 

17 

18from typing import Any, Callable 

19 

20import torch 

21 

22from hyper_parallel.core.dtensor.dtensor import DTensor 

23from hyper_parallel.core.shard._op_dispatch import _OP_DISPATCHER 

24 

25 

26class _LocalCallable: 

27 """Named callable wrapper that exposes the op name to Torch dispatch. 

28 

29 Torch operation-name resolution inspects ``__name__``. Setting it here lets 

30 ``_OP_DISPATCHER`` look up the matching ``DistributedOp``. 

31 

32 Args: 

33 fn: The underlying callable to invoke. 

34 op_name: Canonical op name matching the registered ``DistributedOp``. 

35 """ 

36 

37 def __init__(self, fn: Callable, op_name: str) -> None: 

38 """Initialize the wrapper with a callable and its registered op name.""" 

39 self._fn = fn 

40 self.__name__ = op_name 

41 

42 def __call__(self, *args: Any, **kwargs: Any) -> Any: 

43 """Forward arguments to the wrapped callable.""" 

44 return self._fn(*args, **kwargs) 

45 

46 

47class DFunction(torch.autograd.Function): 

48 """Torch base class for user-defined distributed autograd functions. 

49 

50 Subclass this class and implement ``forward`` and ``backward`` as 

51 ``@staticmethod`` methods that operate on **local** tensors. To enable 

52 multi-device support, also set ``_op_name`` to a string that matches the 

53 ``op_name`` of a registered ``DistributedOp`` subclass. 

54 

55 Dispatch behaviour: 

56 

57 * **No DTensor inputs** — calls ``super().apply()`` directly, going straight 

58 into Torch autograd (single-device path). 

59 * **DTensor inputs + ``_op_name`` set** — delegates to 

60 ``_OP_DISPATCHER.dispatch``. The dispatcher extracts local tensors, calls 

61 ``DistributedOp.infer_layout`` to derive the output layout, invokes the 

62 local callable (which re-enters ``apply`` with plain tensors, triggering 

63 the single-device path), and wraps the result as a ``DTensor``. 

64 

65 Note on non-tensor positional arguments: 

66 Non-tensor positional arguments are preserved via the three-phase dispatch 

67 flow (``preprocess()`` → ``infer_layout(cache_values)`` → 

68 ``get_expand_impl``). Implement ``preprocess()`` in your ``DistributedOp`` 

69 to collect non-tensor values into ``cache_values``. 

70 

71 Example:: 

72 

73 class MyAddDistOp(DistributedOp): 

74 def __init__(self): 

75 super().__init__("MyAdd") 

76 

77 def infer_layout(self, layouts, extra_args=None): 

78 return layouts[0] 

79 

80 MyAddDistOp() # instantiation triggers registration 

81 

82 class MyAdd(DFunction): 

83 _op_name = "MyAdd" 

84 

85 @staticmethod 

86 def forward(ctx, x, y): 

87 ctx.save_for_backward(x, y) 

88 return x + y 

89 

90 @staticmethod 

91 def backward(ctx, grad): 

92 return grad, grad 

93 

94 result = MyAdd.apply(x, y) # plain tensors → single-device 

95 result = MyAdd.apply(dtensor_x, dtensor_y) # DTensors → distributed 

96 """ 

97 

98 _op_name: str = None 

99 

100 @staticmethod 

101 def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: 

102 """Override in subclass to define the forward computation on local tensors.""" 

103 raise NotImplementedError("Subclasses must implement forward()") 

104 

105 @staticmethod 

106 def backward(ctx: Any, *grad_outputs: Any) -> Any: 

107 """Override in subclass to define the backward computation on local tensors.""" 

108 raise NotImplementedError("Subclasses must implement backward()") 

109 

110 @classmethod 

111 def apply(cls, *args: Any, **kwargs: Any) -> Any: 

112 """Execute the function, routing to distributed dispatch when DTensor inputs are present. 

113 

114 Args: 

115 *args: Positional arguments forwarded to ``forward``. May mix 

116 ``DTensor`` and plain ``Tensor``. 

117 **kwargs: Keyword arguments forwarded to ``forward``. 

118 

119 Returns: 

120 Output of the forward computation. Wrapped as ``DTensor`` when the 

121 distributed dispatch path is taken and the ``DistributedOp`` infers a 

122 valid output layout. 

123 

124 Raises: 

125 ValueError: If ``DTensor`` inputs are detected but ``_op_name`` is not 

126 set on the subclass. 

127 """ 

128 has_dtensor = any(isinstance(a, DTensor) for a in args) 

129 if has_dtensor: 

130 if cls._op_name is None: 

131 raise ValueError( 

132 f"{cls.__name__} received DTensor inputs but '_op_name' is not set. " 

133 "Set '_op_name' on the subclass and register a matching DistributedOp." 

134 ) 

135 return _OP_DISPATCHER.dispatch(cls._get_local_callable(), args, kwargs) 

136 return super().apply(*args, **kwargs) 

137 

138 @classmethod 

139 def _get_local_callable(cls) -> _LocalCallable: 

140 """Return (and lazily create) the named local callable for this subclass. 

141 

142 The callable re-enters ``cls.apply`` with plain tensors so the single-device 

143 autograd path is taken without recursion. Cached per subclass in 

144 ``cls.__dict__`` to avoid repeated object creation. 

145 

146 Returns: 

147 A ``_LocalCallable`` whose ``__name__`` equals ``cls._op_name``, 

148 enabling Torch operation-name resolution to find the correct 

149 ``DistributedOp``. 

150 """ 

151 if '_local_callable' not in cls.__dict__: 

152 def _local_fn(*a, **kw): 

153 # Called by _OP_DISPATCHER with extracted local (non-DTensor) tensors. 

154 # has_dtensor will be False here, so super().apply() is taken directly. 

155 return cls.apply(*a, **kw) 

156 

157 cls._local_callable = _LocalCallable(_local_fn, cls._op_name) 

158 return cls._local_callable