Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / _sync_hook.py: 87%

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"""Autograd synchronization hooks for pipeline overlap.""" 

16from typing import Any 

17 

18import torch 

19 

20from hyper_parallel.core.pipeline_parallel.hook_coordinator import HookCoordinator, HookRole 

21 

22 

23class _SyncHookFunction(torch.autograd.Function): 

24 """Autograd identity that fires HookCoordinator rendezvous on fwd/bwd. 

25 

26 Uses a **4-hook** design (``A``, ``B``, ``C``, ``D``) with pure 

27 COMM / COMPUTE roles — no NONE role. Every rendezvous is a strict 

28 COMM + COMPUTE pair, guaranteeing NCCL-first dispatch ordering at 

29 **all** points including layer boundaries. 

30 

31 Hook placement per MoE layer:: 

32 

33 [A] → dispatch → [B] → module → [C] → combine → [D] → (Attention) → [A_next] 

34 

35 At layer boundaries (D / A hooks), the Attention that runs between 

36 layers is treated as COMPUTE, and the combine / combine.bwd is treated 

37 as COMM, so the coordinator enforces comm-first ordering even across 

38 layer transitions. 

39 """ 

40 

41 # 4-hook role tables: (prev_role_idx, next_role_idx). 

42 # Index encoding: 1 = COMM, 2 = COMPUTE. 

43 # 

44 _FWD_ROLES = { 

45 # (prev, next) prev op next op 

46 "A": (2, 1), # COMPUTE, COMM Attention | dispatch 

47 "B": (1, 2), # COMM, COMPUTE dispatch | module 

48 "C": (2, 1), # COMPUTE, COMM module | combine 

49 "D": (1, 2), # COMM, COMPUTE combine | Attention 

50 } 

51 _BWD_ROLES = { 

52 "D": (2, 1), # COMPUTE, COMM Attn.bwd | combine.bwd 

53 "C": (1, 2), # COMM, COMPUTE combine.bwd | module.bwd 

54 "B": (2, 1), # COMPUTE, COMM module.bwd | dispatch.bwd 

55 "A": (1, 2), # COMM, COMPUTE dispatch.bwd| Attn.bwd 

56 } 

57 

58 _ROLE_CACHE = (None, HookRole.COMM, HookRole.COMPUTE) 

59 

60 @staticmethod 

61 def _role_enum(idx: int): 

62 return _SyncHookFunction._ROLE_CACHE[idx] 

63 

64 @staticmethod 

65 def forward(ctx: Any, x: torch.Tensor, hook_name: str, coordinator: HookCoordinator) -> torch.Tensor: # pylint: disable=arguments-differ 

66 """Identity forward that fires a HookCoordinator rendezvous. 

67 

68 Notifies the previous op's role and rendezvouses for the next op's 

69 role per the ``_FWD_ROLES`` table. ``"D_LAST"`` is a sentinel 

70 meaning "skip this rendezvous" (last layer's closing D — no 

71 Attention follows). 

72 

73 Args: 

74 ctx: Autograd context, stores ``hook_name`` and 

75 ``coordinator`` for the backward pass. 

76 x: Input tensor, returned unchanged. 

77 hook_name: One of ``"A"``, ``"B"``, ``"C"``, ``"D"``, 

78 ``"D_LAST"``. 

79 coordinator: The :class:`HookCoordinator` driving the rendezvous. 

80 

81 Returns: 

82 ``x`` unchanged. 

83 """ 

84 ctx.hook_name = hook_name 

85 ctx.coordinator = coordinator 

86 

87 if not coordinator.is_enabled(): 

88 return x 

89 

90 if hook_name == "D_LAST": 

91 # ``D_LAST`` marks the last layer's closing D hook — no 

92 # Attention follows in this chunk, so the rendezvous is 

93 # meaningless and is skipped. We still 

94 # ``notify_dispatched(COMM)`` so the COMPUTE side of the 

95 # preceding ``C`` rendezvous unblocks early, letting 

96 # BWD's Attn.bwd_last overlap with FWD's post-combine 

97 # work — Torch autograd is thread-safe so this concurrent 

98 # FWD-record + BWD-replay is fine. 

99 prev_idx, _ = _SyncHookFunction._FWD_ROLES["D"] 

100 role_of = _SyncHookFunction._role_enum 

101 coordinator.notify_dispatched(role_of(prev_idx)) 

102 return x 

103 

104 prev_idx, next_idx = _SyncHookFunction._FWD_ROLES[hook_name] 

105 role_of = _SyncHookFunction._role_enum 

106 coordinator.notify_dispatched(role_of(prev_idx)) 

107 coordinator.rendezvous(role_of(next_idx)) 

108 return x 

109 

110 @staticmethod 

111 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: 

112 """Identity backward that fires a HookCoordinator rendezvous. 

113 

114 Mirror of :meth:`forward` using the ``_BWD_ROLES`` table. 

115 ``"D_LAST"`` skips the rendezvous because this is the first BWD 

116 hook to fire and ``combine.bwd`` has already dispatched freely 

117 before any rendezvous can happen. 

118 

119 Args: 

120 ctx: Autograd context with ``hook_name`` and 

121 ``coordinator`` saved during forward. 

122 grad_output: Gradient w.r.t. the forward output, returned 

123 unchanged. 

124 

125 Returns: 

126 ``(grad_output, None, None)`` — gradients only flow back to 

127 the tensor input, ``hook_name`` and ``coordinator`` are 

128 non-tensor inputs. 

129 """ 

130 hook_name = ctx.hook_name 

131 coordinator = ctx.coordinator 

132 

133 if not coordinator.is_enabled(): 

134 return grad_output, None, None 

135 

136 if hook_name == "D_LAST": 

137 # First BWD hook to fire; combine.bwd has already 

138 # dispatched freely before any rendezvous can happen. 

139 # Skipping here is safe on Torch because CUDA streams 

140 # are process-wide and the NCCL FIFO order is consistent 

141 # across ranks regardless of which thread launched 

142 # combine.bwd. 

143 return grad_output, None, None 

144 

145 prev_idx, next_idx = _SyncHookFunction._BWD_ROLES[hook_name] 

146 role_of = _SyncHookFunction._role_enum 

147 coordinator.notify_dispatched(role_of(prev_idx)) 

148 coordinator.rendezvous(role_of(next_idx)) 

149 return grad_output, None, None