Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / function_module.py: 80%

1# Copyright 2025-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"""FunctionModule — a module shell for autograd.Function (tutorial §10.8). 

16 

17The boundary mechanism acts at the granularity of ``nn.Module.forward``; a 

18custom ``autograd.Function`` invoked as a bare ``A.apply(...)`` call is 

19invisible to the framework (no FQN means no spec mount point). 

20``FunctionModule`` gives a Function a module form: parameterless, stateless, 

21forward passes through to ``fn.apply``, and backward goes through the 

22Function's own static backward — the shell is transparent to autograd. 

23 

24Usage (declare the boundary in plan_overrides; custom Functions are not 

25covered by DTensor dispatch overrides, so ``region_dispatch=False`` is 

26required):: 

27 

28 self.a_fn = FunctionModule(A) 

29 plan_overrides={"...a_fn": ModuleShardingSpec( 

30 params={}, region_dispatch=False, 

31 in_src={"x": ...}, in_dst={"x": ...}, 

32 out_src={...}, out_dst={...})} 

33 

34Contract key binding (same as all boundaries, ``_bind_input_indices``): first 

35bind by the shell forward's parameter names; **a single-input contract falls 

36back to binding to the 0th positional argument** — so a single-tensor input 

37can be passed straight through this shell's ``*args``. For multiple inputs 

38(e.g. an extra weight tensor), subclass and provide an explicit signature:: 

39 

40 class SeqNormWithWeight(FunctionModule): 

41 def forward(self, x, weight): 

42 return self._fn.apply(x, weight) 

43""" 

44 

45from typing import Any 

46 

47from torch import nn 

48 

49 

50class FunctionModule(nn.Module): 

51 """Wrap an ``autograd.Function`` class as an ``nn.Module``. 

52 

53 This boundary FQN mount point is transparent to autograd: backward is the 

54 Function's own static ``backward``. 

55 """ 

56 

57 def __init__(self, fn: type) -> None: 

58 """Store the ``autograd.Function`` class wrapped by this module.""" 

59 super().__init__() 

60 self._fn = fn 

61 

62 def forward(self, *args: Any, **kwargs: Any) -> Any: 

63 """Invoke the wrapped Function's ``apply`` with the given inputs.""" 

64 return self._fn.apply(*args, **kwargs) 

65 

66 def extra_repr(self) -> str: 

67 """Show the wrapped Function class name in the module repr.""" 

68 return f"fn={self._fn.__name__}"