Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / overlap / overlap_param_gather.py: 89%
202 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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"""Prefetch FSDP parameters in a joint forward/backward FX graph.
17Adapted from https://gitcode.com/Ascend/pytorch/pull/38283/diffs:
18group gathers by their first real consumer, launch ahead of compute, and sink
19waits through parameter views. HyperParallel identifies gathers by FSDPPass
20metadata rather than by placeholder inputs (which would also match TP).
21Forward and backward are scheduled separately; storage-release boundaries
22remain explicit constraints. Like manual overlap scheduling, collective
23reordering is part of this pass and requires consistent participating graphs.
24"""
26__all__ = ["overlap_param_gather"]
28import logging
29from dataclasses import dataclass
31import torch
32from torch import fx
33from torch.ops import _c10d_functional
35from ...pass_config import PassConfig
36from .utils import bucket_preparation, is_barrier, is_op, is_view
38_LOG = logging.getLogger(__name__)
39_GATHER_TYPES = ("fsdp_all_gather", "fsdp_all_gather_backward")
40_HEAVY_OPS = {"aten.mm", "aten.addmm", "aten.bmm", "aten.matmul", "aten.linear"}
42# Mutable module policies keep graph consumer lookahead separate from eager FSDP configuration.
43fsdp_forward_prefetch_depth: int = 1
44fsdp_backward_prefetch_depth: int = 1
47@dataclass
48class _Gather:
49 """One tagged FSDP gather and its complete wait/view realization DAG."""
51 launch: fx.Node
52 realization: list[fx.Node]
53 first_user: fx.Node
56def _indices(graph: fx.Graph) -> dict[fx.Node, int]:
57 return {node: index for index, node in enumerate(graph.nodes)}
60def _is_attention(node: fx.Node) -> bool:
61 return node.op == "call_function" and any(
62 name in str(node.target)
63 for name in ("npu_fusion_attention", "scaled_dot_product", "flash_attention")
64 )
67def _is_heavy(node: fx.Node) -> bool:
68 name = str(node.target)
69 return node.op == "call_function" and (
70 name.rsplit(".", 1)[0] in _HEAVY_OPS or _is_attention(node)
71 )
74def _collect_gathers(graph: fx.Graph) -> list[_Gather]:
75 """Find real consumers through views and the bucket's explicit copy-out nodes."""
76 order = _indices(graph)
77 matches = []
78 for launch in graph.nodes:
79 if launch.meta.get("comm_type") not in _GATHER_TYPES:
80 continue
81 if not is_op(launch, _c10d_functional.all_gather_into_tensor):
82 continue
83 waits = list(launch.users)
84 if len(waits) != 1 or not is_op(waits[0], _c10d_functional.wait_tensor):
85 continue
86 realization = {waits[0]}
87 stack = [waits[0]]
88 consumers = set()
89 while stack:
90 current = stack.pop()
91 for user in current.users:
92 copy_out = (launch.meta.get("fsdp_bucket_id") is not None
93 and user.meta.get("fsdp_bucket_copy_out") == launch.meta["fsdp_bucket_id"]
94 and user.target is torch.ops.aten.clone.default)
95 if is_view(user) or copy_out:
96 if user not in realization:
97 realization.add(user)
98 stack.append(user)
99 else:
100 consumers.add(user)
101 if consumers:
102 matches.append(_Gather(launch, sorted(realization, key=order.get), min(consumers, key=order.get)))
103 return matches
106def _launch_chain(launch: fx.Node) -> list[fx.Node]:
107 """Move the private bucket pack DAG or an original private shard cast."""
108 if launch.meta.get("fsdp_bucket_id"):
109 return [*bucket_preparation(launch, _indices(launch.graph)), launch]
110 shard = launch.args[0]
111 if isinstance(shard, fx.Node) and len(shard.users) == 1 and shard.op == "call_function":
112 schema = getattr(shard.target, "_schema", None)
113 if schema is not None and schema.name in ("prims::convert_element_type", "aten::_to_copy"):
114 if shard.all_input_nodes and all(node.op == "placeholder" for node in shard.all_input_nodes):
115 return [shard, launch]
116 return [launch]
119def _gather_nbytes(launch: fx.Node) -> int | None:
120 nbytes = launch.meta.get("fsdp_unsharded_nbytes")
121 return nbytes if isinstance(nbytes, int) and not isinstance(nbytes, bool) and nbytes >= 0 else None
124def _within_budget(
125 match: _Gather,
126 target: fx.Node,
127 prefetched: list[_Gather],
128 order: dict[fx.Node, int],
129 budget: int,
130) -> bool:
131 """Bound full replicas launched early but not yet consumed, across anchors."""
132 nbytes = _gather_nbytes(match.launch)
133 if nbytes is None:
134 return False
135 events = [(2 * order[target] - 1, nbytes), (2 * order[match.first_user], -nbytes)]
136 for previous in prefetched:
137 if previous is match:
138 continue
139 previous_bytes = _gather_nbytes(previous.launch)
140 if previous_bytes is None:
141 return False
142 events.extend(((2 * order[previous.launch], previous_bytes),
143 (2 * order[previous.first_user], -previous_bytes)))
144 live_bytes = 0
145 for _, delta in sorted(events):
146 live_bytes += delta
147 if live_bytes > budget:
148 return False
149 return True
152def _add_attention_interval_anchors(
153 groups: dict[fx.Node, list[_Gather]],
154 nodes: list[fx.Node],
155 order: dict[fx.Node, int],
156 anchors: dict[fx.Node, fx.Node],
157) -> None:
158 """Merge earlier attention-interval candidates into the lookahead anchors."""
159 users = list(groups)
160 attention = [node for node in nodes if _is_attention(node)]
161 boundaries = [-1, *[order[node] for node in attention], len(order)]
162 for index in range(len(boundaries) - 1):
163 left, right = boundaries[index:index + 2]
164 heavy = [node for node in nodes if _is_heavy(node) and max(left, 0) <= order[node] <= right]
165 if not attention or not heavy:
166 continue
167 end = boundaries[index + 2] if index + 2 < len(boundaries) else right
168 targets = [user for user in users if right < order[user] <= end]
169 if index == 0:
170 targets = [user for user in users if left < order[user] <= right] + targets
171 for target_index, user in enumerate(targets):
172 anchor = heavy[target_index * len(heavy) // len(targets)]
173 for match in groups[user]:
174 existing = anchors.get(match.launch, user)
175 if order[anchor] < order[existing]:
176 anchors[match.launch] = anchor
179def _prefetch_anchors(
180 matches: list[_Gather], nodes: list[fx.Node], depth: int, order: dict[fx.Node, int]
181) -> dict[fx.Node, fx.Node]:
182 """Combine consumer lookahead with the PR's attention-interval anchors."""
183 groups: dict[fx.Node, list[_Gather]] = {}
184 for match in sorted(matches, key=lambda item: order[item.first_user]):
185 groups.setdefault(match.first_user, []).append(match)
186 users = list(groups)
187 anchors = {}
188 for index, user in enumerate(users):
189 if index > 0:
190 for match in groups[user]:
191 anchors[match.launch] = users[max(0, index - depth)]
193 _add_attention_interval_anchors(groups, nodes, order, anchors)
194 return anchors
197def _launch_constraints(graph: fx.Graph) -> dict[fx.Node, list[fx.Node]]:
198 """Preserve effects and parameter lifetimes during collective reordering."""
199 constraints = {}
200 previous_effect = None
201 frees = {}
202 for node in graph.nodes:
203 if node.meta.get("comm_type") in _GATHER_TYPES:
204 names = node.meta.get("fsdp_bucket_members", (node.meta.get("param_name"),))
205 predecessors = [previous_effect, *(frees.get(name) for name in names)]
206 constraints[node] = [item for item in predecessors if item is not None]
207 if is_barrier(node):
208 previous_effect = node
209 if node.meta.get("comm_type") in ("fsdp_reshard_free", "fsdp_reshard_free_backward"):
210 frees[node.meta.get("param_name")] = node
211 return constraints
214def _prefetch_phase(
215 graph: fx.Graph, matches: list[_Gather], depth: int, budget: int, constraints: dict[fx.Node, list[fx.Node]]
216) -> int:
217 """Prefetch one phase within input, effect, collective, and byte constraints."""
218 if not matches or depth == 0 or budget == 0:
219 return 0
220 backward = bool(matches[0].launch.meta.get("autograd_backward", False))
221 nodes = [node for node in graph.nodes if node.op not in ("placeholder", "get_attr")
222 and bool(node.meta.get("autograd_backward", False)) == backward]
223 order = _indices(graph)
224 phase_start = nodes[0]
225 anchors = _prefetch_anchors(matches, nodes, depth, order)
226 prefetched = [match for match in matches if "fsdp_prefetch_anchor" in match.launch.meta]
227 changed = 0
228 for match in matches:
229 anchor = anchors.get(match.launch)
230 if anchor is None:
231 continue
232 order = _indices(graph)
233 chain = _launch_chain(match.launch)
234 predecessors = list(constraints[match.launch])
235 predecessors.extend(arg for node in chain for arg in node.all_input_nodes if arg not in chain)
236 position = max(order[anchor], order[phase_start], *(order[node] + 1 for node in predecessors))
237 if position >= order[chain[0]]:
238 continue
239 target = list(graph.nodes)[position]
240 if not _within_budget(match, target, prefetched, order, budget):
241 continue
242 for node in chain:
243 target.prepend(node)
244 match.launch.meta["fsdp_prefetch_anchor"] = anchor.name
245 if match not in prefetched:
246 prefetched.append(match)
247 changed += 1
248 return changed
251def _sink_realization(graph: fx.Graph, matches: list[_Gather]) -> int:
252 """Sink every branch in reverse topological order, keeping all users valid."""
253 changed = 0
254 for match in matches:
255 next_realization = None
256 for node in reversed(match.realization):
257 if not node.users:
258 continue
259 nodes = list(graph.nodes)
260 order = {item: index for index, item in enumerate(nodes)}
261 # Preserve independent branch order, including copy-outs sharing a
262 # consumer; otherwise reverse traversal flips them on every run.
263 successors = list(node.users)
264 if next_realization is not None:
265 successors.append(next_realization)
266 target = min(successors, key=order.get)
267 # Other parameter frees may be crossed, but this storage's free
268 # is already a user of its realization and bounds the sinking.
269 for between in nodes[order[node] + 1:order[target]]:
270 if is_barrier(between) or (
271 bool(between.meta.get("autograd_backward", False))
272 != bool(match.launch.meta.get("autograd_backward", False))
273 ):
274 target = between
275 break
276 if order[target] > order[node] + 1:
277 target.prepend(node)
278 changed += 1
279 next_realization = node
280 return changed
283def overlap_param_gather(graph_module: fx.GraphModule, pass_config: PassConfig) -> fx.GraphModule:
284 """Schedule tagged FSDP gathers without changing tensor or collective semantics.
286 Args:
287 graph_module: Joint graph after FSDPPass (including reshard/rematerialization).
288 pass_config: Independent prefetch switch and maximum early replicas in GiB.
290 Returns:
291 The transformed graph. Unknown replica sizes are conservatively not prefetched.
292 The budget covers early launches until first use, not total model/activation memory.
294 Note:
295 Collective reordering requires identical static graphs and
296 configuration across participating ranks. Only FSDP gathers move; other
297 launches retain their relative order. Input, effect and free boundaries
298 still apply. Module variables fsdp_forward_prefetch_depth and
299 fsdp_backward_prefetch_depth set consumer-group lookahead (default 1).
300 Patch them consistently on all ranks before compilation; existing
301 compiled graphs are unaffected. Zero disables early gathers in that
302 phase, including attention-interval prefetch, but still sinks waits/views.
303 """
304 if not pass_config.fsdp_allgather_prefetch:
305 return graph_module
306 depths = (fsdp_forward_prefetch_depth, fsdp_backward_prefetch_depth)
307 for phase, depth in zip(("forward", "backward"), depths):
308 if not isinstance(depth, int) or isinstance(depth, bool) or depth < 0:
309 raise ValueError(f"fsdp_{phase}_prefetch_depth must be a non-negative integer")
310 graph = graph_module.graph
311 graph.lint()
312 matches = _collect_gathers(graph)
313 constraints = _launch_constraints(graph)
314 budget = int(pass_config.max_in_flight_gb * 1024 * 1024 * 1024)
315 counts = []
316 for backward, depth in zip((False, True), depths):
317 phase_matches = [match for match in matches
318 if bool(match.launch.meta.get("autograd_backward", False)) == backward]
319 counts.append(_prefetch_phase(graph, phase_matches, depth, budget, constraints))
320 sunk = _sink_realization(graph, matches)
321 graph.lint()
322 graph_module.recompile()
323 graph_module.meta["fsdp_prefetch"] = {"forward": counts[0], "backward": counts[1], "sunk": sunk}
324 _LOG.info("FSDP parameter prefetch: forward=%s backward=%s sunk=%s budget_bytes=%s depths=%s",
325 counts[0], counts[1], sunk, budget, depths)
326 return graph_module