Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / bucketing / bucketing_pass.py: 85%
162 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"""Shared FSDP bucket graph rewriting with kind-specific packing and ownership."""
17import logging
18import math
19import operator
20from functools import partial
21from typing import Any
23import torch
24import torch.distributed as dist
25from torch import fx
26from torch.ops import _c10d_functional as c10d
27from torch.distributed.distributed_c10d import _resolve_process_group
29from ...pass_config import PassConfig
30from ..base import GraphPass
31from ..overlap.utils import is_view, rs_layout_preparation
32from .planner import (
33 Bucket, BucketPlan, Candidate, build_module_bucket_plans, graph_signature, plan_fsdp_buckets,
34 validate_group_plans,
35)
37_LOG = logging.getLogger(__name__)
38_GATHERS = ("fsdp_all_gather", "fsdp_all_gather_backward")
41def _create_node(graph: fx.Graph, bucket: Bucket, first: Candidate, target: object, args: tuple,
42 *, preparation: bool = False, **kwargs: object) -> fx.Node:
43 """Create a node carrying phase, scope and optional private preparation ownership."""
44 node = graph.call_function(target, args, kwargs)
45 node.meta.update(autograd_backward=first.backward, fsdp_bucket_id=bucket.bucket_id,
46 fsdp_bucket_module_fqns=first.module_fqns)
47 if preparation:
48 node.meta["fsdp_bucket_prepare"] = bucket.bucket_id
49 return node
52def _prepare_rs_layouts(members: tuple, bucket: Bucket) -> None:
53 """Expose private RS preparation nodes as pure contiguous calls."""
54 for member in members:
55 for source in rs_layout_preparation(member.launch):
56 # Preserve upstream contiguous and expose its purity to scheduling.
57 source.op = "call_function"
58 source.target = torch.ops.aten.contiguous.default
59 source.meta.update(autograd_backward=True, fsdp_bucket_prepare=bucket.bucket_id,
60 fsdp_bucket_id=bucket.bucket_id, fsdp_reduce_scatter_prepare=member.launch.name)
63def _bucket_anchor(order: dict, members: tuple) -> tuple:
64 """Return (insertion anchor, first launch) placing the bucket after all inputs."""
65 first_launch = min((member.launch for member in members), key=order.get)
66 latest_input = max((member.launch.args[0] for member in members), key=order.get)
67 anchor = first_launch if order[latest_input] < order[first_launch] else latest_input.next
68 return anchor, first_launch
71def _delayed_result_views(members: tuple, order: dict, anchor: fx.Node) -> list:
72 """Collect immutable member views scheduled before the insertion anchor."""
73 pending = [member.wait for member in members]
74 views = set()
75 while pending:
76 for user in pending.pop().users:
77 if is_view(user) and user not in views:
78 views.add(user)
79 pending.append(user)
80 return sorted((node for node in views if order[node] < order[anchor]), key=order.get)
83def _launch_packed_collective(create, packed, first: Candidate, members: tuple, degree: int,
84 gather: bool, numels: list) -> tuple:
85 """Create the fused collective launch and its wait node at the current insertion point."""
86 if gather:
87 launch = create(c10d.all_gather_into_tensor.default, (packed, degree, first.group))
88 replica_bytes = sum(member.size_bytes for member in members)
89 launch.meta.update(fsdp_unsharded_nbytes=replica_bytes, fsdp_bucket_workspace_nbytes=replica_bytes)
90 else:
91 packed = create(torch.ops.aten.reshape.default, (packed, (-1,)), preparation=True)
92 launch = create(c10d.reduce_scatter_tensor.default, (packed, first.reduce_op, degree, first.group))
93 launch.meta.update(comm_type=first.kind, comm_group=first.group, fsdp_degree=degree,
94 fsdp_bucket_members=tuple(member.param_name for member in members),
95 fsdp_bucket_spec={"input_shape": (sum(numels) * (1 if gather else degree),),
96 "output_shape": (sum(numels) * (degree if gather else 1),),
97 "dtype": first.dtype, "device": first.device})
98 wait = create(c10d.wait_tensor.default, (launch,))
99 wait.meta["wait_for"] = launch.name
100 return launch, wait
103def _unpack_member_outputs(create, wait: fx.Node, numels: list, members: tuple,
104 gather: bool, degree: int, bucket: Bucket) -> list:
105 """Split the waited packed result back into per-member reshape outputs."""
106 unpack = create(torch.ops.aten.reshape.default, (wait, (degree, -1))) if gather else wait
107 parts = create(torch.ops.aten.split_with_sizes.default, (unpack, numels, 1 if gather else 0))
108 outputs = []
109 for index, member in enumerate(members):
110 part = create(operator.getitem, (parts, index))
111 if gather:
112 # reshape/contiguous may alias; each original storage free needs
113 # a copy independent of both the workspace and other members.
114 part = create(torch.ops.aten.clone.default, (part,), memory_format=torch.contiguous_format)
115 part.meta["fsdp_bucket_copy_out"] = bucket.bucket_id
116 output = create(torch.ops.aten.reshape.default, (part, member.output_shape))
117 output.meta.update(param_name=member.param_name, state_idx=member.state_idx)
118 outputs.append(output)
119 return outputs
122def _erase_members(graph: fx.Graph, members: tuple, launch: fx.Node, gather: bool) -> None:
123 """Remove original wait/launch pairs and retarget gather back-pointers."""
124 for member in members:
125 if gather:
126 for node in graph.nodes:
127 if node.meta.get("fsdp_ag_node") == member.launch.name:
128 node.meta["fsdp_ag_node"] = launch.name
129 graph.erase_node(member.wait)
130 graph.erase_node(member.launch)
133def _merge_bucket(graph: fx.Graph, bucket: Bucket) -> fx.Node:
134 """Place one packed collective after its inputs and repair immutable result views."""
135 members = bucket.members
136 first = members[0]
137 gather = bucket.supports_ag_rewrite
138 degree = len(first.ranks)
139 if not gather:
140 _prepare_rs_layouts(members, bucket)
141 order = {node: index for index, node in enumerate(graph.nodes)}
142 anchor, first_launch = _bucket_anchor(order, members)
143 delayed_views = _delayed_result_views(members, order, anchor)
144 numels = [math.prod(member.input_shape if gather else member.output_shape) for member in members]
145 _LOG.debug("FSDP bucket preparing: bucket=%s placement=%s anchor=%s members=%s",
146 bucket.bucket_id, "before_first_launch" if anchor is first_launch else "after_latest_input",
147 anchor.name, [member.param_name for member in members])
148 create = partial(_create_node, graph, bucket, first)
149 with graph.inserting_before(anchor):
150 rows = [create(torch.ops.aten.reshape.default,
151 (member.launch.args[0], (-1,) if gather else (degree, -1)), preparation=True)
152 for member in members]
153 packed = create(torch.ops.aten.cat.default, (rows, 0 if gather else 1), preparation=True)
154 launch, wait = _launch_packed_collective(create, packed, first, members, degree, gather, numels)
155 outputs = _unpack_member_outputs(create, wait, numels, members, gather, degree, bucket)
156 for member, output in zip(members, outputs):
157 member.wait.replace_all_uses_with(output)
158 cursor = outputs[-1]
159 for view in delayed_views:
160 cursor.append(view)
161 cursor = view
162 _erase_members(graph, members, launch, gather)
163 return launch
166def _validate_plan_snapshot(graph_module: fx.GraphModule, plan: BucketPlan) -> None:
167 """Reject plans built against a different or outdated graph snapshot."""
168 if any(member.launch.graph is not graph_module.graph for bucket in plan.buckets for member in bucket.members):
169 raise ValueError("Bucket plan belongs to a different graph")
170 if graph_signature(graph_module.graph) != plan.signature:
171 raise ValueError("Stale bucket plan: re-plan after changing the graph or its metadata")
174def _rewrite_stats(graph_module: fx.GraphModule, buckets: list, kinds: tuple, before: int,
175 gather: bool) -> dict:
176 """Summarize one completed kind-specific rewrite over the updated graph."""
177 kind = "ag" if gather else "rs"
178 member_bytes = sum(member.size_bytes for bucket in buckets for member in bucket.members)
179 collectives_after = sum(node.meta.get("comm_type") in kinds for node in graph_module.graph.nodes)
180 bytes_key = "copy_out_bytes" if gather else "packed_input_bytes"
181 return {"buckets": len(buckets),
182 "original_collectives": sum(len(bucket.members) for bucket in buckets),
183 bytes_key: member_bytes,
184 f"{kind}_collectives_before": before,
185 f"{kind}_collectives_after": collectives_after}
188def _apply_buckets(graph_module: fx.GraphModule, plan: BucketPlan, *, gather: bool) -> fx.GraphModule:
189 """Validate the snapshot and publish kind-specific statistics only after success."""
190 _validate_plan_snapshot(graph_module, plan)
191 if gather:
192 buckets = [bucket for bucket in plan.buckets if bucket.supports_ag_rewrite]
193 else:
194 buckets = [bucket for bucket in plan.buckets if bucket.supports_rs_rewrite]
195 kinds = _GATHERS if gather else ("fsdp_reduce_scatter",)
196 before = sum(node.meta.get("comm_type") in kinds for node in graph_module.graph.nodes)
197 launches = [_merge_bucket(graph_module.graph, bucket) for bucket in buckets]
198 if buckets:
199 graph_module.graph.lint()
200 graph_module.recompile()
201 kind = "ag" if gather else "rs"
202 stats = _rewrite_stats(graph_module, buckets, kinds, before, gather)
203 graph_module.meta[f"fsdp_{kind}_bucketing"] = stats
204 _LOG.info("FSDP bucket rewrite completed: kind=%s graph=%s stats=%s", kind, plan.signature[:12], stats)
205 if _LOG.isEnabledFor(logging.DEBUG):
206 order = {node: index for index, node in enumerate(graph_module.graph.nodes)}
207 for launch in launches:
208 _LOG.debug("FSDP bucket applied: kind=%s bucket=%s launch=%s index=%s coordinate_space=rewritten_graph",
209 kind, launch.meta["fsdp_bucket_id"], launch.name, order[launch])
210 return graph_module
213def apply_ag_buckets(graph_module: fx.GraphModule, plan: BucketPlan) -> fx.GraphModule:
214 """Fuse non-out AG with independent copies, preserving original consumers and frees.
216 Args:
217 graph_module: Exact graph inspected by the planner, before overlap.
218 plan: Fresh rank-consistent plan; re-plan after RS or any other graph change.
220 Returns:
221 Updated graph with each padded parameter owning independent storage.
222 Singleton/RS buckets stay unchanged. The workspace uses normal backend
223 stream/reference lifetime; prefetch bytes do not bound peak memory.
225 Raises:
226 ValueError: The plan belongs to another graph or is stale.
228 Note:
229 Caller owns cross-rank validation. Only immutable result views may move.
230 Stats describe this invocation; failed rewrites do not automatically roll back.
231 """
232 return _apply_buckets(graph_module, plan, gather=True)
235def apply_rs_buckets(graph_module: fx.GraphModule, plan: BucketPlan) -> fx.GraphModule:
236 """Fuse RS with rank-major packing and shared output views, preserving output slots.
238 Args:
239 graph_module: Exact graph inspected by the planner, before overlap.
240 plan: Fresh rank-consistent plan; re-plan after AG or any other graph change.
242 Returns:
243 Updated graph. Singleton/AG buckets stay unchanged; stats describe this call.
244 Run overlap afterwards to schedule private preparation and launch together.
246 Raises:
247 ValueError: The plan belongs to another graph or is stale.
249 Note:
250 Caller owns cross-rank validation. Failed rewrites do not automatically roll back.
251 """
252 return _apply_buckets(graph_module, plan, gather=False)
255class FsdpBucketingPass(GraphPass):
256 """Plan and fuse FSDP communication after partitioning and before overlap."""
258 name = "fsdp_bucketing"
260 def run(self, graph_module: fx.GraphModule, pass_config: PassConfig, **kwargs: Any) -> fx.GraphModule:
261 """Apply Python-configured module buckets with compile-time rank validation.
263 Args:
264 graph_module: Joint graph annotated by FSDPPass.
265 pass_config: Bucketing switch and optional module scopes; identical across ranks.
266 **kwargs: Live model and registered dense/expert FSDP group names from the compiler.
268 Returns:
269 Graph with fused AG/RS, or unchanged graph when disabled/unsupported.
270 """
271 if not pass_config.fsdp_bucketing or not pass_config.fsdp_enabled or not dist.is_initialized():
272 return graph_module
273 if dist.get_world_size() == 1:
274 return graph_module
275 if pass_config.dp_mode != "fsdp":
276 _LOG.info("Skipped FSDP bucketing: requires dp_mode=fsdp")
277 return graph_module
278 modules = pass_config.fsdp_bucket_modules
279 if modules is not None:
280 modules = build_module_bucket_plans(kwargs["model"], module_bucket_plans=modules)
281 names = [kwargs.get("fsdp_group_name", "fsdp")]
282 if kwargs.get("expert_param_names"):
283 names.append(kwargs.get("expert_fsdp_group_name", "fsdp_expert"))
284 groups = {name: _resolve_process_group(name) for name in names}
285 ranks = {name: tuple(dist.get_process_group_ranks(group)) for name, group in groups.items()}
286 for apply in (apply_ag_buckets, apply_rs_buckets):
287 plan = plan_fsdp_buckets(graph_module, group_ranks=ranks, module_bucket_plans=modules)
288 for name, group in groups.items():
289 report = {"group_ranks": {name: ranks[name]}, "group_digests": {name: plan.group_digest(name)}}
290 reports = [None] * len(ranks[name])
291 dist.all_gather_object(reports, report, group=group)
292 validate_group_plans(dict(zip(ranks[name], reports)), name)
293 apply(graph_module, plan)
294 return graph_module