Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / bucketing / planner.py: 97%
371 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"""Deterministic, read-only planning for explicitly tagged FSDP collectives."""
17import copy
18import hashlib
19import json
20import logging
21import math
22from collections import Counter
23from dataclasses import dataclass, field, replace
24from typing import Mapping, Optional, Sequence, Union
26import torch
27from torch import fx, nn
28from torch.ops import _c10d_functional as c10d
30from ..overlap.utils import is_barrier, is_collective, is_op, is_view, rs_layout_preparation
31from ..parallel.fsdp_pass import _free_tensor_storage
33_LOG = logging.getLogger(__name__)
35_KINDS = {"fsdp_all_gather": c10d.all_gather_into_tensor,
36 "fsdp_all_gather_backward": c10d.all_gather_into_tensor,
37 "fsdp_reduce_scatter": c10d.reduce_scatter_tensor}
38_FREES = {"fsdp_reshard_free", "fsdp_reshard_free_backward"}
41ModuleBucketPlans = Sequence[Union[str, Sequence[str]]]
44def normalize_module_bucket_plans(plans: ModuleBucketPlans) -> tuple[tuple[str, ...], ...]:
45 """Freeze explicit module groups and reject ambiguous or wildcard ownership.
47 Args:
48 plans: Each entry is a module FQN or a nonempty sequence of module FQNs.
49 Prefixes include descendants on a dot boundary. Root scope is unsupported.
51 Returns:
52 Disjoint groups in caller order. An empty plan selects no modules.
54 Raises:
55 ValueError: A path is invalid, repeated, or overlaps another path.
56 """
57 if not isinstance(plans, (list, tuple)):
58 raise ValueError("module_bucket_plans must be a list or tuple of module groups")
59 groups, prefixes = [], []
60 for entry in plans:
61 group = (entry,) if isinstance(entry, str) else entry
62 if not isinstance(group, (list, tuple)) or not group:
63 raise ValueError("Each module bucket group must be a FQN or a nonempty list/tuple")
64 for fqn in group:
65 if (not isinstance(fqn, str) or not fqn or any(char in fqn for char in "*?[]")
66 or any(not part or part.strip() != part for part in fqn.split("."))):
67 raise ValueError(f"Invalid explicit module FQN: {fqn!r}")
68 if any(fqn == old or fqn.startswith(old + ".") or old.startswith(fqn + ".") for old in prefixes):
69 raise ValueError(f"Overlapping module bucket scopes: {fqn!r}")
70 prefixes.append(fqn)
71 groups.append(tuple(group))
72 return tuple(groups)
75def match_module_bucket(param_name: str, plans: tuple[tuple[str, ...], ...]) -> Optional[tuple[str, ...]]:
76 """Return the logical module group owning a parameter, or None if unselected.
78 Args:
79 param_name: Fully qualified parameter name.
80 plans: Normalized explicit module groups.
81 """
82 for group in plans:
83 if any(param_name.startswith(fqn + ".") for fqn in group):
84 return group
85 return None
88def build_module_bucket_plans(
89 model: nn.Module, *, block_module_types: tuple[type[nn.Module], ...] = (),
90 module_bucket_plans: ModuleBucketPlans = (),
91) -> tuple[tuple[str, ...], ...]:
92 """Build groups from model-owned block types and explicit additional modules.
94 Args:
95 model: Model whose named_modules paths match FSDPPass parameter FQNs.
96 block_module_types: Concrete decoder/block classes supplied by the model
97 adapter. Each outermost matching non-root module becomes one group.
98 module_bucket_plans: Additional groups, e.g. embedding and norm/head.
99 No naming heuristic or model-wide wildcard is used.
101 Returns:
102 Validated immutable groups. Their order describes ownership, not the
103 collective execution order, which is derived from the graph.
105 Raises:
106 ValueError: Types, module paths or group ownership are invalid.
107 """
108 if not isinstance(model, nn.Module):
109 raise ValueError("model must be a torch.nn.Module")
110 if (not isinstance(block_module_types, tuple)
111 or any(not isinstance(kind, type) or not issubclass(kind, nn.Module) for kind in block_module_types)):
112 raise ValueError("block_module_types must be a tuple of torch.nn.Module subclasses")
113 explicit = normalize_module_bucket_plans(module_bucket_plans)
114 modules = dict(model.named_modules())
115 for group in explicit:
116 for fqn in group:
117 if fqn not in modules:
118 raise ValueError(f"Module bucket scope does not exist in model: {fqn!r}")
119 blocks = []
120 for fqn, module in modules.items():
121 if (fqn and isinstance(module, block_module_types)
122 and not any(fqn.startswith(parent + ".") for parent in blocks)):
123 blocks.append(fqn)
124 return normalize_module_bucket_plans([*blocks, *explicit])
127def _static_int(value: object) -> bool:
128 return isinstance(value, int) and not isinstance(value, bool)
131@dataclass(frozen=True)
132class BucketPolicy:
133 """Optional physical bucket limits inside a logical module group.
135 Args:
136 cap_bytes: Maximum packed input/output bytes; None leaves capacity uncapped.
137 max_input_distance: Optional FX-node distance between input producers.
138 max_consumer_distance: Optional FX-node distance between first consumers.
139 All defaults are None, matching module-driven rather than size-driven
140 grouping. Safety barriers and communication compatibility always apply.
141 """
143 cap_bytes: Optional[int] = None
144 max_input_distance: Optional[int] = None
145 max_consumer_distance: Optional[int] = None
147 def __post_init__(self):
148 for name in ("cap_bytes", "max_input_distance", "max_consumer_distance"):
149 value = getattr(self, name)
150 if value is not None and (not _static_int(value) or value < (1 if name == "cap_bytes" else 0)):
151 raise ValueError(f"{name} must be None or a non-negative integer (cap_bytes must be positive)")
154@dataclass(frozen=True)
155class Candidate:
156 """One communication instance, with static layout and scheduling bounds."""
158 launch: fx.Node = field(repr=False)
159 wait: fx.Node = field(repr=False)
160 source: fx.Node = field(repr=False)
161 param_name: str
162 state_idx: int
163 kind: str
164 backward: bool
165 scope: str
166 module_fqns: tuple[str, ...]
167 group: str
168 ranks: tuple[int, ...]
169 dtype: torch.dtype
170 device: str
171 input_shape: tuple[int, ...]
172 output_shape: tuple[int, ...]
173 ready: int
174 first_use: int
175 reduce_op: Optional[str]
177 @property
178 def size_bytes(self) -> int:
179 """Return the larger of the input and output buffer sizes."""
180 return max(math.prod(self.input_shape), math.prod(self.output_shape)) * self.dtype.itemsize
182 @property
183 def key(self) -> tuple:
184 """Return the complete compatibility key for this local graph."""
185 return (self.kind, self.backward, self.scope, self.group, self.ranks,
186 self.dtype, self.device, self.reduce_op)
188 def describe(self) -> dict:
189 """Return rank-independent wire layout plus useful local diagnostics."""
190 return {"param_name": self.param_name, "state_idx": self.state_idx,
191 "kind": self.kind, "backward": self.backward, "scope": self.scope, "module_fqns": self.module_fqns,
192 "group": self.group, "ranks": self.ranks, "dtype": str(self.dtype),
193 "device_type": torch.device(self.device).type,
194 "input_shape": self.input_shape, "output_shape": self.output_shape,
195 "reduce_op": self.reduce_op}
198@dataclass(frozen=True)
199class Bucket:
200 """Ordered members of a proposed bucket; singleton means no fusion."""
202 members: tuple[Candidate, ...]
203 bucket_id: str = ""
205 @property
206 def supports_rs_rewrite(self) -> bool:
207 """Whether this bucket has an executable multi-member RS rewrite."""
208 return len(self.members) > 1 and self.members[0].kind == "fsdp_reduce_scatter"
210 @property
211 def supports_ag_rewrite(self) -> bool:
212 """Whether independent copy-out can materialize this multi-member AG bucket."""
213 return len(self.members) > 1 and self.members[0].kind in ("fsdp_all_gather", "fsdp_all_gather_backward")
215 def describe(self) -> dict:
216 """Describe packed offsets, bytes and scheduling bounds without mutation."""
217 offset = 0
218 members = []
219 for member in self.members:
220 local_numel = math.prod(member.output_shape if member.kind == "fsdp_reduce_scatter"
221 else member.input_shape)
222 members.append({**member.describe(), "local_offset": offset, "local_numel": local_numel})
223 offset += local_numel
224 return {"bucket_id": self.bucket_id, "rewrite_supported": self.supports_rs_rewrite or self.supports_ag_rewrite,
225 "members": members, "size_bytes": sum(m.size_bytes for m in self.members),
226 "latest_input": max(m.ready for m in self.members),
227 "earliest_consumer": min(m.first_use for m in self.members),
228 "fusable": len(self.members) > 1}
231@dataclass(frozen=True)
232class BucketPlan:
233 """Dry-run result tied to one graph snapshot; never changes that graph."""
235 buckets: tuple[Bucket, ...]
236 decisions: tuple[dict, ...]
237 signature: str
238 group_ranks: dict[str, tuple[int, ...]]
239 wire_collectives: dict[str, list]
240 module_bucket_plans: Optional[tuple[tuple[str, ...], ...]]
241 summary: dict = field(default_factory=dict)
243 def to_dict(self) -> dict:
244 """Export planning facts on demand; never reads erased nodes' arguments.
246 Returns an independent JSON-serializable dict, not an application receipt.
247 Original readiness bounds and layout remain available after rewriting.
248 """
249 return copy.deepcopy({"buckets": [bucket.describe() for bucket in self.buckets],
250 "decisions": self.decisions, "summary": self.summary,
251 "signature": self.signature, "module_bucket_plans": self.module_bucket_plans,
252 "group_ranks": self.group_ranks,
253 "group_digests": {group: self.group_digest(group) for group in self.group_ranks}})
255 def group_digest(self, group: str) -> str:
256 """Hash layout/order for one PG, excluding rank-local device indices.
258 Args:
259 group: Process group name in this plan.
260 """
261 layouts = [[member.describe() for member in bucket.members]
262 for bucket in self.buckets if bucket.members[0].group == group]
263 return _digest((self.group_ranks[group], layouts, self.wire_collectives.get(group, [])))
266def _digest(value: object) -> str:
267 return hashlib.sha256(json.dumps(value, sort_keys=True, default=str).encode()).hexdigest()
270def graph_signature(graph: fx.Graph) -> str:
271 """Fingerprint graph structure and planning facts to reject stale plans.
273 Args:
274 graph: Graph whose structure and FSDP metadata define the snapshot.
275 """
276 records = []
277 for node in graph.nodes:
278 metadata = {key: value for key, value in node.meta.items()
279 if key.startswith("fsdp_") or key in
280 ("comm_type", "comm_group", "param_name", "state_idx", "autograd_backward")}
281 records.append((node.name, str(node.target), str(node.args), str(node.kwargs), metadata))
282 return _digest(records)
285def validate_group_plans(plans_by_rank: Mapping[int, Union[BucketPlan, dict]], group: str) -> None:
286 """Validate reports exchanged by the caller, without creating collectives.
288 Args:
289 plans_by_rank: Plans or JSON-serializable to_dict() reports for every
290 global rank in this process group. Exchange reports, not FX nodes.
291 group: Exact process-group name used by the communication nodes.
293 Raises:
294 ValueError: Reports are missing ranks or disagree on membership/order/layout.
295 """
296 if not plans_by_rank:
297 raise ValueError("No rank reports supplied")
298 reports = {rank: plan.to_dict() if isinstance(plan, BucketPlan) else plan
299 for rank, plan in plans_by_rank.items()}
300 first = next(iter(reports.values()))
301 expected = first["group_ranks"].get(group)
302 if expected is None or set(plans_by_rank) != set(expected):
303 raise ValueError(f"Expected reports for group {group} ranks {expected}")
304 for report in reports.values():
305 if (tuple(report["group_ranks"].get(group, ())) != tuple(expected)
306 or report["group_digests"].get(group) != first["group_digests"].get(group)):
307 raise ValueError(f"Bucket plans disagree for group {group}")
310def _consumers(wait: fx.Node) -> set[fx.Node]:
311 """Return real consumers of a wait, traversing only immutable view chains."""
312 consumers, seen, stack = set(), {wait}, [wait]
313 while stack:
314 for user in stack.pop().users:
315 if is_view(user):
316 if user not in seen:
317 seen.add(user)
318 stack.append(user)
319 else:
320 consumers.add(user)
321 return consumers
324def _signature_reason(node: fx.Node, kind: str, groups: dict) -> Optional[str]:
325 """Return why the node signature, group or parameter identity blocks bucketing."""
326 backward = bool(node.meta.get("autograd_backward", False))
327 if backward != (kind != "fsdp_all_gather"):
328 return "phase_mismatch"
329 if node.kwargs or len(node.args) != (4 if kind == "fsdp_reduce_scatter" else 3):
330 return "unsupported_signature"
331 source, group = node.args[0], node.args[-1]
332 if not isinstance(source, fx.Node) or not isinstance(group, str) or group not in groups:
333 return "unknown_process_group"
334 degree = node.args[-2]
335 if not _static_int(degree) or degree <= 1 or degree != len(groups[group]):
336 return "group_size_mismatch"
337 if node.meta.get("comm_group", group) != group:
338 return "group_metadata_mismatch"
339 name, index = node.meta.get("param_name"), node.meta.get("state_idx")
340 if not isinstance(name, str) or not name or not _static_int(index) or index < 0:
341 return "missing_parameter_identity"
342 return None
345def _layout_reason(kind: str, degree: int, spec: dict) -> Optional[str]:
346 """Return why the recorded bucket spec cannot be packed statically."""
347 input_shape, output_shape = spec.get("input_shape"), spec.get("output_shape")
348 for shape in (input_shape, output_shape):
349 if (not isinstance(shape, (list, tuple)) or not shape
350 or not all(_static_int(dim) and dim > 0 for dim in shape)):
351 return "unknown_or_dynamic_shape"
352 dtype, device = spec.get("dtype"), spec.get("device")
353 if not isinstance(dtype, torch.dtype) or dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64):
354 return "unsupported_dtype"
355 if not isinstance(device, str):
356 return "unknown_device"
357 try:
358 torch.device(device)
359 except (RuntimeError, ValueError):
360 return "unknown_device"
361 smaller, larger = (output_shape, input_shape) if kind == "fsdp_reduce_scatter" else (input_shape, output_shape)
362 if tuple(larger) != (smaller[0] * degree, *smaller[1:]):
363 return "invalid_dim0_layout"
364 return None
367def _candidate(node: fx.Node, module_fqns: Optional[tuple], groups: dict, order: dict) -> tuple:
368 """Return (candidate, None) for a bucketable collective, else (None, skip reason)."""
369 kind = node.meta.get("comm_type")
370 if not is_op(node, _KINDS[kind]) or node.meta.get("fsdp_bucket_members"):
371 return None, "unsupported_or_already_bucketed"
372 reason = _signature_reason(node, kind, groups)
373 if reason:
374 return None, reason
375 spec = node.meta.get("fsdp_bucket_spec", {})
376 reason = _layout_reason(kind, node.args[-2], spec)
377 if reason:
378 return None, reason
379 waits = list(node.users)
380 if len(waits) != 1 or not is_op(waits[0], c10d.wait_tensor):
381 return None, "unsupported_wait_users"
382 consumers = _consumers(waits[0])
383 if not consumers:
384 return None, "no_consumer"
385 if kind == "fsdp_reduce_scatter":
386 if node.args[1] not in ("sum", "avg"):
387 return None, "unsupported_reduce_op"
388 if any(user.op != "output" for user in consumers):
389 return None, "non_output_rs_consumer"
390 if module_fqns is None:
391 return None, "outside_module_plan"
392 name = node.meta.get("param_name")
393 index = node.meta.get("state_idx")
394 backward = bool(node.meta.get("autograd_backward", False))
395 candidate = Candidate(node, waits[0], node.args[0], name, index, kind, backward, module_fqns[0], module_fqns,
396 node.args[-1], groups[node.args[-1]], spec.get("dtype"), spec.get("device"),
397 tuple(spec.get("input_shape")), tuple(spec.get("output_shape")), order[node.args[0]],
398 min(order[user] for user in consumers),
399 node.args[1] if kind == "fsdp_reduce_scatter" else None)
400 return candidate, None
403def _membership_reason(members: list[Candidate], candidate: Candidate, policy: BucketPolicy) -> Optional[str]:
404 """Return why compatibility, uniqueness or capacity rejects the candidate."""
405 if members[-1].key != candidate.key:
406 return "incompatible_key"
407 combined = [*members, candidate]
408 param_names = {m.param_name for m in combined}
409 state_indices = {m.state_idx for m in combined}
410 if len(param_names) != len(combined) or len(state_indices) != len(combined):
411 return "repeated_parameter_instance"
412 if policy.cap_bytes is not None and sum(m.size_bytes for m in combined) > policy.cap_bytes:
413 return "capacity"
414 return None
417def _window_reason(members: list[Candidate], candidate: Candidate, policy: BucketPolicy) -> Optional[str]:
418 """Return why scheduling windows or dependencies reject the fusion."""
419 combined = [*members, candidate]
420 if (policy.max_input_distance is not None
421 and max(m.ready for m in combined) - min(m.ready for m in combined) > policy.max_input_distance):
422 return "input_distance"
423 if (policy.max_consumer_distance is not None
424 and max(m.first_use for m in combined) - min(m.first_use for m in combined) > policy.max_consumer_distance):
425 return "consumer_distance"
426 # Only immutable result views may move after the latest input. Actual AG
427 # consumers (including frees) must remain after that insertion point.
428 if max(m.ready for m in combined) >= min(m.first_use for m in combined):
429 return "dependency_conflict"
430 return None
433def _owned_free(node: fx.Node, candidate: Candidate, combined: list) -> bool:
434 """Whether node frees exactly one member's waited result via copy-out."""
435 return (candidate.kind != "fsdp_reduce_scatter"
436 and node.op == "call_function" and node.target is _free_tensor_storage
437 and not node.kwargs and any(node.args == (member.wait,) for member in combined))
440def _intervening_reason(members: list[Candidate], candidate: Candidate, nodes: list, order: dict) -> Optional[str]:
441 """Return why nodes between the first member and the candidate block fusion."""
442 combined = [*members, candidate]
443 owned = {m.launch for m in combined}
444 preparations = set()
445 for member in combined:
446 if member.kind == "fsdp_reduce_scatter":
447 preparations.update(rs_layout_preparation(member.launch))
448 for node in nodes[order[members[0].launch] + 1:order[candidate.launch]]:
449 if node in preparations:
450 # Only these default, private RS layout nodes are known pure here.
451 continue
452 if node.meta.get("comm_type") in _FREES:
453 # Copy-out preserves each member's independent storage and original
454 # free point. Only that exact materialization may be crossed.
455 if not _owned_free(node, candidate, combined):
456 return "effect_or_storage_release"
457 elif is_barrier(node):
458 return "effect_or_storage_release"
459 if is_collective(node) and node not in owned:
460 return "intervening_collective"
461 if node.op not in ("placeholder", "get_attr", "output") and (
462 bool(node.meta.get("autograd_backward", False)) != candidate.backward):
463 return "phase_boundary"
464 return None
467def _split_reason(members: list[Candidate], candidate: Candidate, nodes: list, order: dict,
468 policy: BucketPolicy) -> Optional[str]:
469 """Return why candidate cannot join members, or None when fusion is safe."""
470 reason = _membership_reason(members, candidate, policy)
471 if reason:
472 return reason
473 reason = _window_reason(members, candidate, policy)
474 if reason:
475 return reason
476 return _intervening_reason(members, candidate, nodes, order)
479def _finalize_bucket(members: list[Candidate]) -> Bucket:
480 """Freeze bucket members in deterministic (state_idx, param_name) order."""
481 ordered = sorted(members, key=lambda member: (member.state_idx, member.param_name))
482 return Bucket(tuple(ordered))
485def _physical_buckets(candidates: list[Candidate], nodes: list, order: dict,
486 policy: BucketPolicy, decisions: list) -> list[Bucket]:
487 """Group by module/phase/compatibility, then split where legal motion stops."""
488 compatible = {}
489 for candidate in candidates:
490 compatible.setdefault(candidate.key, []).append(candidate)
491 buckets = []
492 for group in compatible.values():
493 current = []
494 for candidate in group:
495 reason = _split_reason(current, candidate, nodes, order, policy) if current else None
496 if reason:
497 buckets.append(_finalize_bucket(current))
498 decisions.append({"node": candidate.launch.name, "previous_node": current[-1].launch.name,
499 "action": "split", "reason": reason})
500 current = []
501 current.append(candidate)
502 if policy.cap_bytes is not None and candidate.size_bytes > policy.cap_bytes:
503 decisions.append({"node": candidate.launch.name, "action": "singleton", "reason": "oversized"})
504 if current:
505 buckets.append(_finalize_bucket(current))
506 # Model declaration order never silently reorders the collective stream.
507 buckets.sort(key=lambda bucket: min(order[member.launch] for member in bucket.members))
508 return buckets
511def _parameter_scope(name: object, module_plans: Optional[tuple]) -> Optional[tuple[str, ...]]:
512 if not isinstance(name, str) or not name:
513 return None
514 if module_plans is not None:
515 return match_module_bucket(name, module_plans)
516 return (name.rpartition(".")[0] or f"<parameter:{name}>",)
519def _log_plan(plan: BucketPlan, selected: dict, counts: Counter, order: dict) -> None:
520 """Log diagnostics where the plan is produced, without a second representation."""
521 _LOG.info("FSDP bucket plan: graph=%s summary=%s", plan.signature[:12], plan.summary)
522 if not _LOG.isEnabledFor(logging.DEBUG):
523 return
524 for scope, names in selected.items():
525 _LOG.debug("FSDP bucket scope: graph=%s modules=%s parameters=%s collectives=%s",
526 plan.signature[:12], scope, sorted(names), counts[scope])
527 for decision in plan.decisions:
528 _LOG.debug("FSDP bucket decision: graph=%s %s", plan.signature[:12], decision)
529 for bucket in plan.buckets:
530 members = bucket.members
531 first = members[0]
532 _LOG.debug("FSDP bucket planned: bucket=%s modules=%s kind=%s group=%s ranks=%s dtype=%s device=%s "
533 "reduce_op=%s members=%s input_span=%s original_launches=%s action=%s",
534 bucket.bucket_id, first.module_fqns, first.kind, first.group, first.ranks, first.dtype,
535 first.device, first.reduce_op, [member.param_name for member in members],
536 (min(member.ready for member in members), max(member.ready for member in members)),
537 [(member.launch.name, order[member.launch]) for member in members],
538 "rewrite_rs" if bucket.supports_rs_rewrite else
539 ("rewrite_ag_copy_out" if bucket.supports_ag_rewrite else "singleton"))
542def _validate_group_ranks(groups: dict) -> None:
543 """Reject process-group names or rank lists that cannot define a consistent plan."""
544 for group, ranks in groups.items():
545 if (not isinstance(group, str) or not ranks or len(set(ranks)) != len(ranks)
546 or any(not _static_int(rank) or rank < 0 for rank in ranks)):
547 raise ValueError("group_ranks must map PG names to distinct ordered non-negative global ranks")
550def _record_wire_collective(wire: dict, node: fx.Node, group: str) -> None:
551 """Append one rank-comparable wire record for a tagged collective."""
552 spec = dict(node.meta.get("fsdp_bucket_spec", {}))
553 if "device" in spec:
554 spec["device"] = str(spec["device"]).split(":", 1)[0]
555 reduce_op = (node.args[1] if node.meta.get("comm_type") == "fsdp_reduce_scatter"
556 and len(node.args) > 1 else None)
557 wire[group].append((node.meta.get("comm_type"), node.meta.get("param_name"),
558 node.meta.get("state_idx"), spec, reduce_op))
561def _scan_collectives(nodes: list, groups: dict, module_plans: Optional[tuple],
562 pp_degree: int, dp_mode: str) -> tuple:
563 """Scan tagged collectives into wire records, skip decisions and candidates.
565 Returns:
566 (candidates, wire, selected, counts, decisions, tagged_count); physical
567 bucketing appends further split/singleton decisions to ``decisions``.
568 """
569 order = {node: index for index, node in enumerate(nodes)}
570 decisions, candidates = [], []
571 wire = {group: [] for group in groups}
572 selected = {scope: set() for scope in module_plans or ()}
573 counts = Counter()
574 tagged_count = 0
575 for node in nodes:
576 if node.meta.get("comm_type") not in _KINDS:
577 continue
578 tagged_count += 1
579 name = node.meta.get("param_name")
580 scope = _parameter_scope(name, module_plans)
581 if scope is not None:
582 selected.setdefault(scope, set()).add(name)
583 counts[scope] += 1
584 group = node.args[-1] if node.args else None
585 if isinstance(group, str) and group in wire:
586 _record_wire_collective(wire, node, group)
587 if pp_degree != 1 or dp_mode != "fsdp":
588 candidate, reason = None, "unsupported_parallel_mode"
589 else:
590 candidate, reason = _candidate(node, scope, groups, order)
591 if reason:
592 decisions.append({"node": node.name, "action": "skip", "reason": reason})
593 else:
594 candidates.append(candidate)
595 return candidates, wire, selected, counts, decisions, tagged_count
598def _unmatched_module_paths(selected: dict) -> list:
599 """List scope paths that ended up selecting no parameter at all."""
600 unmatched = []
601 for scope, names in selected.items():
602 for path in scope:
603 if not any(name.startswith(path + ".") for name in names):
604 unmatched.append(path)
605 return unmatched
608def plan_fsdp_buckets(graph_module: fx.GraphModule, *, group_ranks: Mapping[str, tuple[int, ...]],
609 module_bucket_plans: Optional[ModuleBucketPlans] = None,
610 policy: Optional[BucketPolicy] = None, pp_degree: int = 1,
611 dp_mode: str = "fsdp") -> BucketPlan:
612 """Collect by model module and compatibility, then plan safe physical buckets.
614 Args:
615 graph_module: Joint graph after FSDPPass and before overlap scheduling.
616 group_ranks: Exact PG names mapped to ordered global ranks, supplied by caller.
617 module_bucket_plans: Explicit FQN groups, as in TorchTitan. None falls back
618 to immediate parameter parents; an empty list selects nothing. Unmatched
619 parameters remain untouched. Prefix ownership must be unambiguous.
620 policy: Optional capacity/distance limits; no size/distance cap by default.
621 pp_degree: Pipeline degree; this first implementation supports only one stage.
622 dp_mode: Only pure FSDP is supported; HSDP/DDP receive explicit skip decisions.
624 Returns:
625 Read-only plan with stable parameter layouts and explicit skip/split reasons.
626 RS and AG may launch after the latest input, before any real consumer.
627 AG rewriting uses independent copies, allowing exact member frees inside a bucket.
628 """
629 policy = policy or BucketPolicy()
630 module_plans = None if module_bucket_plans is None else normalize_module_bucket_plans(module_bucket_plans)
631 groups = {group: tuple(ranks) for group, ranks in group_ranks.items()}
632 _validate_group_ranks(groups)
633 graph_module.graph.lint()
634 nodes = list(graph_module.graph.nodes)
635 order = {node: index for index, node in enumerate(nodes)}
636 candidates, wire, selected, counts, decisions, tagged_count = _scan_collectives(
637 nodes, groups, module_plans, pp_degree, dp_mode)
638 buckets = _physical_buckets(candidates, nodes, order, policy, decisions)
639 signature = graph_signature(graph_module.graph)
640 buckets = tuple(replace(bucket, bucket_id=f"{signature[:12]}:bucket_{index}")
641 for index, bucket in enumerate(buckets))
642 unmatched_module_paths = (_unmatched_module_paths(selected)
643 if module_plans is not None else [])
644 summary = {"tagged_collectives": tagged_count, "selected_parameters": sum(map(len, selected.values())),
645 "selected_collectives": sum(counts.values()), "candidate_collectives": len(candidates),
646 "buckets": len(buckets), "singletons": sum(len(bucket.members) == 1 for bucket in buckets),
647 "rs_rewrite_candidates": sum(bucket.supports_rs_rewrite for bucket in buckets),
648 "ag_rewrite_candidates": sum(bucket.supports_ag_rewrite for bucket in buckets),
649 "reason_counts": {action: dict(Counter(item["reason"] for item in decisions if item["action"] == action))
650 for action in ("skip", "split", "singleton")},
651 "unmatched_module_paths": unmatched_module_paths}
652 plan = BucketPlan(buckets, tuple(decisions), signature, groups, wire, module_plans, summary)
653 _log_plan(plan, selected, counts, order)
654 return plan