Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / overlap / overlap_grad_comm.py: 96%

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"""Overlap FSDP gradient reduce-scatter with the remaining backward compute. 

16 

17Adapted from https://gitcode.com/Ascend/pytorch/pull/37311: defer gradient 

18waits until consumption. Parameter all-gather lookahead is already provided 

19by overlap_param_gather. FSDPPass emits reductions at the joint graph's tail, 

20so this pass also launches each reduction as soon as its gradient is ready. 

21FSDP metadata handles decomposed matmuls, norm and accumulated/tied gradients 

22without guessing which output of an operator is a weight gradient. 

23""" 

24 

25__all__ = ["overlap_grad_comm"] 

26 

27import logging 

28from dataclasses import dataclass 

29 

30from torch import fx 

31from torch.ops import _c10d_functional, aten 

32 

33from ...pass_config import PassConfig 

34from .utils import bucket_preparation, is_barrier, is_collective, is_op, is_view, rs_layout_preparation 

35 

36_LOG = logging.getLogger(__name__) 

37_PARAM_GATHERS = ("fsdp_all_gather", "fsdp_all_gather_backward") 

38_STORAGE_RELEASES = ("fsdp_reshard_free", "fsdp_reshard_free_backward") 

39 

40 

41@dataclass 

42class _Reduction: 

43 """A tagged launch and the full alias DAG materializing its waited result.""" 

44 

45 launch: fx.Node 

46 realization: list[fx.Node] 

47 input_aliases: set[fx.Node] 

48 preparation: list[fx.Node] 

49 

50 

51def _is_contiguous_method(node: fx.Node) -> bool: 

52 return node.op == "call_method" and node.target == "contiguous" 

53 

54 

55def _alias_ancestors(inputs: list[fx.Node]) -> set[fx.Node]: 

56 """Track input storage through immutable aliases, without traversing compute.""" 

57 aliases = set(inputs) 

58 stack = list(inputs) 

59 while stack: 

60 node = stack.pop() 

61 # Even an unmovable/shared Tensor.contiguous may return its input. 

62 if is_view(node) or _is_contiguous_method(node): 

63 for arg in node.all_input_nodes: 

64 if arg not in aliases: 

65 aliases.add(arg) 

66 stack.append(arg) 

67 return aliases 

68 

69 

70def _is_input_release(node: fx.Node, match: _Reduction) -> bool: 

71 return node.meta.get("comm_type") in _STORAGE_RELEASES and bool( 

72 _alias_ancestors(node.all_input_nodes) & match.input_aliases 

73 ) 

74 

75 

76def _collect_reductions(graph: fx.Graph) -> tuple[list[_Reduction], int]: 

77 """Collect reductions and normalize their private method preparations.""" 

78 order = {node: index for index, node in enumerate(graph.nodes)} 

79 matches = [] 

80 normalized = 0 

81 for launch in graph.nodes: 

82 if launch.meta.get("comm_type") != "fsdp_reduce_scatter": 

83 continue 

84 if not launch.meta.get("autograd_backward", False): 

85 continue 

86 if not is_op(launch, _c10d_functional.reduce_scatter_tensor): 

87 continue 

88 users = list(launch.users) 

89 if len(users) != 1 or not is_op(users[0], _c10d_functional.wait_tensor): 

90 continue 

91 realization = {users[0]} 

92 stack = [users[0]] 

93 while stack: 

94 node = stack.pop() 

95 for user in node.users: 

96 if is_view(user) and user not in realization: 

97 realization.add(user) 

98 stack.append(user) 

99 preparation = bucket_preparation(launch, order) or rs_layout_preparation(launch) 

100 for source in preparation: 

101 if source.op == "call_method": 

102 # Only a matched backward RS owns this default Tensor method. 

103 # Its ATen schema exposes both purity and possible input aliasing. 

104 source.op = "call_function" 

105 source.target = aten.contiguous.default 

106 source.meta.update(autograd_backward=True, fsdp_reduce_scatter_prepare=launch.name) 

107 normalized += 1 

108 external = [arg for node in [*preparation, launch] for arg in node.all_input_nodes 

109 if arg not in preparation] 

110 matches.append(_Reduction(launch, sorted(realization, key=order.get), 

111 _alias_ancestors([*external, *launch.all_input_nodes]), preparation)) 

112 return matches, normalized 

113 

114 

115def _group(node: fx.Node) -> object: 

116 return node.meta.get("comm_group", node.kwargs.get("group_name", node.args[-1] if node.args else None)) 

117 

118 

119def _launch_boundaries(nodes: list[fx.Node], matches: list[_Reduction]) -> dict[fx.Node, int]: 

120 """Find the last effect, phase boundary or unowned same-group collective. 

121 

122 Owned FSDP gathers and reductions may interleave in a new deterministic 

123 order. All other launches retain their relative order, including TP. 

124 A same-group collective not owned by this scheduler remains a barrier. 

125 """ 

126 launches = {match.launch: match for match in matches} 

127 groups = {_group(node) for node in launches} 

128 last_effect = -1 

129 last_unowned = dict.fromkeys(groups, -1) 

130 boundaries = {} 

131 releases = [] 

132 for index, node in enumerate(nodes): 

133 if node in launches: 

134 related = [position for position, release in releases if _is_input_release(release, launches[node])] 

135 boundaries[node] = max(last_effect, last_unowned[_group(node)], *related) 

136 elif node.op not in ("placeholder", "get_attr", "output"): 

137 if node.meta.get("comm_type") in _STORAGE_RELEASES: 

138 releases.append((index, node)) 

139 if is_barrier(node) or not node.meta.get("autograd_backward", False): 

140 last_effect = index 

141 if is_collective(node) and node.meta.get("comm_type") not in _PARAM_GATHERS: 

142 group = _group(node) 

143 if group in last_unowned: 

144 last_unowned[group] = index 

145 return boundaries 

146 

147 

148def _advance_launches(graph: fx.Graph, matches: list[_Reduction]) -> int: 

149 """Launch each reduction after its inputs and communication boundaries are ready.""" 

150 nodes = list(graph.nodes) 

151 order = {node: index for index, node in enumerate(nodes)} 

152 boundaries = _launch_boundaries(nodes, matches) 

153 moving = {node for match in matches for node in [*match.preparation, match.launch]} 

154 moves = [] 

155 for match in matches: 

156 launch = match.launch 

157 chain = [*match.preparation, launch] 

158 inputs = [arg for node in chain for arg in node.all_input_nodes if arg not in chain] 

159 ready = max(boundaries[launch], *(order[arg] for arg in inputs)) 

160 # Targets remain stationary until all launches have moved. A target 

161 # that is itself a moving reduction could invalidate dependencies. 

162 position = ready + 1 

163 while position < len(nodes) and nodes[position] in moving: 

164 position += 1 

165 if position < order[chain[0]]: 

166 moves.append((position, order[launch], chain)) 

167 for position, _, chain in sorted(moves): 

168 target = nodes[position] 

169 for node in chain: 

170 target.prepend(node) 

171 chain[-1].meta["fsdp_grad_overlap_anchor"] = target.name 

172 return len(moves) 

173 

174 

175def _sink_waits(graph: fx.Graph, matches: list[_Reduction]) -> int: 

176 """Delay reduction wait/alias DAGs while protecting effects and input storage.""" 

177 changed = 0 

178 order = {node: index for index, node in enumerate(graph.nodes)} 

179 for match in sorted(matches, key=lambda item: order[item.launch]): 

180 # Moving each branch separately preserves early consumers of a 

181 # branched view and the mandatory RS-wait-AR dependency in HSDP. 

182 next_realization = None 

183 for node in reversed(match.realization): 

184 if not node.users: 

185 continue 

186 nodes = list(graph.nodes) 

187 positions = {item: index for index, item in enumerate(nodes)} 

188 # Keep independent unpack branches stable across repeated scheduling. 

189 successors = list(node.users) 

190 if next_realization is not None: 

191 successors.append(next_realization) 

192 target = min(successors, key=positions.get) 

193 for between in nodes[positions[node] + 1:positions[target]]: 

194 if is_barrier(between) or _is_input_release(between, match) or ( 

195 between.op != "output" and not between.meta.get("autograd_backward", False) 

196 ): 

197 target = between 

198 break 

199 if positions[target] > positions[node] + 1: 

200 target.prepend(node) 

201 changed += 1 

202 next_realization = node 

203 return changed 

204 

205 

206def overlap_grad_comm(graph_module: fx.GraphModule, pass_config: PassConfig) -> fx.GraphModule: 

207 """Advance tagged gradient launches and sink waits through immutable aliases. 

208 

209 Args: 

210 graph_module: Joint graph after FSDPPass and parameter prefetch. 

211 pass_config: fsdp_overlap_grad_reduce controls this pass independently. 

212 

213 Returns: 

214 The same graph module, with scheduling statistics in fsdp_grad_overlap. 

215 

216 Note: 

217 FSDP collective order changes to static gradient readiness order. As 

218 with the original static FSDP graph, participating ranks must execute 

219 the same graph and configuration. No tensor-dependent decisions occur. 

220 TP launches, tensor arithmetic, gradient averaging and outputs are unchanged. 

221 """ 

222 if not pass_config.fsdp_overlap_grad_reduce: 

223 return graph_module 

224 graph = graph_module.graph 

225 graph.lint() 

226 before = {node: index for index, node in enumerate(graph.nodes)} if _LOG.isEnabledFor(logging.DEBUG) else {} 

227 matches, normalized = _collect_reductions(graph) 

228 advanced = _advance_launches(graph, matches) 

229 sunk = _sink_waits(graph, matches) 

230 graph.lint() 

231 if normalized or advanced or sunk: 

232 graph_module.recompile() 

233 if _LOG.isEnabledFor(logging.DEBUG): 

234 after = {node: index for index, node in enumerate(graph.nodes)} 

235 for match in matches: 

236 launch, wait = match.launch, match.realization[0] 

237 _LOG.debug("FSDP gradient overlap completed: bucket=%s launch=%s index_before=%s index_after=%s " 

238 "wait=%s wait_before=%s wait_after=%s", 

239 launch.meta.get("fsdp_bucket_id"), launch.name, before.get(launch), after[launch], 

240 wait.name, before.get(wait), after[wait]) 

241 graph_module.meta["fsdp_grad_overlap"] = { 

242 "matched": len(matches), "launches_advanced": advanced, "wait_view_nodes_sunk": sunk, 

243 } 

244 _LOG.info("FSDP gradient overlap: matched=%s launches_advanced=%s wait_view_nodes_sunk=%s", 

245 len(matches), advanced, sunk) 

246 return graph_module