Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / fsdp_pass.py: 95%
635 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"""
16FSDP / DDP / HSDP Pass - Data Parallel Partitioning Pass
18Operates on the joint fwd+bwd graph produced by the tracer, where
19parameters/buffers are **static inputs** (leading placeholders), not get_attr
20nodes. The data-parallel mode comes from ``PassConfig.dp_mode`` and mirrors
21simplefsdp's ``data_parallel`` modes:
23- ``"fsdp"`` (fully_shard): parameters sharded on dim 0; forward all_gather,
24 backward reduce_scatter.
25- ``"ddp"`` (replicate): parameters stay replicated; no all_gather, gradients
26 are all-reduced on the ``dp_replicate`` axis.
27- ``"hsdp"`` (hybrid_shard): parameters replicated on ``dp_replicate`` and
28 sharded on ``fsdp``; forward all_gather on ``fsdp``, backward reduce_scatter
29 on ``fsdp`` then all_reduce on ``dp_replicate``.
31Responsibilities:
321. Identify parameter placeholders belonging to DP-wrapped modules
33 (via GraphParallelPlan, exact FQN or pattern)
342. Insert AllGather after each such placeholder (Shard -> Replicate), so the
35 computation body operates on full parameters while the graph input stays
36 sharded (skipped for ``"ddp"``)
373. Insert gradient reduction on the gradient outputs of DP parameters —
38 reduce_scatter (FSDP/HSDP) and/or all_reduce (DDP/HSDP)
394. Physically shard the *live model's* parameters in place (dim 0, by FSDP
40 rank) for FSDP/HSDP, so ``model.parameters()`` already holds the local
41 shard and the trainer / optimizer need no DP awareness at all
425. Reshard (``fsdp_reshard_after_forward``): sink each AllGather to its first
43 forward use, free the replicated parameter after its last forward read,
44 and, for parameters the backward still needs, re-gather plus rematerialize
45 the saved view chain just before the first backward consumer. Peak memory
46 tracks the forward working set instead of every replicated parameter.
48All DP logic (which parameters, the collectives, and the sharding itself)
49lives in this pass; the trainer simply feeds ``model.parameters()``.
50"""
52__all__ = ["FSDPPass"]
54import logging
55import operator
56from contextlib import contextmanager
57from dataclasses import dataclass
58from typing import Any, Dict, Iterator, List, Optional, Set, Tuple
60import torch
61import torch.distributed as dist
62from torch import fx, nn
63from torch._subclasses import FakeTensor
64from torch.distributed.distributed_c10d import _resolve_process_group
65from torch.fx.node import _side_effectful_functions
66from torch.ops import _c10d_functional
68from ...graph_parallel_plan import GraphParallelPlan
69from ...pass_config import GradientReduction, PassConfig, normalize_dp_mode
70from ..base import GraphPass
71from ..overlap.utils import rs_layout_preparation
73_LOG = logging.getLogger(__name__)
75_GETITEM = operator.getitem
76_DEFAULT_GRADIENT_REDUCTION = GradientReduction()
79def tag_gradient_identity(wait: fx.Node, param_name: str, state_idx: int) -> None:
80 """Attach the known output-to-parameter identity to a reduction chain.
82 Args:
83 wait: Final waited gradient returned by FSDPPass.
84 param_name: Canonical traced parameter FQN.
85 state_idx: Position in the traced state inputs.
86 """
87 node = wait
88 while node.op == "call_function" and str(node.target).startswith(
89 "_c10d_functional."
90 ):
91 node.meta.update(param_name=param_name, state_idx=state_idx)
92 if not node.args or not isinstance(node.args[0], fx.Node):
93 break
94 node = node.args[0]
97def _annotate_rs_spec(node: fx.Node, degree: int) -> None:
98 """Attach packed input/output shapes to one reduce-scatter node."""
99 source = node.args[0]
100 preparation = rs_layout_preparation(node)
101 if preparation:
102 # Contiguous changes strides, not the padded shape/dtype/device.
103 source = preparation[0].args[0]
104 value = source.meta.get("val")
105 spec = source.meta.get("fsdp_padded_input_spec")
106 if spec is None and value is not None and hasattr(value, "shape"):
107 spec = {
108 "shape": tuple(value.shape),
109 "dtype": value.dtype,
110 "device": str(value.device),
111 }
112 if spec is None:
113 return
114 shape = spec["shape"]
115 static = shape and all(
116 isinstance(dim, int) and not isinstance(dim, bool) and dim > 0 for dim in shape
117 )
118 if not static or shape[0] % degree:
119 return
120 node.meta["fsdp_bucket_spec"] = {
121 "input_shape": tuple(shape),
122 "output_shape": (shape[0] // degree, *shape[1:]),
123 "dtype": spec["dtype"],
124 "device": spec["device"],
125 }
128def _annotate_ag_spec(node: fx.Node, degree: int, params: dict) -> None:
129 """Attach shard/full shapes of one all-gather node from its parameter."""
130 param = params.get(node.meta.get("param_name"))
131 if param is None:
132 return
133 shape = tuple(param.shape)
134 if not shape or any(
135 not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 for dim in shape
136 ):
137 return
138 node.meta["fsdp_bucket_spec"] = {
139 "input_shape": shape,
140 "output_shape": (shape[0] * degree, *shape[1:]),
141 "dtype": param.dtype,
142 "device": str(param.device),
143 }
146def annotate_fsdp_bucket_specs(graph_module: fx.GraphModule, model: nn.Module) -> None:
147 """Record actual padded shard shapes after the live model has been sharded.
149 Args:
150 graph_module: Graph transformed by FSDPPass.
151 model: Live model whose parameters now have their rank-local shard shapes.
152 """
153 params = dict(model.named_parameters())
154 for node in graph_module.graph.nodes:
155 kind = node.meta.get("comm_type")
156 if kind not in (
157 "fsdp_all_gather",
158 "fsdp_all_gather_backward",
159 "fsdp_reduce_scatter",
160 ):
161 continue
162 degree = node.meta.get("fsdp_degree")
163 if (not isinstance(degree, int) or isinstance(degree, bool)) or degree <= 1:
164 continue
165 if kind == "fsdp_reduce_scatter":
166 _annotate_rs_spec(node, degree)
167 else:
168 _annotate_ag_spec(node, degree, params)
171def _free_tensor_storage(tensor: torch.Tensor) -> None:
172 """Release ``tensor``'s storage in place (FSDP reshard free step).
174 Inserted as a side-effecting FX node between an unsharded parameter's last
175 forward reader and its backward re-gather, so the replicated full
176 parameter (and every view aliasing it) does not stay resident for the
177 whole joint graph. Returns ``None``; the node exists only for its effect.
178 Mirrors ``fully_shard``'s ``free_unsharded_param`` storage release.
179 """
180 # FakeTensorProp (InductorPass shape propagation) executes this node on
181 # fake values whose storage cannot be resized; there is nothing to free.
182 if isinstance(tensor, FakeTensor):
183 return
184 if isinstance(tensor, torch.Tensor) and tensor.untyped_storage().size() > 0:
185 tensor.untyped_storage().resize_(0)
188# FX dead-code elimination drops call_function nodes it judges pure; our free
189# op is pure to FX but has a real side effect. Register it so a later
190# ``eliminate_dead_code`` (a future pass, or a user) cannot delete the free.
191_side_effectful_functions.add(_free_tensor_storage)
194@dataclass
195class _UnshardedParam:
196 """Forward all_gather that materializes one FSDP parameter.
198 ``param_node`` is the sharded placeholder (graph input); ``ag_node`` /
199 ``wait_node`` produce and expose the replicated (unsharded) full parameter
200 the forward — and, after a reshard re-gather, the backward — read.
201 """
203 param_node: fx.Node
204 ag_node: fx.Node
205 wait_node: fx.Node
208@dataclass
209class _DpParamClassification:
210 """How DP handles each trainable parameter's gradient.
212 Produced by ``_classify_dp_params``:
214 - ``sharded_nodes``: marked, trainable parameter placeholders that get an
215 AllGather (Shard -> Replicate) and a reduce_scatter gradient, and are
216 physically sharded in the live model. Dim-0-divisible params are chunked
217 directly; non-divisible ones are pad-sharded (see ``_PaddedParam``).
218 - ``replicated_indices``: state indices of parameters kept replicated (never
219 chunked) whose gradient is averaged over the DP group with ``all_reduce``
220 instead of reduce-scattered — a marked 0-dim scalar (cannot be sharded), or
221 a partial-plan leftover that falls back to DDP.
222 - ``unmarked_trainable``: ``(state_idx, fqn)`` for the partial-plan leftovers
223 (a subset of ``replicated_indices``), kept only so the caller can log that
224 the plan did not cover them.
225 """
227 sharded_nodes: List[fx.Node]
228 replicated_indices: Set[int]
229 unmarked_trainable: List[Tuple[int, str]]
232@dataclass
233class _PaddedParam:
234 """A marked parameter that cannot be sharded evenly, pad-sharded instead.
236 Mirrors eager ``fully_shard`` (``core/fully_shard/hsdp_param.py``): dim 0 of
237 ``n`` elements is zero-padded to ``n_pad = ceil(n / degree) * degree`` so
238 every rank owns a uniform ``n_pad / degree``-element shard. The graph
239 narrows the gathered ``[n_pad, ...]`` back to ``[n, ...]`` for the compute,
240 and the backward pads the ``[n, ...]`` gradient back to ``[n_pad, ...]``
241 before reduce_scatter, so the returned local gradient matches the padded
242 shard (the padding region gets an all-zero gradient, hence never moves).
243 """
245 n: int
246 n_pad: int
247 dim0_pad: int
248 ndim: int
251@dataclass(frozen=True)
252class _GradReductionInfo:
253 """A parameter's selected mesh view and independent normalization policy."""
255 group_name: str
256 degree: int
257 reduction: GradientReduction
260def _iter_arg_nodes(node: fx.Node) -> Iterator[fx.Node]:
261 """Yield every ``fx.Node`` in ``node.args`` / ``node.kwargs`` (nested)."""
262 stack: List[Any] = list(node.args) + list(node.kwargs.values())
263 while stack:
264 item = stack.pop()
265 if isinstance(item, fx.Node):
266 yield item
267 elif isinstance(item, (tuple, list)):
268 stack.extend(item)
269 elif isinstance(item, dict):
270 stack.extend(item.values())
273@contextmanager
274def _insertion_point(
275 graph: fx.Graph, anchor: Optional[fx.Node], fallback: fx.Node
276) -> Iterator[None]:
277 """Yield a graph insertion context.
279 Owns the ``anchor``-vs-``fallback`` choice so the gather-insertion body
280 stays a single code path: when ``anchor`` is given, new nodes land right
281 before it; otherwise they land right after ``fallback``.
282 """
283 if anchor is None:
284 with graph.inserting_after(fallback):
285 yield
286 else:
287 with graph.inserting_before(anchor):
288 yield
291class FSDPPass(GraphPass):
292 """
293 FSDP Partitioning Pass (Module-Level, static-input graph)
295 Responsibilities:
296 1. Only shard parameters in modules marked for FSDP wrapping
297 2. Insert AllGather after each FSDP parameter placeholder
298 (Shard -> Replicate)
299 3. Insert ReduceScatter after Backward (Replicate -> Shard)
300 4. Physically shard the live model's parameters in place so the
301 trainer / optimizer stay FSDP-agnostic
303 Key Difference from Old Implementation:
304 - Old: shard all parameters via get_attr nodes and physically rewrite the
305 GraphModule's stored tensors
306 - New: parameters are static graph inputs; the pass inserts collectives
307 into the graph and shards the live model's parameters in place
308 """
310 name = "fsdp_parallel"
312 def __init__(
313 self,
314 fsdp_group_name: Optional[str] = None,
315 parallel_plan: Optional[GraphParallelPlan] = None,
316 ) -> None:
317 """Initialize FSDP pass state.
319 Args:
320 fsdp_group_name: Process-group name registered for FSDP
321 collectives. Defaults to ``"fsdp"``. May be overridden
322 per-run via the ``fsdp_group_name`` kwarg in ``run``.
323 parallel_plan: Declarative plan identifying which modules to
324 shard. When ``None``, all parameters are sharded.
325 """
326 super().__init__()
327 self._fsdp_group_name = fsdp_group_name or "fsdp"
328 self._parallel_plan = parallel_plan
329 # Resolved from ``pass_config.fsdp_degree`` (or world_size) at
330 # ``run`` entry; left as ``None`` here so a stray access before
331 # ``run`` fails loudly instead of silently using a wrong default.
332 self._fsdp_degree: Optional[int] = None
333 self._expert_param_names: Set[str] = set()
334 self._expert_fsdp_group_name = "fsdp_expert"
335 self._expert_fsdp_degree: Optional[int] = None
336 self._gradient_reductions: Dict[str, GradientReduction] = {}
337 self._source_gradient_groups: Dict[str, str] = {}
338 self._parameter_groups: Dict[str, tuple[str, int]] = {}
339 # Data-parallel mode ("fsdp" | "ddp" | "hsdp") and the replicate-axis
340 # group/degree used by DDP / HSDP gradient reduction.
341 self._dp_mode = "fsdp"
342 self._dp_replicate_group_name = "dp_replicate"
343 self._dp_replicate_degree: Optional[int] = None
344 self._processed_params: Set[str] = set()
345 self._fsdp_modules: Set[str] = set()
346 # Per-param forward all_gather record (the unsharded parameter),
347 # populated by ``_insert_all_gather_for_params`` and consumed by the
348 # reshard step.
349 self._unsharded_params: Dict[str, _UnshardedParam] = {}
350 # Non-divisible marked params keyed by FQN; populated by
351 # ``_classify_dp_params`` and consumed by the all_gather / grad-reduce /
352 # live-shard steps so they are pad-sharded like eager ``fully_shard``.
353 self._padded_params: Dict[str, _PaddedParam] = {}
355 def run(
356 self,
357 graph_module: fx.GraphModule,
358 pass_config: PassConfig,
359 **kwargs: Any,
360 ) -> fx.GraphModule:
361 """Insert AllGather/ReduceScatter and shard the live model.
363 Args:
364 graph_module: Joint fwd+bwd FX graph from the tracer.
365 pass_config: Parallel configuration. ``fsdp_degree`` is
366 read directly; ``None`` falls back to ``world_size`` (the
367 FSDP-only path). A ``TypeError`` here means a non-Protocol
368 config was passed — fix at the caller, do not paper over.
369 **kwargs: Must include ``model`` (the live ``nn.Module``) and
370 ``fsdp_group_name`` / ``parallel_plan`` as needed.
372 Returns:
373 The transformed graph module.
374 """
375 if not dist.is_initialized() or dist.get_world_size() == 1:
376 _LOG.info("Skipped: distributed not initialized or world_size=1")
377 return graph_module
379 model = self._setup_dp_config(pass_config, kwargs)
381 # Identify DP parameter placeholders. The joint graph's
382 # parameters/buffers are static inputs (leading placeholders, not
383 # get_attr nodes): locate them by position via the state layout the
384 # tracer attached, then keep only those in DP-marked modules. For
385 # FSDP/HSDP a marked param whose leading dim is not divisible by the
386 # shard degree is pad-sharded (like eager fully_shard) so it is still
387 # sharded and gradient-reduced; scalars stay replicated. DDP has no
388 # shard axis so every marked parameter is reduced.
389 state_fqns = getattr(graph_module, "state_fqns", [])
390 num_state_inputs = getattr(graph_module, "num_state_inputs", 0)
391 # Param-vs-buffer flag per leading state input. Absent on graphs
392 # traced before this flag existed; default to all-params so old
393 # traces keep the previous (parameter-only) behaviour.
394 state_is_param = getattr(graph_module, "state_is_param", None)
396 self._reset_run_state()
397 classification = self._classify_dp_params(
398 graph_module, state_fqns, num_state_inputs, state_is_param, model
399 )
400 param_nodes = classification.sharded_nodes
401 graph_module.fsdp_sharded_param_names = frozenset(
402 node.meta["param_name"] for node in param_nodes if self._dp_mode != "ddp"
403 )
405 _LOG.info(
406 "Identified %d sharded (incl. %d pad-sharded) + %d replicated DP "
407 "parameter nodes out of %d total state inputs",
408 len(param_nodes),
409 len(self._padded_params),
410 len(classification.replicated_indices),
411 num_state_inputs,
412 )
414 # A partial plan that leaves trainable params unmarked keeps their local
415 # gradient (never reduced) — a silent divergence source across DP ranks,
416 # so warn loudly before continuing.
417 if classification.unmarked_trainable:
418 self._warn_unmarked_trainable(classification.unmarked_trainable)
420 if not param_nodes and not classification.replicated_indices and not self._gradient_reductions:
421 _LOG.warning(
422 "No DP parameters found, check GraphParallelPlan or model structure"
423 )
424 return graph_module
426 graph_module = self._apply_dp_transforms(
427 graph_module,
428 model,
429 classification,
430 pass_config,
431 state_fqns,
432 num_state_inputs,
433 state_is_param,
434 )
436 _LOG.info(
437 "Completed dp_mode=%s, processed %d parameters",
438 self._dp_mode,
439 len(param_nodes),
440 )
442 graph_module.recompile()
443 return graph_module
445 def _setup_dp_config(
446 self, pass_config: PassConfig, kwargs: Dict[str, Any]
447 ) -> nn.Module:
448 """Resolve run-time DP config, stash it on ``self``, return the model.
450 Reads ``dp_mode`` / ``fsdp_degree`` / ``dp_replicate_degree`` off
451 ``pass_config`` (the last two fall back to ``world_size`` when unset)
452 and the group names / ``parallel_plan`` / ``model`` off ``kwargs``.
454 Returns:
455 The live ``nn.Module`` to physically shard.
457 Raises:
458 ValueError: If ``model`` is missing from ``kwargs``.
459 """
460 self._dp_mode = normalize_dp_mode(getattr(pass_config, "dp_mode", "fsdp"))
461 # FSDP shard-axis group size: use the explicitly configured degree when
462 # present (required for TP+FSDP / HSDP, where the FSDP group is a
463 # proper sub-group of the world — using world_size would over-shard
464 # along the TP axis).
465 configured = pass_config.fsdp_degree
466 self._fsdp_degree = configured if configured else dist.get_world_size()
467 self._fsdp_group_name = kwargs.get("fsdp_group_name", self._fsdp_group_name)
468 self._expert_param_names = set(kwargs.get("expert_param_names", ()))
469 self._gradient_reductions = dict(kwargs.get("gradient_reductions", {}))
470 self._source_gradient_groups = dict(kwargs.get("source_gradient_groups", {}))
471 self._parameter_groups = dict(kwargs.get("parameter_groups", {}))
472 if not self._expert_param_names.issubset(self._gradient_reductions):
473 raise ValueError("Expert parameters require explicit gradient reduction policies")
474 self._expert_fsdp_group_name = kwargs.get(
475 "expert_fsdp_group_name", self._expert_fsdp_group_name
476 )
477 self._expert_fsdp_degree = (
478 dist.get_world_size(_resolve_process_group(self._expert_fsdp_group_name))
479 if self._expert_param_names
480 else None
481 )
482 self._dp_replicate_group_name = kwargs.get(
483 "dp_replicate_group_name", self._dp_replicate_group_name
484 )
485 # Replicate-axis degree used by DDP / HSDP gradient reduction.
486 replicate_degree = getattr(pass_config, "dp_replicate_degree", None)
487 if self._dp_mode == "ddp":
488 self._dp_replicate_degree = replicate_degree or dist.get_world_size()
489 elif self._dp_mode == "hsdp":
490 self._dp_replicate_degree = replicate_degree or (
491 dist.get_world_size() // self._fsdp_degree
492 )
493 else:
494 self._dp_replicate_degree = replicate_degree
495 self._parallel_plan = kwargs.get("parallel_plan", self._parallel_plan)
496 model = kwargs.get("model")
497 if model is None:
498 raise ValueError(
499 "FSDPPass requires the live model via kwargs (model=...) so it "
500 "can physically shard parameters; the trainer passes it through "
501 "in compile()"
502 )
504 _LOG.info(
505 "Running dp_mode=%s fsdp_degree=%s dp_replicate_degree=%s world_size=%s",
506 self._dp_mode,
507 self._fsdp_degree,
508 self._dp_replicate_degree,
509 dist.get_world_size(),
510 )
511 return model
513 def _reset_run_state(self) -> None:
514 """Clear per-run state so a reused pass does not carry names across graphs."""
515 self._processed_params = set()
516 self._fsdp_modules = set()
517 self._unsharded_params = {}
518 self._padded_params = {}
520 def _apply_dp_transforms(
521 self,
522 graph_module: fx.GraphModule,
523 model: nn.Module,
524 classification: _DpParamClassification,
525 pass_config: PassConfig,
526 state_fqns: List[str],
527 num_state_inputs: int,
528 state_is_param: Optional[List[bool]],
529 ) -> fx.GraphModule:
530 """Insert gather / grad-reduce / reshard and shard the live model.
532 DDP replicates parameters: no all_gather, no live-model sharding and no
533 reshard — only the gradient all-reduce path runs. FSDP/HSDP also
534 all_gather each parameter for the forward and pad/shard the live model
535 in place.
536 """
537 param_nodes = classification.sharded_nodes
538 if self._dp_mode != "ddp":
539 graph_module = self._insert_all_gather_for_params(
540 graph_module, param_nodes, pass_config
541 )
543 sharded_param_indices = frozenset(
544 node.meta["state_idx"] for node in param_nodes
545 )
546 graph_module = self._insert_grad_reduction(
547 graph_module,
548 sharded_param_indices,
549 classification.replicated_indices,
550 state_fqns,
551 num_state_inputs,
552 state_is_param,
553 model,
554 )
556 if self._dp_mode != "ddp":
557 # Reshard: free each replicated parameter once forward is done and
558 # re-gather + rematerialize it for the backward. Runs after the
559 # grad reduction so the grad path is untouched; no-op when disabled.
560 graph_module = self._insert_reshard_logic(
561 graph_module, param_nodes, pass_config
562 )
564 # Shard the live model's parameters in place (dim 0, by this rank's
565 # index in the FSDP group). ``model.parameters()`` then yields the
566 # shards, so the trainer's optimizer / grad accumulation stay
567 # DP-agnostic; the graph re-gathers each step.
568 self._shard_live_model_params(model)
569 annotate_fsdp_bucket_specs(graph_module, model)
571 return graph_module
573 def _param_group(self, name: str) -> tuple[str, int]:
574 """Choose the expert or dense shard axis for one parameter."""
575 if name in self._parameter_groups:
576 return self._parameter_groups[name]
577 if name in self._expert_param_names:
578 return self._expert_fsdp_group_name, self._expert_fsdp_degree
579 return self._fsdp_group_name, self._fsdp_degree
581 def _classify_dp_params(
582 self,
583 graph_module: fx.GraphModule,
584 state_fqns: List[str],
585 num_state_inputs: int,
586 state_is_param: Optional[List[bool]] = None,
587 model: Optional[nn.Module] = None,
588 ) -> _DpParamClassification:
589 """Classify parameter placeholders by how DP handles their gradient.
591 Parameters are the leading ``num_state_inputs`` placeholders of the
592 joint graph, in ``state_fqns`` order. A parameter is DP-marked when any
593 ancestor module FQN (e.g. ``layers.0.attention.wq`` for
594 ``layers.0.attention.wq.weight``) is marked in the GraphParallelPlan —
595 either exactly or via pattern (e.g. ``layers.*``).
597 Buffers (e.g. RoPE's non-persistent ``cache``) and frozen parameters
598 are skipped: the tracer omits them from the gradient list and the
599 live-model sharder leaves them full-size. The pad-shard and unmarked-
600 fallback decisions are delegated to ``_marked_param_stays_replicated``
601 and ``_unmarked_disposition``.
603 Returns:
604 A ``_DpParamClassification`` (see its docstring).
605 """
606 sharded_nodes: List[fx.Node] = []
607 replicated_indices: Set[int] = set()
608 unmarked_trainable: List[Tuple[int, str]] = []
610 placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
611 if len(placeholders) < num_state_inputs:
612 return _DpParamClassification(
613 sharded_nodes, replicated_indices, unmarked_trainable
614 )
616 # Build a FQN -> Parameter lookup so the divisibility gate uses the
617 # same source of truth as ``_shard_live_model_params`` (which iterates
618 # ``model.named_parameters``). ``state_fqns`` covers both parameters
619 # and buffers; only parameters are reachable here because buffers are
620 # filtered out by ``state_is_param`` above.
621 param_lookup: Dict[str, nn.Parameter] = (
622 dict(model.named_parameters(remove_duplicate=False))
623 if model is not None
624 else {}
625 )
627 for idx in range(num_state_inputs):
628 fqn, param, skip = self._resolve_dp_param(
629 idx, state_fqns, state_is_param, param_lookup, model
630 )
631 if skip:
632 continue
634 _, degree = self._param_group(fqn)
635 if fqn in self._expert_param_names and degree == 1:
636 continue
638 fallback = self._unmarked_disposition(fqn, param, model)
639 if fallback is not None:
640 if fallback:
641 replicated_indices.add(idx)
642 unmarked_trainable.append((idx, fqn))
643 continue
645 if self._marked_param_stays_replicated(fqn, param):
646 replicated_indices.add(idx)
647 continue
649 sharded_nodes.append(
650 self._mark_sharded_param(idx, fqn, degree, placeholders[idx], param)
651 )
653 return _DpParamClassification(
654 sharded_nodes, replicated_indices, unmarked_trainable
655 )
657 def _resolve_dp_param(
658 self,
659 idx: int,
660 state_fqns: List[str],
661 state_is_param: Optional[List[bool]],
662 param_lookup: Dict[str, nn.Parameter],
663 model: Optional[nn.Module],
664 ) -> Tuple[Optional[str], Optional[nn.Parameter], bool]:
665 """Resolve one state placeholder to ``(fqn, param, skip)``.
667 Buffers and frozen parameters are skipped: the tracer omits them from
668 the gradient list and the live-model sharder leaves them full-size.
669 """
670 if state_is_param is not None and not state_is_param[idx]:
671 # Buffer, not a parameter -> never FSDP-sharded.
672 return None, None, True
674 fqn = state_fqns[idx]
675 param = param_lookup.get(fqn)
676 if model is not None and param is None:
677 return None, None, True
679 if param is not None and not param.requires_grad:
680 _LOG.info("Skip %s: requires_grad=False (frozen)", fqn)
681 return None, None, True
682 return fqn, param, False
684 def _mark_sharded_param(
685 self,
686 idx: int,
687 fqn: str,
688 degree: int,
689 node: fx.Node,
690 param: Optional[nn.Parameter],
691 ) -> fx.Node:
692 """Tag a DP-sharded placeholder and record its live-model shard size."""
693 node.meta["state_idx"] = idx
694 node.meta["param_name"] = fqn
695 node.meta["fsdp_degree"] = degree
696 node.meta["is_param"] = True
697 # Record the TP-local, unsharded size before live-model FSDP slicing.
698 # Newly inserted collective nodes have no shape-propagated ``val``.
699 value = param if param is not None else node.meta.get("val")
700 if isinstance(value, torch.Tensor) and isinstance(value.numel(), int):
701 node.meta["fsdp_unsharded_nbytes"] = value.numel() * value.element_size()
702 self._fsdp_modules.add(self._get_parent_module_fqn(fqn))
703 return node
705 def _unmarked_disposition(
706 self,
707 fqn: str,
708 param: Optional[nn.Parameter],
709 model: Optional[nn.Module],
710 ) -> Optional[bool]:
711 """Decide how a trainable param a partial plan left unmarked is handled.
713 Rather than leave it with a local, unreduced gradient (silent
714 divergence across DP ranks), fall it back to DDP: keep it replicated
715 and average its gradient over the DP group.
717 Returns:
718 ``True`` to replicate it and warn the caller, ``False`` to skip it
719 silently (``param is None`` with a live model means a buffer, only
720 reachable when an old trace lacks ``state_is_param``), or ``None``
721 when the param is marked and the caller should shard it normally.
722 """
723 if self._parallel_plan is None or self._param_belongs_to_fsdp_module(fqn):
724 return None
725 return model is None or param is not None
727 def _marked_param_stays_replicated(
728 self, fqn: str, param: Optional[nn.Parameter]
729 ) -> bool:
730 """Decide whether a marked param is kept replicated instead of sharded.
732 A marked param whose dim 0 is not divisible by ``fsdp_degree`` cannot be
733 chunked evenly. A 0-dim scalar cannot be padded along dim 0, so it stays
734 replicated and its gradient is averaged over the shard axis. Every
735 other non-divisible param is recorded for pad-sharding (mirroring eager
736 ``fully_shard``): dim 0 is zero-padded to a multiple of the degree, so
737 it is still sharded and gradient-reduced. DDP has no shard axis, so
738 every marked param is handled the same way regardless of divisibility.
739 """
740 if self._dp_mode == "ddp" or param is None:
741 return False
742 _, degree = self._param_group(fqn)
743 shape = param.shape
744 if fqn in self._expert_param_names and (not shape or shape[0] % degree):
745 raise ValueError(
746 f"Graph FSDP+EP expert parameter {fqn} with shape {tuple(shape)} "
747 f"cannot shard across edp_shard degree {degree}"
748 )
749 if shape and shape[0] % degree == 0:
750 return False
751 if not shape:
752 _LOG.info(
753 "Scalar param %s stays replicated: 0-dim cannot be "
754 "pad-sharded; its gradient will be all-reduced",
755 fqn,
756 )
757 return True
758 n = int(shape[0])
759 n_pad = -(-n // degree) * degree
760 self._padded_params[fqn] = _PaddedParam(
761 n=n, n_pad=n_pad, dim0_pad=n_pad - n, ndim=param.ndim
762 )
763 _LOG.info(
764 "Param %s pad-sharded: dim 0 %d -> %d (fsdp_degree %d)",
765 fqn,
766 n,
767 n_pad,
768 degree,
769 )
770 return False
772 @staticmethod
773 def _warn_unmarked_trainable(entries: List[Tuple[int, str]]) -> None:
774 """Report trainable params a partial plan left unmarked.
776 They are not FSDP-sharded, so they fall back to DDP: kept replicated and
777 their gradient averaged over the DP group. The warning is informational —
778 the plan simply did not cover them, and they cost no sharded memory but a
779 replicated gradient all-reduce.
780 """
781 shown = [fqn for _, fqn in entries[:10]]
782 remaining = len(entries) - len(shown)
783 _LOG.warning(
784 "%d trainable parameter(s) are not marked for DP; falling back to "
785 "DDP (kept replicated, gradient all-reduced over the DP group): "
786 "%s%s. Mark them via GraphParallelPlan.fsdp_mark()/"
787 "fsdp_mark_pattern() to FSDP-shard them instead.",
788 len(entries),
789 ", ".join(shown),
790 f" (+{remaining} more)" if remaining > 0 else "",
791 )
793 def _shard_live_model_params(self, model: nn.Module) -> None:
794 """
795 Physically shard the live model's parameters in place (dim 0).
797 After this, model.parameters() yields the local shards.
798 """
799 # Use the FSDP group's local rank (NOT the global rank) as the
800 # chunk index — when fsdp_degree < world_size (TP+FSDP), the global
801 # rank exceeds the chunk count and causes IndexError.
802 fsdp_pg = _resolve_process_group(self._fsdp_group_name)
803 group_ranks = {self._fsdp_group_name: dist.get_rank(group=fsdp_pg)}
804 sharded_count = 0
806 for name, param in model.named_parameters():
807 if not param.requires_grad:
808 continue
810 group_name, degree = self._param_group(name)
811 if name in self._expert_param_names and degree == 1:
812 continue
814 # Same plan gate as ``_classify_dp_params``: only parameters inside
815 # FSDP-marked modules are sharded. Skipping the gate here would
816 # shard params the graph still expects full-rank
817 # (no AllGather inserted), causing a shape mismatch at
818 # ``run_traced_graph`` time on any non-``*`` plan.
819 if (
820 self._parallel_plan is not None
821 and not self._param_belongs_to_fsdp_module(name)
822 ):
823 _LOG.info(
824 "Skip %s: not in an FSDP-wrapped module",
825 name,
826 )
827 continue
829 if group_name not in group_ranks:
830 group_ranks[group_name] = dist.get_rank(
831 group=_resolve_process_group(group_name)
832 )
833 rank = group_ranks[group_name]
835 # Pad-shard a non-divisible param (recorded by
836 # ``_classify_dp_params``): zero-pad dim 0 to a multiple of the
837 # degree, then take this rank's equal-size chunk — exactly what
838 # eager ``fully_shard`` does. A 0-dim scalar cannot be padded and
839 # stays replicated.
840 pad_info = self._padded_params.get(name)
841 if pad_info is not None:
842 original_shape = param.shape
843 full = param.detach()
844 if pad_info.dim0_pad:
845 full = torch.cat(
846 [full, full.new_zeros((pad_info.dim0_pad, *full.shape[1:]))],
847 dim=0,
848 )
849 param.data = full.chunk(degree, dim=0)[rank].clone()
850 sharded_count += 1
851 _LOG.info(
852 "Pad-sharded %s: %s -> %s",
853 name,
854 list(original_shape),
855 list(param.shape),
856 )
857 continue
859 # Mirror ``_classify_dp_params``: a 0-dim scalar stays replicated.
860 if not param.shape or param.shape[0] % degree != 0:
861 _LOG.info(
862 "Skip %s: dim 0 (%s) not divisible by fsdp_degree (%s)",
863 name,
864 param.shape[0] if param.shape else "scalar",
865 degree,
866 )
867 continue
869 original_shape = param.shape
870 param.data = param.detach().chunk(degree, dim=0)[rank].clone()
871 sharded_count += 1
872 _LOG.info(
873 "Sharded %s: %s -> %s",
874 name,
875 list(original_shape),
876 list(param.shape),
877 )
879 _LOG.info("Total sharded parameters: %d", sharded_count)
881 def _param_belongs_to_fsdp_module(self, param_fqn: str) -> bool:
882 """Check if a parameter FQN belongs to an FSDP-wrapped module.
884 The parameter's own module FQN (``layers.0.attention.wq.weight`` ->
885 ``layers.0.attention.wq``) plus every ancestor (``layers.0``, ...) is
886 tested, so both ``fsdp_mark("layers.0.attention.wq")`` and
887 ``fsdp_mark_pattern("layers.*")`` match.
889 A parameter directly on the root (``weight``) has no module ancestor;
890 its own FQN is still tested first so it can be marked by its name
891 (``fsdp_mark("weight")``) or by a ``*`` pattern.
892 """
893 if self._parallel_plan.is_marked_for_fsdp(param_fqn):
894 return True
896 parts = param_fqn.split(".")
897 for i in range(len(parts) - 1, 0, -1):
898 module_fqn = ".".join(parts[:i])
899 if self._parallel_plan.is_marked_for_fsdp(module_fqn):
900 return True
901 return False
903 def _get_parent_module_fqn(self, param_fqn: str) -> str:
904 """
905 Get parent module FQN from parameter FQN.
907 e.g., "layers.0.attention.wq.weight" -> "layers.0.attention.wq"
908 """
909 parts = param_fqn.split(".")
910 return ".".join(parts[:-1]) if len(parts) > 1 else param_fqn
912 def _insert_all_gather_for_params(
913 self,
914 graph_module: fx.GraphModule,
915 param_nodes: List[fx.Node],
916 pass_config: PassConfig,
917 ) -> fx.GraphModule:
918 """
919 Insert AllGather for each FSDP parameter placeholder.
921 Parameter state: Shard -> Replicate. All subsequent uses of the
922 placeholder are rewired to the unsharded (replicated) tensor, so the
923 computation body keeps operating on full parameters while the graph
924 input stays sharded. The gather is placed after the placeholder, or
925 sunk to just before the first consumer when reshard is enabled.
927 When reshard is enabled the gather is **sunk** to just before the
928 parameter's first consumer instead of sitting next to the placeholder.
929 All state placeholders lead the graph, so an unsunk gather would
930 materialize every full parameter before any compute runs, leaving the
931 reshard free step nothing to reclaim; sinking is what makes peak
932 memory track the forward working set.
933 """
934 graph = graph_module.graph
936 for param_node in param_nodes:
937 if param_node.name in self._processed_params:
938 continue
940 sink = pass_config.fsdp_reshard_after_forward
941 if sink and not param_node.users:
942 # Dead parameter: a gather would materialize a full replica
943 # nobody reads, and reshard would never free it. Leave the
944 # placeholder sharded.
945 self._processed_params.add(param_node.name)
946 continue
947 anchor = self._first_user(param_node) if sink else None
949 group_name, degree = self._param_group(
950 param_node.meta.get("param_name", "")
951 )
953 # AllGather + immediate wait for correctness; FsdpOverlapPass may
954 # sink the wait past independent compute later.
955 with _insertion_point(graph, anchor, param_node):
956 ag_node = graph.call_function(
957 _c10d_functional.all_gather_into_tensor,
958 args=(param_node, degree, group_name),
959 )
960 ag_node.meta["comm_type"] = "fsdp_all_gather"
961 ag_node.meta["comm_group"] = group_name
962 ag_node.meta["param_node"] = param_node.name
963 ag_node.meta["param_name"] = param_node.meta.get("param_name")
964 ag_node.meta["state_idx"] = param_node.meta.get("state_idx")
965 ag_node.meta["fsdp_degree"] = degree
966 ag_node.meta["fsdp_unsharded_nbytes"] = param_node.meta.get(
967 "fsdp_unsharded_nbytes"
968 )
970 # Insert the wait in its own block: a wait placed in the same block
971 # as the gather would land *before* the gather it depends on (the
972 # first insert inside ``inserting_after(X)`` lands right after
973 # ``X``), and codegen would emit
974 # ``wait_tensor(all_gather_into_tensor)`` before that variable is
975 # assigned.
976 with graph.inserting_after(ag_node):
977 wait_node = graph.call_function(
978 _c10d_functional.wait_tensor,
979 args=(ag_node,),
980 )
981 wait_node.meta["wait_for"] = ag_node.name
983 param_node.meta["fsdp_sharded"] = True
984 param_node.meta["fsdp_ag_node"] = ag_node.name
985 self._unsharded_params[param_node.name] = _UnshardedParam(
986 param_node=param_node, ag_node=ag_node, wait_node=wait_node
987 )
989 # Pad-sharded parameter: the gather produced the padded full tensor
990 # ``[n_pad, ...]``; narrow it back to the real ``[n, ...]`` and route
991 # every consumer to that view. Use the ``aten`` overload (not the
992 # ``torch.narrow`` builtin) so its schema's ``alias_info`` marks it a
993 # view: ``_propagates_alias`` keys off that schema, so the reshard
994 # alias walk (free / rematerialize) tracks the narrow and does not
995 # free the padded gather out from under its forward consumers.
996 replacement = wait_node
997 pad_info = self._padded_params.get(param_node.meta.get("param_name"))
998 if pad_info is not None:
999 with graph.inserting_after(wait_node):
1000 replacement = graph.call_function(
1001 torch.ops.aten.narrow.default,
1002 args=(wait_node, 0, 0, pad_info.n),
1003 )
1004 replacement.meta["param_name"] = param_node.meta.get("param_name")
1005 replacement.meta["state_idx"] = param_node.meta.get("state_idx")
1007 for user in list(param_node.users.keys()):
1008 if user not in (ag_node, wait_node):
1009 user.replace_input_with(param_node, replacement)
1011 self._processed_params.add(param_node.name)
1013 return graph_module
1015 @staticmethod
1016 def _first_user(param_node: fx.Node) -> Optional[fx.Node]:
1017 """Earliest (in graph order) consumer of ``param_node``.
1019 Placeholders are all at the top; their consumers are spread through
1020 the forward. Returning the earliest consumer lets the gather sink
1021 next to the first real use. ``None`` when the placeholder is unused
1022 (then the gather stays next to the placeholder).
1023 """
1024 users = [u for u in param_node.users if u.op != "output"]
1025 if not users:
1026 return None
1027 forward_users = [u for u in users if not u.meta.get("autograd_backward", False)]
1028 candidates = forward_users or users
1029 order = {n: i for i, n in enumerate(param_node.graph.nodes)}
1030 return min(candidates, key=lambda u: order.get(u, len(order)))
1032 def _insert_grad_reduction( # pylint: disable=too-many-locals
1033 self,
1034 graph_module: fx.GraphModule,
1035 sharded_param_indices: Set[int],
1036 replicated_param_indices: Set[int],
1037 state_fqns: List[str],
1038 num_state_inputs: int,
1039 state_is_param: Optional[List[bool]] = None,
1040 model: Optional[nn.Module] = None,
1041 ) -> fx.GraphModule:
1042 """Insert gradient reduction on the gradient outputs of DP parameters.
1044 The joint graph returns ``[loss, grad0, grad1, ...]`` from the fwd+bwd
1045 function; gradient ``i`` (output index ``i+1``) corresponds to the
1046 ``i``-th trainable parameter (``state_fqns`` order, buffers and frozen
1047 params skipped, mirrored by ``_build_trainable_state_indices``). For a
1048 parameter in ``sharded_param_indices`` (a real shard axis) the mode
1049 decides the collective chain: ``fsdp`` reduce_scatters on the ``fsdp``
1050 group (Replicate -> Shard); ``ddp`` all_reduces on ``dp_replicate``;
1051 ``hsdp`` reduce_scatters on ``fsdp`` then all_reduces on
1052 ``dp_replicate`` (shard axis first, so the cross-axis message is
1053 already sharded). Pad-sharded params (``_padded_params``) zero-pad their
1054 ``[n, ...]`` gradient back to ``[n_pad, ...]`` first so the local
1055 gradient matches the padded shard.
1057 For a parameter in ``replicated_param_indices`` (a marked 0-dim scalar,
1058 or a partial-plan leftover that fell back to DDP) no shard axis exists,
1059 so the gradient is reduced with all_reduce instead. Ordinary parameters
1060 average over the DP axes; source-layout parameters use their explicit
1061 reduction policy. Gradients outside the DP modules stay full.
1062 """
1063 graph = graph_module.graph
1065 output_node = next((n for n in graph.nodes if n.op == "output"), None)
1066 if output_node is None:
1067 return graph_module
1069 returned = output_node.args[0]
1070 if not isinstance(returned, (list, tuple)):
1071 # Model returned a single value (e.g. just loss) -> no grads.
1072 return graph_module
1074 trainable_state_indices = self._build_trainable_state_indices(
1075 state_fqns, num_state_inputs, state_is_param, model
1076 )
1078 new_returned = list(returned)
1079 # ``train_fn`` may append the traced ``loss_dict`` outputs AFTER the
1080 # gradients (see ``trace_model_graph``); those are not gradients and
1081 # must stay outside the reduce_scatter range.
1082 num_loss_outputs = len(getattr(graph_module, "loss_dict_keys", []))
1083 num_grads = len(new_returned) - 1 - num_loss_outputs # index 0 is the loss
1084 if num_grads != len(trainable_state_indices):
1085 raise ValueError(
1086 f"Gradient count ({num_grads}) does not match trainable "
1087 f"parameter count ({len(trainable_state_indices)}). The "
1088 f"traced graph and the live model disagree on which "
1089 f"parameters are trainable; refusing to insert gradient "
1090 f"reduction to avoid silent gradient/state misalignment."
1091 )
1093 # Index 0 is the loss; gradients start at index 1. Gradient i+1
1094 # corresponds to trainable parameter i, whose state_idx is
1095 # trainable_state_indices[i]. The trailing loss_dict outputs are
1096 # named loss values, not gradients.
1097 for i in range(1, len(new_returned) - num_loss_outputs):
1098 grad_node = new_returned[i]
1099 if not isinstance(grad_node, fx.Node):
1100 continue
1102 state_idx = trainable_state_indices[i - 1]
1103 fqn = state_fqns[state_idx]
1104 group_name, degree = self._param_group(fqn)
1105 if state_idx in sharded_param_indices:
1106 replicated = False
1107 pad_info = self._padded_params.get(fqn)
1108 elif state_idx in replicated_param_indices:
1109 replicated = True
1110 pad_info = None
1111 elif degree == 1 and fqn in self._gradient_reductions:
1112 replicated = True
1113 pad_info = None
1114 else:
1115 continue
1117 reduction_info = _GradReductionInfo(
1118 group_name, degree, self._gradient_reductions.get(fqn, _DEFAULT_GRADIENT_REDUCTION)
1119 )
1120 with graph.inserting_before(output_node):
1121 new_returned[i] = self._build_grad_reduction_nodes(
1122 graph,
1123 grad_node,
1124 reduction_info,
1125 replicated=replicated,
1126 pad_info=pad_info,
1127 )
1128 tag_gradient_identity(new_returned[i], fqn, state_idx)
1129 source_group = self._source_gradient_groups.get(fqn)
1130 if source_group is not None:
1131 new_returned[i] = self._insert_source_gradient_reduction(
1132 graph, new_returned[i], source_group
1133 )
1134 tag_gradient_identity(new_returned[i], state_fqns[state_idx], state_idx)
1136 output_node.args = (type(returned)(new_returned),) + tuple(output_node.args[1:])
1138 return graph_module
1140 @staticmethod
1141 def _insert_source_gradient_reduction(
1142 graph: fx.Graph, gradient: fx.Node, group_name: str
1143 ) -> fx.Node:
1144 """Sum replicated source-layout contributions after completed DP reduction."""
1145 contiguous = graph.call_method("contiguous", args=(gradient,))
1146 reduced = graph.call_function(
1147 _c10d_functional.all_reduce, args=(contiguous, "sum", group_name)
1148 )
1149 reduced.meta.update(comm_type="source_all_reduce", comm_group=group_name, autograd_backward=True)
1150 return FSDPPass._insert_wait(graph, reduced)
1152 def _build_grad_reduction_nodes(
1153 self,
1154 graph: fx.Graph,
1155 grad_node: fx.Node,
1156 info: _GradReductionInfo,
1157 replicated: bool = False,
1158 pad_info: Optional[_PaddedParam] = None,
1159 ) -> fx.Node:
1160 """Build and return the waited gradient-reduction node for ``grad_node``.
1162 The returned (wait) node replaces ``grad_node`` in the graph output.
1163 Emits the mode-specific chain (see ``_insert_grad_reduction``). When
1164 ``pad_info`` is given the parameter is pad-sharded: the real ``[n, ...]``
1165 gradient is zero-padded on dim 0 to ``[n_pad, ...]`` before the
1166 reduce_scatter, so the scattered local gradient matches the padded live
1167 shard (mirrors eager ``fully_shard``'s padded grad buffer). When
1168 ``replicated`` is True the parameter is a 0-dim scalar that could not be
1169 sharded, so its gradient is reduced over the shard group with
1170 ``all_reduce`` instead of reduce-scattered — the graph and the live
1171 model both keep the parameter full, so a reduce_scatter would change the
1172 shape and desync them.
1174 Every collective is waited on before its result is consumed. A chained
1175 collective runs on a different process-group stream (HSDP's
1176 reduce_scatter is on ``fsdp`` while the all_reduce is on
1177 ``dp_replicate``), so without the intermediate wait the all_reduce could
1178 read a partially written shard. This mirrors eager ``fully_shard``,
1179 which orders the all_reduce stream after the reduce_scatter stream.
1180 """
1181 group_name, degree = info.group_name, info.degree
1182 reduced = grad_node
1183 if info.reduction.scale is not None:
1184 reduced = graph.call_function(torch.mul, args=(reduced, info.reduction.scale))
1185 if degree == 1 and self._dp_mode == "fsdp":
1186 return reduced
1187 if pad_info is not None:
1188 # constant_pad_nd's pad tuple runs last-dim-first: pad dim 0 on the
1189 # right by ``dim0_pad`` and leave every other dim untouched.
1190 pad_tuple = tuple([0, 0] * (pad_info.ndim - 1) + [0, pad_info.dim0_pad])
1191 reduced = graph.call_function(
1192 torch.constant_pad_nd, args=(reduced, pad_tuple)
1193 )
1194 reduced.meta["comm_type"] = "fsdp_grad_pad"
1195 value = grad_node.meta.get("val")
1196 if value is not None and hasattr(value, "shape"):
1197 reduced.meta["fsdp_padded_input_spec"] = {
1198 "shape": (pad_info.n_pad, *value.shape[1:]),
1199 "dtype": value.dtype,
1200 "device": str(value.device),
1201 }
1202 if replicated:
1203 # Replicated parameter (a 0-dim scalar, or a partial-plan leftover
1204 # that fell back to DDP): average the full gradient over the shard
1205 # group so every replica takes the same optimizer step. DDP mode has
1206 # no shard group — the replicate all_reduce below handles it.
1207 if self._dp_mode in ("fsdp", "hsdp"):
1208 ar_node = graph.call_function(
1209 _c10d_functional.all_reduce,
1210 args=(reduced, info.reduction.reduce_op, group_name),
1211 )
1212 ar_node.meta["comm_type"] = "fsdp_all_reduce"
1213 ar_node.meta["comm_group"] = group_name
1214 ar_node.meta["fsdp_degree"] = degree
1215 reduced = ar_node
1216 # HSDP chains the replicate all_reduce (a different process-group
1217 # stream) after the shard all_reduce, so wait first.
1218 if self._dp_mode == "hsdp":
1219 reduced = self._insert_wait(graph, reduced)
1220 elif self._dp_mode in ("fsdp", "hsdp"):
1221 # Matmul backward can return a transposed view on NPU, including dense router weights.
1222 reduced = graph.call_method("contiguous", args=(reduced,))
1223 rs_node = graph.call_function(
1224 _c10d_functional.reduce_scatter_tensor,
1225 args=(
1226 reduced,
1227 info.reduction.reduce_op,
1228 degree,
1229 group_name,
1230 ),
1231 )
1232 rs_node.meta["comm_type"] = "fsdp_reduce_scatter"
1233 rs_node.meta["comm_group"] = group_name
1234 rs_node.meta["fsdp_degree"] = degree
1235 rs_node.meta["autograd_backward"] = True
1236 reduced = rs_node
1237 # HSDP chains all_reduce (a different process-group stream) after
1238 # the reduce_scatter, so wait first; plain FSDP waits on the final
1239 # node only.
1240 if self._dp_mode == "hsdp":
1241 reduced = self._insert_wait(graph, reduced)
1243 if self._dp_mode in ("ddp", "hsdp"):
1244 ar_node = graph.call_function(
1245 _c10d_functional.all_reduce,
1246 args=(reduced, info.reduction.reduce_op, self._dp_replicate_group_name),
1247 )
1248 ar_node.meta["comm_type"] = "dp_all_reduce"
1249 ar_node.meta["comm_group"] = self._dp_replicate_group_name
1250 ar_node.meta["dp_replicate_degree"] = self._dp_replicate_degree
1251 ar_node.meta["autograd_backward"] = True
1252 reduced = ar_node
1254 return self._insert_wait(graph, reduced)
1256 @staticmethod
1257 def _insert_wait(graph: fx.Graph, tensor_node: fx.Node) -> fx.Node:
1258 """Append ``wait_tensor(tensor_node)`` and return the waited node."""
1259 wait_node = graph.call_function(
1260 _c10d_functional.wait_tensor,
1261 args=(tensor_node,),
1262 )
1263 wait_node.meta["wait_for"] = tensor_node.name
1264 wait_node.meta["autograd_backward"] = True
1265 return wait_node
1267 def _build_trainable_state_indices(
1268 self,
1269 state_fqns: List[str],
1270 num_state_inputs: int,
1271 state_is_param: Optional[List[bool]],
1272 model: Optional[nn.Module],
1273 ) -> List[int]:
1274 """
1275 Return state indices of trainable parameters, in ``state_fqns`` order.
1277 Mirrors the tracer's ``params`` list construction (skip buffers via
1278 ``state_is_param``, skip frozen params via ``requires_grad``) so
1279 gradient ``i`` maps to ``trainable_state_indices[i]``. Reproducing this
1280 filter here is what keeps reduce_scatter aligned for models with
1281 buffers or frozen params.
1282 """
1283 if model is None:
1284 # Without the model we cannot inspect requires_grad; fall back to
1285 # the all-trainable assumption. The grad-count check in the caller
1286 # raises if that assumption is wrong, so misalignment is not silent.
1287 return list(range(num_state_inputs))
1289 param_lookup = dict(model.named_parameters(remove_duplicate=False))
1290 trainable: List[int] = []
1291 for idx in range(num_state_inputs):
1292 if state_is_param is not None and not state_is_param[idx]:
1293 continue
1294 param = param_lookup.get(state_fqns[idx])
1295 if param is not None and param.requires_grad:
1296 trainable.append(idx)
1297 return trainable
1299 def _insert_reshard_logic(
1300 self,
1301 graph_module: fx.GraphModule,
1302 param_nodes: List[fx.Node],
1303 pass_config: PassConfig,
1304 ) -> fx.GraphModule:
1305 """Reshard each FSDP parameter after forward (FSDP release).
1307 For every FSDP parameter whose all_gather was inserted by
1308 ``_insert_all_gather_for_params``:
1310 1. Free the replicated full parameter once its last forward reader —
1311 including forward views aliasing it — has run.
1312 2. If the backward still needs the parameter (directly or through a
1313 saved forward view), insert a fresh all_gather just before the
1314 first backward consumer and rematerialize the saved view chain, so
1315 the backward computes the exact same values from the re-gathered
1316 (unsharded) parameter.
1317 3. Free that backward re-gather once its last backward reader runs.
1319 Steps 1-3 keep only the forward working set resident during forward and
1320 only the active layer's parameter resident during backward, instead of
1321 every replicated parameter for the whole joint graph.
1323 No-op when ``fsdp_reshard_after_forward`` is disabled.
1324 """
1325 if not pass_config.fsdp_reshard_after_forward:
1326 return graph_module
1328 for param_node in param_nodes:
1329 record = self._unsharded_params.get(param_node.name)
1330 if record is None:
1331 continue
1332 self._reshard_param(graph_module.graph, record)
1333 return graph_module
1335 def _reshard_param(self, graph: fx.Graph, record: _UnshardedParam) -> None:
1336 """Reshard one FSDP parameter (free + backward re-gather/remat)."""
1337 wait = record.wait_node
1338 alias_nodes = self._alias_descendants(wait)
1339 order = {n: i for i, n in enumerate(graph.nodes)}
1341 # Alias (view) descendants the backward reads must be rematerialized
1342 # from the re-gathered parameter; the unsharded tensor itself may be
1343 # consumed directly by the backward in some models.
1344 needed = [
1345 n for n in [wait, *alias_nodes] if self._feeds_backward(n, alias_nodes)
1346 ]
1348 # Remat plan first: the forward free is only safe once every backward
1349 # consumer has been rewired onto fresh (re-gathered) storage. If the
1350 # view chain cannot be cloned, keep the previous resident-for-the-graph
1351 # behavior — a missed free costs memory, a misordered one is unsound.
1352 if self._rematerialize_for_backward(graph, record, needed, order):
1353 # 1. Free after the last forward data read of the unsharded storage.
1354 self._insert_reshard_free(graph, wait, record, order, backward=False)
1356 def _rematerialize_for_backward(
1357 self,
1358 graph: fx.Graph,
1359 record: _UnshardedParam,
1360 needed: List[fx.Node],
1361 order: Dict[fx.Node, int],
1362 ) -> bool:
1363 """Re-gather ahead of the first backward consumer and remat views.
1365 The saved forward view chain is cloned against the re-gathered
1366 parameter so the backward reads identical values from a fresh,
1367 correctly-sized tensor.
1369 Returns:
1370 ``True`` when the free may proceed — no backward consumer exists,
1371 or every one of them was rewired onto rematerialized storage.
1372 ``False`` when a needed view cannot be cloned; the caller must
1373 keep the unsharded parameter resident.
1374 """
1375 backward_users = [
1376 u for n in needed for u in n.users if u.meta.get("autograd_backward", False)
1377 ]
1378 if not backward_users:
1379 return True
1381 wait = record.wait_node
1382 uncloneable = [n for n in needed if n is not wait and not self._cloneable(n)]
1383 if uncloneable:
1384 _LOG.warning(
1385 "FSDP reshard: cannot rematerialize %s for '%s'; keeping the "
1386 "unsharded parameter resident",
1387 [n.format_node() for n in uncloneable],
1388 record.param_node.meta.get("param_name"),
1389 )
1390 return False
1392 anchor = min(backward_users, key=lambda u: order.get(u, len(order)))
1393 wait2 = self._insert_backward_gather(graph, record, anchor)
1395 recreate: Dict[fx.Node, fx.Node] = {wait: wait2}
1396 last_inserted = wait2
1397 for node in sorted(
1398 (n for n in needed if n is not wait), key=lambda n: order[n]
1399 ):
1400 last_inserted = self._clone_with_remap(
1401 graph, node, recreate, wait, wait2, last_inserted
1402 )
1404 self._rewire_backward_users(recreate)
1405 self._insert_reshard_free(graph, wait2, record, order, backward=True)
1406 return True
1408 @staticmethod
1409 def _cloneable(node: fx.Node) -> bool:
1410 """Whether ``node`` can be cloned as a rematerialized view op."""
1411 return node.op == "call_function" and callable(node.target)
1413 def _insert_backward_gather(
1414 self, graph: fx.Graph, record: _UnshardedParam, anchor: fx.Node
1415 ) -> fx.Node:
1416 """Insert the backward all_gather + wait pair before ``anchor``."""
1417 group_name, degree = self._param_group(
1418 record.param_node.meta.get("param_name", "")
1419 )
1420 with graph.inserting_before(anchor):
1421 ag_node = graph.call_function(
1422 _c10d_functional.all_gather_into_tensor,
1423 args=(record.param_node, degree, group_name),
1424 )
1425 ag_node.meta["comm_type"] = "fsdp_all_gather_backward"
1426 ag_node.meta["comm_group"] = group_name
1427 ag_node.meta["param_node"] = record.param_node.name
1428 ag_node.meta["param_name"] = record.param_node.meta.get("param_name")
1429 ag_node.meta["state_idx"] = record.param_node.meta.get("state_idx")
1430 ag_node.meta["fsdp_degree"] = degree
1431 ag_node.meta["fsdp_unsharded_nbytes"] = record.param_node.meta.get(
1432 "fsdp_unsharded_nbytes"
1433 )
1434 # This gather lives in the backward half; tag it so downstream
1435 # passes (e.g. PpPass phase classification) do not mistake it and
1436 # its remat clones for forward nodes — their only arg is a forward
1437 # state placeholder, which would otherwise imply the forward phase.
1438 ag_node.meta["autograd_backward"] = True
1439 with graph.inserting_before(anchor):
1440 wait_node = graph.call_function(
1441 _c10d_functional.wait_tensor, args=(ag_node,)
1442 )
1443 wait_node.meta["wait_for"] = ag_node.name
1444 wait_node.meta["autograd_backward"] = True
1445 return wait_node
1447 def _clone_with_remap(
1448 self,
1449 graph: fx.Graph,
1450 node: fx.Node,
1451 recreate: Dict[fx.Node, fx.Node],
1452 wait: fx.Node,
1453 wait2: fx.Node,
1454 insert_after: fx.Node,
1455 ) -> fx.Node:
1456 """Clone ``node`` after ``insert_after``, remapping args to re-gathers."""
1457 args = tuple(self._remap_arg(a, recreate, wait, wait2) for a in node.args)
1458 kwargs = {
1459 k: self._remap_arg(v, recreate, wait, wait2) for k, v in node.kwargs.items()
1460 }
1461 with graph.inserting_after(insert_after):
1462 clone = graph.call_function(node.target, args=args, kwargs=kwargs)
1463 self._copy_recreated_meta(node, clone)
1464 # Rematerialized views serve the backward half, so they carry the
1465 # backward phase (see ``_insert_backward_gather``).
1466 clone.meta["autograd_backward"] = True
1467 recreate[node] = clone
1468 return clone
1470 @staticmethod
1471 def _rewire_backward_users(recreate: Dict[fx.Node, fx.Node]) -> None:
1472 """Point every backward consumer of an original node at its clone."""
1473 for original, clone in recreate.items():
1474 for user in list(original.users):
1475 if user.meta.get("autograd_backward", False):
1476 user.replace_input_with(original, clone)
1478 def _insert_reshard_free(
1479 self,
1480 graph: fx.Graph,
1481 tensor: fx.Node,
1482 record: _UnshardedParam,
1483 order: Dict[fx.Node, int],
1484 backward: bool,
1485 ) -> None:
1486 """Free ``tensor``'s storage after its last reader in the phase."""
1487 last_reader = self._last_data_reader(tensor, order, backward=backward)
1488 if last_reader is None:
1489 return
1490 comm_type = "fsdp_reshard_free_backward" if backward else "fsdp_reshard_free"
1491 with graph.inserting_after(last_reader):
1492 free_node = graph.call_function(_free_tensor_storage, args=(tensor,))
1493 free_node.meta["comm_type"] = comm_type
1494 free_node.meta["param_name"] = record.param_node.meta.get("param_name")
1495 if backward:
1496 free_node.meta["autograd_backward"] = True
1498 def _alias_descendants(self, root: fx.Node) -> Set[fx.Node]:
1499 """Forward view descendants of ``root`` (nodes aliasing its storage).
1501 Only aliasing ops propagate: an op that allocates a new tensor
1502 (``mm``, ``addmm``, ``relu``, ...) copies the data out, so freeing
1503 ``root``'s storage does not invalidate it and it need not be
1504 rematerialized. Backward nodes terminate the walk.
1505 """
1506 alias: Set[fx.Node] = set()
1507 stack = [root]
1508 while stack:
1509 current = stack.pop()
1510 for user in current.users:
1511 if user is root or user in alias:
1512 continue
1513 if user.meta.get("autograd_backward", False):
1514 continue
1515 if current not in set(_iter_arg_nodes(user)):
1516 continue
1517 if not self._propagates_alias(user):
1518 continue
1519 alias.add(user)
1520 stack.append(user)
1521 return alias
1523 @staticmethod
1524 def _propagates_alias(node: fx.Node) -> bool:
1525 """Whether ``node``'s output aliases its already-aliasing input.
1527 Two forms propagate. Schema-declared view returns (``aten.t``,
1528 ``aten.slice``, ...), and ``operator.getitem`` picking one output of a
1529 multi-output view op such as ``aten.chunk`` / ``aten.split`` /
1530 ``aten.unbind`` — getitem itself has no schema, so the alias must be
1531 inherited from the indexed producer (which the caller has already
1532 established is in the alias closure).
1533 """
1534 if node.op != "call_function":
1535 return False
1536 if node.target is _GETITEM:
1537 return True
1538 schema = getattr(node.target, "_schema", None)
1539 if schema is None:
1540 return False
1541 return any(ret.alias_info is not None for ret in schema.returns)
1543 def _feeds_backward(self, node: fx.Node, alias_nodes: Set[fx.Node]) -> bool:
1544 """Whether ``node`` (or an aliasing descendant) is read by backward."""
1545 memo: Dict[fx.Node, bool] = {}
1547 def visit(current: fx.Node) -> bool:
1548 """Whether ``current`` or an aliasing descendant reaches backward.
1550 Args:
1551 current: Node to test.
1553 Returns:
1554 True when a backward node consumes the node or its aliases.
1555 """
1556 if current in memo:
1557 return memo[current]
1558 memo[current] = False # cycle guard (aliasing is a DAG)
1559 result = False
1560 for user in current.users:
1561 if user.meta.get("autograd_backward", False):
1562 result = True
1563 break
1564 if user in alias_nodes and visit(user):
1565 result = True
1566 break
1567 memo[current] = result
1568 return result
1570 return visit(node)
1572 def _last_data_reader(
1573 self, root: fx.Node, order: Dict[fx.Node, int], backward: bool
1574 ) -> Optional[fx.Node]:
1575 """Latest, in graph order, op that reads ``root``'s storage data.
1577 Walks the alias (view) chain so a free is not placed after a view
1578 *creation* but before the non-view op that actually reads it: freeing
1579 between ``t_2 = t(w)`` and ``mm(grad, t_2)`` would invalidate ``t_2``.
1580 Multi-output views reach their consumers through ``getitem``, so the
1581 walk continues through those picks instead of mistaking them for the
1582 final reader. ``backward=False`` picks the last forward read (free
1583 point for the unsharded parameter); ``backward=True`` the last
1584 backward read (free point for the backward re-gather).
1585 """
1586 last: Optional[fx.Node] = None
1587 last_pos = -1
1588 seen: Set[fx.Node] = set()
1589 stack = [root]
1590 while stack:
1591 current = stack.pop()
1592 if current in seen:
1593 continue
1594 seen.add(current)
1595 for user in current.users:
1596 if user.op == "output":
1597 continue
1598 if self._propagates_alias(user) and current in set(
1599 _iter_arg_nodes(user)
1600 ):
1601 stack.append(user)
1602 continue
1603 if bool(user.meta.get("autograd_backward", False)) != backward:
1604 continue
1605 pos = order.get(user, -1)
1606 if pos > last_pos:
1607 last_pos = pos
1608 last = user
1609 return last
1611 @classmethod
1612 def _remap_arg(
1613 cls,
1614 arg: Any,
1615 recreate: Dict[fx.Node, fx.Node],
1616 wait: fx.Node,
1617 wait2: fx.Node,
1618 ) -> Any:
1619 """Rewrite an op arg for a rematerialized clone (nested containers)."""
1620 if arg is wait:
1621 return wait2
1622 if isinstance(arg, fx.Node):
1623 return recreate.get(arg, arg)
1624 if isinstance(arg, tuple):
1625 return tuple(cls._remap_arg(a, recreate, wait, wait2) for a in arg)
1626 if isinstance(arg, list):
1627 return [cls._remap_arg(a, recreate, wait, wait2) for a in arg]
1628 if isinstance(arg, dict):
1629 return {k: cls._remap_arg(v, recreate, wait, wait2) for k, v in arg.items()}
1630 return arg
1632 @staticmethod
1633 def _copy_recreated_meta(original: fx.Node, clone: fx.Node) -> None:
1634 """Copy the module/stack/shape metadata a remat clone needs."""
1635 for key in (
1636 "nn_module_stack",
1637 "custom",
1638 "stack_trace",
1639 "tensor_meta",
1640 "val",
1641 ):
1642 value = original.meta.get(key)
1643 if value is not None:
1644 clone.meta[key] = value