Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / passes / parallel / fsdp_pass.py: 93%
504 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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.distributed.distributed_c10d import _resolve_process_group
64from torch.fx.node import _side_effectful_functions
65from torch.ops import _c10d_functional
67from ...pass_config import PassConfig, normalize_dp_mode
68from ..base import GraphPass
69from ...graph_parallel_plan import GraphParallelPlan
71_LOG = logging.getLogger(__name__)
73_GETITEM = operator.getitem
75# Gradient reduction op for every DP axis. Mirrors eager ``fully_shard``'s
76# default (``core/fully_shard/hsdp_state.py``: ``set_reduce_op_type("avg")``):
77# the backward produces the full (replicated) gradient on every rank, so each
78# axis must be AVERAGED (not summed) to keep the step gradient at the
79# DP-world-mean scale. HSDP applies it on both axes, yielding the mean over the
80# shard times the mean over the replicate axis.
81_GRAD_REDUCE_OP = "avg"
84def _free_tensor_storage(tensor: torch.Tensor) -> None:
85 """Release ``tensor``'s storage in place (FSDP reshard free step).
87 Inserted as a side-effecting FX node between an unsharded parameter's last
88 forward reader and its backward re-gather, so the replicated full
89 parameter (and every view aliasing it) does not stay resident for the
90 whole joint graph. Returns ``None``; the node exists only for its effect.
91 Mirrors ``fully_shard``'s ``free_unsharded_param`` storage release.
92 """
93 if isinstance(tensor, torch.Tensor) and tensor.untyped_storage().size() > 0:
94 tensor.untyped_storage().resize_(0)
97# FX dead-code elimination drops call_function nodes it judges pure; our free
98# op is pure to FX but has a real side effect. Register it so a later
99# ``eliminate_dead_code`` (a future pass, or a user) cannot delete the free.
100_side_effectful_functions.add(_free_tensor_storage)
103@dataclass
104class _UnshardedParam:
105 """Forward all_gather that materializes one FSDP parameter.
107 ``param_node`` is the sharded placeholder (graph input); ``ag_node`` /
108 ``wait_node`` produce and expose the replicated (unsharded) full parameter
109 the forward — and, after a reshard re-gather, the backward — read.
110 """
112 param_node: fx.Node
113 ag_node: fx.Node
114 wait_node: fx.Node
117@dataclass
118class _DpParamClassification:
119 """How DP handles each trainable parameter's gradient.
121 Produced by ``_classify_dp_params``:
123 - ``sharded_nodes``: marked, trainable parameter placeholders that get an
124 AllGather (Shard -> Replicate) and a reduce_scatter gradient, and are
125 physically sharded in the live model. Dim-0-divisible params are chunked
126 directly; non-divisible ones are pad-sharded (see ``_PaddedParam``).
127 - ``replicated_indices``: state indices of parameters kept replicated (never
128 chunked) whose gradient is averaged over the DP group with ``all_reduce``
129 instead of reduce-scattered — a marked 0-dim scalar (cannot be sharded), or
130 a partial-plan leftover that falls back to DDP.
131 - ``unmarked_trainable``: ``(state_idx, fqn)`` for the partial-plan leftovers
132 (a subset of ``replicated_indices``), kept only so the caller can log that
133 the plan did not cover them.
134 """
136 sharded_nodes: List[fx.Node]
137 replicated_indices: Set[int]
138 unmarked_trainable: List[Tuple[int, str]]
141@dataclass
142class _PaddedParam:
143 """A marked parameter that cannot be sharded evenly, pad-sharded instead.
145 Mirrors eager ``fully_shard`` (``core/fully_shard/hsdp_param.py``): dim 0 of
146 ``n`` elements is zero-padded to ``n_pad = ceil(n / degree) * degree`` so
147 every rank owns a uniform ``n_pad / degree``-element shard. The graph
148 narrows the gathered ``[n_pad, ...]`` back to ``[n, ...]`` for the compute,
149 and the backward pads the ``[n, ...]`` gradient back to ``[n_pad, ...]``
150 before reduce_scatter, so the returned local gradient matches the padded
151 shard (the padding region gets an all-zero gradient, hence never moves).
152 """
154 n: int
155 n_pad: int
156 dim0_pad: int
157 ndim: int
160def _iter_arg_nodes(node: fx.Node) -> Iterator[fx.Node]:
161 """Yield every ``fx.Node`` in ``node.args`` / ``node.kwargs`` (nested)."""
162 stack: List[Any] = list(node.args) + list(node.kwargs.values())
163 while stack:
164 item = stack.pop()
165 if isinstance(item, fx.Node):
166 yield item
167 elif isinstance(item, (tuple, list)):
168 stack.extend(item)
169 elif isinstance(item, dict):
170 stack.extend(item.values())
173@contextmanager
174def _insertion_point(
175 graph: fx.Graph, anchor: Optional[fx.Node], fallback: fx.Node
176) -> Iterator[None]:
177 """Yield a graph insertion context.
179 Owns the ``anchor``-vs-``fallback`` choice so the gather-insertion body
180 stays a single code path: when ``anchor`` is given, new nodes land right
181 before it; otherwise they land right after ``fallback``.
182 """
183 if anchor is None:
184 with graph.inserting_after(fallback):
185 yield
186 else:
187 with graph.inserting_before(anchor):
188 yield
191class FSDPPass(GraphPass):
192 """
193 FSDP Partitioning Pass (Module-Level, static-input graph)
195 Responsibilities:
196 1. Only shard parameters in modules marked for FSDP wrapping
197 2. Insert AllGather after each FSDP parameter placeholder
198 (Shard -> Replicate)
199 3. Insert ReduceScatter after Backward (Replicate -> Shard)
200 4. Physically shard the live model's parameters in place so the
201 trainer / optimizer stay FSDP-agnostic
203 Key Difference from Old Implementation:
204 - Old: shard all parameters via get_attr nodes and physically rewrite the
205 GraphModule's stored tensors
206 - New: parameters are static graph inputs; the pass inserts collectives
207 into the graph and shards the live model's parameters in place
208 """
210 name = "fsdp_parallel"
212 def __init__(
213 self,
214 fsdp_group_name: Optional[str] = None,
215 parallel_plan: Optional[GraphParallelPlan] = None,
216 ) -> None:
217 """Initialize FSDP pass state.
219 Args:
220 fsdp_group_name: Process-group name registered for FSDP
221 collectives. Defaults to ``"fsdp"``. May be overridden
222 per-run via the ``fsdp_group_name`` kwarg in ``run``.
223 parallel_plan: Declarative plan identifying which modules to
224 shard. When ``None``, all parameters are sharded.
225 """
226 super().__init__()
227 self._fsdp_group_name = fsdp_group_name or "fsdp"
228 self._parallel_plan = parallel_plan
229 # Resolved from ``pass_config.fsdp_degree`` (or world_size) at
230 # ``run`` entry; left as ``None`` here so a stray access before
231 # ``run`` fails loudly instead of silently using a wrong default.
232 self._fsdp_degree: Optional[int] = None
233 # Data-parallel mode ("fsdp" | "ddp" | "hsdp") and the replicate-axis
234 # group/degree used by DDP / HSDP gradient reduction.
235 self._dp_mode = "fsdp"
236 self._dp_replicate_group_name = "dp_replicate"
237 self._dp_replicate_degree: Optional[int] = None
238 self._processed_params: Set[str] = set()
239 self._fsdp_modules: Set[str] = set()
240 # Per-param forward all_gather record (the unsharded parameter),
241 # populated by ``_insert_all_gather_for_params`` and consumed by the
242 # reshard step.
243 self._unsharded_params: Dict[str, _UnshardedParam] = {}
244 # Non-divisible marked params keyed by FQN; populated by
245 # ``_classify_dp_params`` and consumed by the all_gather / grad-reduce /
246 # live-shard steps so they are pad-sharded like eager ``fully_shard``.
247 self._padded_params: Dict[str, _PaddedParam] = {}
249 def run(
250 self,
251 graph_module: fx.GraphModule,
252 pass_config: PassConfig,
253 **kwargs: Any,
254 ) -> fx.GraphModule:
255 """Insert AllGather/ReduceScatter and shard the live model.
257 Args:
258 graph_module: Joint fwd+bwd FX graph from the tracer.
259 pass_config: Parallel configuration. ``fsdp_degree`` is
260 read directly; ``None`` falls back to ``world_size`` (the
261 FSDP-only path). A ``TypeError`` here means a non-Protocol
262 config was passed — fix at the caller, do not paper over.
263 **kwargs: Must include ``model`` (the live ``nn.Module``) and
264 ``fsdp_group_name`` / ``parallel_plan`` as needed.
266 Returns:
267 The transformed graph module.
268 """
269 if not dist.is_initialized() or dist.get_world_size() == 1:
270 _LOG.info("Skipped: distributed not initialized or world_size=1")
271 return graph_module
273 model = self._setup_dp_config(pass_config, kwargs)
275 # Identify DP parameter placeholders. The joint graph's
276 # parameters/buffers are static inputs (leading placeholders, not
277 # get_attr nodes): locate them by position via the state layout the
278 # tracer attached, then keep only those in DP-marked modules. For
279 # FSDP/HSDP a marked param whose leading dim is not divisible by the
280 # shard degree is pad-sharded (like eager fully_shard) so it is still
281 # sharded and gradient-reduced; scalars stay replicated. DDP has no
282 # shard axis so every marked parameter is reduced.
283 state_fqns = getattr(graph_module, "state_fqns", [])
284 num_state_inputs = getattr(graph_module, "num_state_inputs", 0)
285 # Param-vs-buffer flag per leading state input. Absent on graphs
286 # traced before this flag existed; default to all-params so old
287 # traces keep the previous (parameter-only) behaviour.
288 state_is_param = getattr(graph_module, "state_is_param", None)
290 self._reset_run_state()
291 classification = self._classify_dp_params(
292 graph_module, state_fqns, num_state_inputs, state_is_param, model
293 )
294 param_nodes = classification.sharded_nodes
296 _LOG.info(
297 "Identified %d sharded (incl. %d pad-sharded) + %d replicated DP "
298 "parameter nodes out of %d total state inputs",
299 len(param_nodes),
300 len(self._padded_params),
301 len(classification.replicated_indices),
302 num_state_inputs,
303 )
305 # A partial plan that leaves trainable params unmarked keeps their local
306 # gradient (never reduced) — a silent divergence source across DP ranks,
307 # so warn loudly before continuing.
308 if classification.unmarked_trainable:
309 self._warn_unmarked_trainable(classification.unmarked_trainable)
311 if not param_nodes and not classification.replicated_indices:
312 _LOG.warning(
313 "No DP parameters found, check GraphParallelPlan or model structure"
314 )
315 return graph_module
317 graph_module = self._apply_dp_transforms(
318 graph_module,
319 model,
320 classification,
321 pass_config,
322 state_fqns,
323 num_state_inputs,
324 state_is_param,
325 )
327 _LOG.info(
328 "Completed dp_mode=%s, processed %d parameters",
329 self._dp_mode,
330 len(param_nodes),
331 )
333 graph_module.recompile()
334 return graph_module
336 def _setup_dp_config(
337 self, pass_config: PassConfig, kwargs: Dict[str, Any]
338 ) -> nn.Module:
339 """Resolve run-time DP config, stash it on ``self``, return the model.
341 Reads ``dp_mode`` / ``fsdp_degree`` / ``dp_replicate_degree`` off
342 ``pass_config`` (the last two fall back to ``world_size`` when unset)
343 and the group names / ``parallel_plan`` / ``model`` off ``kwargs``.
345 Returns:
346 The live ``nn.Module`` to physically shard.
348 Raises:
349 ValueError: If ``model`` is missing from ``kwargs``.
350 """
351 self._dp_mode = normalize_dp_mode(getattr(pass_config, "dp_mode", "fsdp"))
352 # FSDP shard-axis group size: use the explicitly configured degree when
353 # present (required for TP+FSDP / HSDP, where the FSDP group is a
354 # proper sub-group of the world — using world_size would over-shard
355 # along the TP axis).
356 configured = pass_config.fsdp_degree
357 self._fsdp_degree = configured if configured else dist.get_world_size()
358 self._fsdp_group_name = kwargs.get("fsdp_group_name", self._fsdp_group_name)
359 self._dp_replicate_group_name = kwargs.get(
360 "dp_replicate_group_name", self._dp_replicate_group_name
361 )
362 # Replicate-axis degree used by DDP / HSDP gradient reduction.
363 replicate_degree = getattr(pass_config, "dp_replicate_degree", None)
364 if self._dp_mode == "ddp":
365 self._dp_replicate_degree = replicate_degree or dist.get_world_size()
366 elif self._dp_mode == "hsdp":
367 self._dp_replicate_degree = replicate_degree or (
368 dist.get_world_size() // self._fsdp_degree
369 )
370 else:
371 self._dp_replicate_degree = replicate_degree
372 self._parallel_plan = kwargs.get("parallel_plan", self._parallel_plan)
373 model = kwargs.get("model")
374 if model is None:
375 raise ValueError(
376 "FSDPPass requires the live model via kwargs (model=...) so it "
377 "can physically shard parameters; the trainer passes it through "
378 "in compile()"
379 )
381 _LOG.info(
382 "Running dp_mode=%s fsdp_degree=%s dp_replicate_degree=%s world_size=%s",
383 self._dp_mode,
384 self._fsdp_degree,
385 self._dp_replicate_degree,
386 dist.get_world_size(),
387 )
388 return model
390 def _reset_run_state(self) -> None:
391 """Clear per-run state so a reused pass does not carry names across graphs."""
392 self._processed_params = set()
393 self._fsdp_modules = set()
394 self._unsharded_params = {}
395 self._padded_params = {}
397 def _apply_dp_transforms(
398 self,
399 graph_module: fx.GraphModule,
400 model: nn.Module,
401 classification: _DpParamClassification,
402 pass_config: PassConfig,
403 state_fqns: List[str],
404 num_state_inputs: int,
405 state_is_param: Optional[List[bool]],
406 ) -> fx.GraphModule:
407 """Insert gather / grad-reduce / reshard and shard the live model.
409 DDP replicates parameters: no all_gather, no live-model sharding and no
410 reshard — only the gradient all-reduce path runs. FSDP/HSDP also
411 all_gather each parameter for the forward and pad/shard the live model
412 in place.
413 """
414 param_nodes = classification.sharded_nodes
415 if self._dp_mode != "ddp":
416 graph_module = self._insert_all_gather_for_params(
417 graph_module, param_nodes, pass_config
418 )
420 sharded_param_indices = frozenset(
421 node.meta["state_idx"] for node in param_nodes
422 )
423 graph_module = self._insert_grad_reduction(
424 graph_module,
425 sharded_param_indices,
426 classification.replicated_indices,
427 state_fqns,
428 num_state_inputs,
429 state_is_param,
430 model,
431 )
433 if self._dp_mode != "ddp":
434 # Reshard: free each replicated parameter once forward is done and
435 # re-gather + rematerialize it for the backward. Runs after the
436 # grad reduction so the grad path is untouched; no-op when disabled.
437 graph_module = self._insert_reshard_logic(
438 graph_module, param_nodes, pass_config
439 )
441 # Shard the live model's parameters in place (dim 0, by this rank's
442 # index in the FSDP group). ``model.parameters()`` then yields the
443 # shards, so the trainer's optimizer / grad accumulation stay
444 # DP-agnostic; the graph re-gathers each step.
445 self._shard_live_model_params(model)
447 return graph_module
449 def _classify_dp_params(
450 self,
451 graph_module: fx.GraphModule,
452 state_fqns: List[str],
453 num_state_inputs: int,
454 state_is_param: Optional[List[bool]] = None,
455 model: Optional[nn.Module] = None,
456 ) -> _DpParamClassification:
457 """Classify parameter placeholders by how DP handles their gradient.
459 Parameters are the leading ``num_state_inputs`` placeholders of the
460 joint graph, in ``state_fqns`` order. A parameter is DP-marked when any
461 ancestor module FQN (e.g. ``layers.0.attention.wq`` for
462 ``layers.0.attention.wq.weight``) is marked in the GraphParallelPlan —
463 either exactly or via pattern (e.g. ``layers.*``).
465 Buffers (e.g. RoPE's non-persistent ``cache``) and frozen parameters
466 are skipped: the tracer omits them from the gradient list and the
467 live-model sharder leaves them full-size. The pad-shard and unmarked-
468 fallback decisions are delegated to ``_marked_param_stays_replicated``
469 and ``_unmarked_disposition``.
471 Returns:
472 A ``_DpParamClassification`` (see its docstring).
473 """
474 sharded_nodes: List[fx.Node] = []
475 replicated_indices: Set[int] = set()
476 unmarked_trainable: List[Tuple[int, str]] = []
478 placeholders = [n for n in graph_module.graph.nodes if n.op == "placeholder"]
479 if len(placeholders) < num_state_inputs:
480 return _DpParamClassification(
481 sharded_nodes, replicated_indices, unmarked_trainable
482 )
484 # Build a FQN -> Parameter lookup so the divisibility gate uses the
485 # same source of truth as ``_shard_live_model_params`` (which iterates
486 # ``model.named_parameters``). ``state_fqns`` covers both parameters
487 # and buffers; only parameters are reachable here because buffers are
488 # filtered out by ``state_is_param`` above.
489 param_lookup: Dict[str, nn.Parameter] = (
490 dict(model.named_parameters(remove_duplicate=False))
491 if model is not None
492 else {}
493 )
495 for idx in range(num_state_inputs):
496 if state_is_param is not None and not state_is_param[idx]:
497 # Buffer, not a parameter -> never FSDP-sharded.
498 continue
500 fqn = state_fqns[idx]
501 param = param_lookup.get(fqn)
503 # Frozen params are not trainable: the tracer omits them from the
504 # gradient list and ``_shard_live_model_params`` leaves them
505 # full-size, so there is no gradient to classify.
506 if param is not None and not param.requires_grad:
507 _LOG.info("Skip %s: requires_grad=False (frozen)", fqn)
508 continue
510 fallback = self._unmarked_disposition(fqn, param, model)
511 if fallback is not None:
512 if fallback:
513 replicated_indices.add(idx)
514 unmarked_trainable.append((idx, fqn))
515 continue
517 if self._marked_param_stays_replicated(fqn, param):
518 replicated_indices.add(idx)
519 continue
521 node = placeholders[idx]
522 node.meta["state_idx"] = idx
523 node.meta["param_name"] = fqn
524 node.meta["fsdp_degree"] = self._fsdp_degree
525 node.meta["is_param"] = True
526 sharded_nodes.append(node)
527 self._fsdp_modules.add(self._get_parent_module_fqn(fqn))
529 return _DpParamClassification(
530 sharded_nodes, replicated_indices, unmarked_trainable
531 )
533 def _unmarked_disposition(
534 self,
535 fqn: str,
536 param: Optional[nn.Parameter],
537 model: Optional[nn.Module],
538 ) -> Optional[bool]:
539 """Decide how a trainable param a partial plan left unmarked is handled.
541 Rather than leave it with a local, unreduced gradient (silent
542 divergence across DP ranks), fall it back to DDP: keep it replicated
543 and average its gradient over the DP group.
545 Returns:
546 ``True`` to replicate it and warn the caller, ``False`` to skip it
547 silently (``param is None`` with a live model means a buffer, only
548 reachable when an old trace lacks ``state_is_param``), or ``None``
549 when the param is marked and the caller should shard it normally.
550 """
551 if self._parallel_plan is None or self._param_belongs_to_fsdp_module(fqn):
552 return None
553 return model is None or param is not None
555 def _marked_param_stays_replicated(
556 self, fqn: str, param: Optional[nn.Parameter]
557 ) -> bool:
558 """Decide whether a marked param is kept replicated instead of sharded.
560 A marked param whose dim 0 is not divisible by ``fsdp_degree`` cannot be
561 chunked evenly. A 0-dim scalar cannot be padded along dim 0, so it stays
562 replicated and its gradient is averaged over the shard axis. Every
563 other non-divisible param is recorded for pad-sharding (mirroring eager
564 ``fully_shard``): dim 0 is zero-padded to a multiple of the degree, so
565 it is still sharded and gradient-reduced. DDP has no shard axis, so
566 every marked param is handled the same way regardless of divisibility.
567 """
568 if self._dp_mode == "ddp" or param is None:
569 return False
570 shape = param.shape
571 if shape and shape[0] % self._fsdp_degree == 0:
572 return False
573 if not shape:
574 _LOG.info(
575 "Scalar param %s stays replicated: 0-dim cannot be "
576 "pad-sharded; its gradient will be all-reduced",
577 fqn,
578 )
579 return True
580 n = int(shape[0])
581 n_pad = -(-n // self._fsdp_degree) * self._fsdp_degree
582 self._padded_params[fqn] = _PaddedParam(
583 n=n, n_pad=n_pad, dim0_pad=n_pad - n, ndim=param.ndim
584 )
585 _LOG.info(
586 "Param %s pad-sharded: dim 0 %d -> %d (fsdp_degree %d)",
587 fqn,
588 n,
589 n_pad,
590 self._fsdp_degree,
591 )
592 return False
594 @staticmethod
595 def _warn_unmarked_trainable(entries: List[Tuple[int, str]]) -> None:
596 """Report trainable params a partial plan left unmarked.
598 They are not FSDP-sharded, so they fall back to DDP: kept replicated and
599 their gradient averaged over the DP group. The warning is informational —
600 the plan simply did not cover them, and they cost no sharded memory but a
601 replicated gradient all-reduce.
602 """
603 shown = [fqn for _, fqn in entries[:10]]
604 remaining = len(entries) - len(shown)
605 _LOG.warning(
606 "%d trainable parameter(s) are not marked for DP; falling back to "
607 "DDP (kept replicated, gradient all-reduced over the DP group): "
608 "%s%s. Mark them via GraphParallelPlan.fsdp_mark()/"
609 "fsdp_mark_pattern() to FSDP-shard them instead.",
610 len(entries),
611 ", ".join(shown),
612 f" (+{remaining} more)" if remaining > 0 else "",
613 )
615 def _shard_live_model_params(self, model: nn.Module) -> None:
616 """
617 Physically shard the live model's parameters in place (dim 0).
619 After this, model.parameters() yields the local shards.
620 """
621 # Use the FSDP group's local rank (NOT the global rank) as the
622 # chunk index — when fsdp_degree < world_size (TP+FSDP), the global
623 # rank exceeds the chunk count and causes IndexError.
624 fsdp_pg = _resolve_process_group(self._fsdp_group_name)
625 rank = dist.get_rank(group=fsdp_pg)
626 sharded_count = 0
628 for name, param in model.named_parameters():
629 if not param.requires_grad:
630 continue
632 # Same plan gate as ``_classify_dp_params``: only parameters inside
633 # FSDP-marked modules are sharded. Skipping the gate here would
634 # shard params the graph still expects full-rank
635 # (no AllGather inserted), causing a shape mismatch at
636 # ``run_traced_graph`` time on any non-``*`` plan.
637 if (
638 self._parallel_plan is not None
639 and not self._param_belongs_to_fsdp_module(name)
640 ):
641 _LOG.info(
642 "Skip %s: not in an FSDP-wrapped module",
643 name,
644 )
645 continue
647 # Pad-shard a non-divisible param (recorded by
648 # ``_classify_dp_params``): zero-pad dim 0 to a multiple of the
649 # degree, then take this rank's equal-size chunk — exactly what
650 # eager ``fully_shard`` does. A 0-dim scalar cannot be padded and
651 # stays replicated.
652 pad_info = self._padded_params.get(name)
653 if pad_info is not None:
654 original_shape = param.shape
655 full = param.detach()
656 if pad_info.dim0_pad:
657 full = torch.cat(
658 [full, full.new_zeros((pad_info.dim0_pad, *full.shape[1:]))],
659 dim=0,
660 )
661 param.data = full.chunk(self._fsdp_degree, dim=0)[rank].clone()
662 sharded_count += 1
663 _LOG.info(
664 "Pad-sharded %s: %s -> %s",
665 name,
666 list(original_shape),
667 list(param.shape),
668 )
669 continue
671 # Mirror ``_classify_dp_params``: a 0-dim scalar stays replicated.
672 if not param.shape or param.shape[0] % self._fsdp_degree != 0:
673 _LOG.info(
674 "Skip %s: dim 0 (%s) not divisible by fsdp_degree (%s)",
675 name,
676 param.shape[0] if param.shape else "scalar",
677 self._fsdp_degree,
678 )
679 continue
681 original_shape = param.shape
682 param.data = param.detach().chunk(self._fsdp_degree, dim=0)[rank].clone()
683 sharded_count += 1
684 _LOG.info(
685 "Sharded %s: %s -> %s",
686 name,
687 list(original_shape),
688 list(param.shape),
689 )
691 _LOG.info("Total sharded parameters: %d", sharded_count)
693 def _param_belongs_to_fsdp_module(self, param_fqn: str) -> bool:
694 """Check if a parameter FQN belongs to an FSDP-wrapped module.
696 The parameter's own module FQN (``layers.0.attention.wq.weight`` ->
697 ``layers.0.attention.wq``) plus every ancestor (``layers.0``, ...) is
698 tested, so both ``fsdp_mark("layers.0.attention.wq")`` and
699 ``fsdp_mark_pattern("layers.*")`` match.
701 A parameter directly on the root (``weight``) has no module ancestor;
702 its own FQN is still tested first so it can be marked by its name
703 (``fsdp_mark("weight")``) or by a ``*`` pattern.
704 """
705 if self._parallel_plan.is_marked_for_fsdp(param_fqn):
706 return True
708 parts = param_fqn.split(".")
709 for i in range(len(parts) - 1, 0, -1):
710 module_fqn = ".".join(parts[:i])
711 if self._parallel_plan.is_marked_for_fsdp(module_fqn):
712 return True
713 return False
715 def _get_parent_module_fqn(self, param_fqn: str) -> str:
716 """
717 Get parent module FQN from parameter FQN.
719 e.g., "layers.0.attention.wq.weight" -> "layers.0.attention.wq"
720 """
721 parts = param_fqn.split(".")
722 return ".".join(parts[:-1]) if len(parts) > 1 else param_fqn
724 def _insert_all_gather_for_params(
725 self,
726 graph_module: fx.GraphModule,
727 param_nodes: List[fx.Node],
728 pass_config: PassConfig,
729 ) -> fx.GraphModule:
730 """
731 Insert AllGather for each FSDP parameter placeholder.
733 Parameter state: Shard -> Replicate. All subsequent uses of the
734 placeholder are rewired to the unsharded (replicated) tensor, so the
735 computation body keeps operating on full parameters while the graph
736 input stays sharded. The gather is placed after the placeholder, or
737 sunk to just before the first consumer when reshard is enabled.
739 When reshard is enabled the gather is **sunk** to just before the
740 parameter's first consumer instead of sitting next to the placeholder.
741 All state placeholders lead the graph, so an unsunk gather would
742 materialize every full parameter before any compute runs, leaving the
743 reshard free step nothing to reclaim; sinking is what makes peak
744 memory track the forward working set.
745 """
746 graph = graph_module.graph
748 for param_node in param_nodes:
749 if param_node.name in self._processed_params:
750 continue
752 sink = pass_config.fsdp_reshard_after_forward
753 if sink and not param_node.users:
754 # Dead parameter: a gather would materialize a full replica
755 # nobody reads, and reshard would never free it. Leave the
756 # placeholder sharded.
757 self._processed_params.add(param_node.name)
758 continue
759 anchor = self._first_user(param_node) if sink else None
761 # AllGather + immediate wait for correctness; AutoOverlapPass may
762 # sink the wait past independent compute later.
763 with _insertion_point(graph, anchor, param_node):
764 ag_node = graph.call_function(
765 _c10d_functional.all_gather_into_tensor,
766 args=(param_node, self._fsdp_degree, self._fsdp_group_name),
767 )
768 ag_node.meta["comm_type"] = "fsdp_all_gather"
769 ag_node.meta["comm_group"] = self._fsdp_group_name
770 ag_node.meta["param_node"] = param_node.name
771 ag_node.meta["param_name"] = param_node.meta.get("param_name")
772 ag_node.meta["state_idx"] = param_node.meta.get("state_idx")
773 ag_node.meta["fsdp_degree"] = self._fsdp_degree
775 # Insert the wait in its own block: a wait placed in the same block
776 # as the gather would land *before* the gather it depends on (the
777 # first insert inside ``inserting_after(X)`` lands right after
778 # ``X``), and codegen would emit
779 # ``wait_tensor(all_gather_into_tensor)`` before that variable is
780 # assigned.
781 with graph.inserting_after(ag_node):
782 wait_node = graph.call_function(
783 _c10d_functional.wait_tensor,
784 args=(ag_node,),
785 )
786 wait_node.meta["wait_for"] = ag_node.name
788 param_node.meta["fsdp_sharded"] = True
789 param_node.meta["fsdp_ag_node"] = ag_node.name
790 self._unsharded_params[param_node.name] = _UnshardedParam(
791 param_node=param_node, ag_node=ag_node, wait_node=wait_node
792 )
794 # Pad-sharded parameter: the gather produced the padded full tensor
795 # ``[n_pad, ...]``; narrow it back to the real ``[n, ...]`` and route
796 # every consumer to that view. Use the ``aten`` overload (not the
797 # ``torch.narrow`` builtin) so its schema's ``alias_info`` marks it a
798 # view: ``_propagates_alias`` keys off that schema, so the reshard
799 # alias walk (free / rematerialize) tracks the narrow and does not
800 # free the padded gather out from under its forward consumers.
801 replacement = wait_node
802 pad_info = self._padded_params.get(param_node.meta.get("param_name"))
803 if pad_info is not None:
804 with graph.inserting_after(wait_node):
805 replacement = graph.call_function(
806 torch.ops.aten.narrow.default,
807 args=(wait_node, 0, 0, pad_info.n),
808 )
809 replacement.meta["param_name"] = param_node.meta.get("param_name")
810 replacement.meta["state_idx"] = param_node.meta.get("state_idx")
812 for user in list(param_node.users.keys()):
813 if user not in (ag_node, wait_node):
814 user.replace_input_with(param_node, replacement)
816 self._processed_params.add(param_node.name)
818 return graph_module
820 @staticmethod
821 def _first_user(param_node: fx.Node) -> Optional[fx.Node]:
822 """Earliest (in graph order) consumer of ``param_node``.
824 Placeholders are all at the top; their consumers are spread through
825 the forward. Returning the earliest consumer lets the gather sink
826 next to the first real use. ``None`` when the placeholder is unused
827 (then the gather stays next to the placeholder).
828 """
829 users = [u for u in param_node.users if u.op != "output"]
830 if not users:
831 return None
832 forward_users = [u for u in users if not u.meta.get("autograd_backward", False)]
833 candidates = forward_users or users
834 order = {n: i for i, n in enumerate(param_node.graph.nodes)}
835 return min(candidates, key=lambda u: order.get(u, len(order)))
837 def _insert_grad_reduction( # pylint: disable=too-many-locals
838 self,
839 graph_module: fx.GraphModule,
840 sharded_param_indices: Set[int],
841 replicated_param_indices: Set[int],
842 state_fqns: List[str],
843 num_state_inputs: int,
844 state_is_param: Optional[List[bool]] = None,
845 model: Optional[nn.Module] = None,
846 ) -> fx.GraphModule:
847 """Insert gradient reduction on the gradient outputs of DP parameters.
849 The joint graph returns ``[loss, grad0, grad1, ...]`` from the fwd+bwd
850 function; gradient ``i`` (output index ``i+1``) corresponds to the
851 ``i``-th trainable parameter (``state_fqns`` order, buffers and frozen
852 params skipped, mirrored by ``_build_trainable_state_indices``). For a
853 parameter in ``sharded_param_indices`` (a real shard axis) the mode
854 decides the collective chain: ``fsdp`` reduce_scatters on the ``fsdp``
855 group (Replicate -> Shard); ``ddp`` all_reduces on ``dp_replicate``;
856 ``hsdp`` reduce_scatters on ``fsdp`` then all_reduces on
857 ``dp_replicate`` (shard axis first, so the cross-axis message is
858 already sharded). Pad-sharded params (``_padded_params``) zero-pad their
859 ``[n, ...]`` gradient back to ``[n_pad, ...]`` first so the local
860 gradient matches the padded shard.
862 For a parameter in ``replicated_param_indices`` (a marked 0-dim scalar,
863 or a partial-plan leftover that fell back to DDP) no shard axis exists,
864 so the gradient is AVERAGED (all_reduce) over the DP group instead;
865 FSDP/HSDP also average the shard axis, HSDP the replicate axis.
866 Gradients of parameters outside the DP modules stay full.
867 """
868 graph = graph_module.graph
870 output_node = next((n for n in graph.nodes if n.op == "output"), None)
871 if output_node is None:
872 return graph_module
874 returned = output_node.args[0]
875 if not isinstance(returned, (list, tuple)):
876 # Model returned a single value (e.g. just loss) -> no grads.
877 return graph_module
879 trainable_state_indices = self._build_trainable_state_indices(
880 state_fqns, num_state_inputs, state_is_param, model
881 )
883 new_returned = list(returned)
884 # ``train_fn`` may append the traced ``loss_dict`` outputs AFTER the
885 # gradients (see ``trace_model_graph``); those are not gradients and
886 # must stay outside the reduce_scatter range.
887 num_loss_outputs = len(getattr(graph_module, "loss_dict_keys", []))
888 num_grads = len(new_returned) - 1 - num_loss_outputs # index 0 is the loss
889 if num_grads != len(trainable_state_indices):
890 raise ValueError(
891 f"Gradient count ({num_grads}) does not match trainable "
892 f"parameter count ({len(trainable_state_indices)}). The "
893 f"traced graph and the live model disagree on which "
894 f"parameters are trainable; refusing to insert gradient "
895 f"reduction to avoid silent gradient/state misalignment."
896 )
898 # Index 0 is the loss; gradients start at index 1. Gradient i+1
899 # corresponds to trainable parameter i, whose state_idx is
900 # trainable_state_indices[i]. The trailing loss_dict outputs are
901 # named loss values, not gradients.
902 for i in range(1, len(new_returned) - num_loss_outputs):
903 grad_node = new_returned[i]
904 if not isinstance(grad_node, fx.Node):
905 continue
907 state_idx = trainable_state_indices[i - 1]
908 if state_idx in sharded_param_indices:
909 replicated = False
910 pad_info = self._padded_params.get(state_fqns[state_idx])
911 elif state_idx in replicated_param_indices:
912 replicated = True
913 pad_info = None
914 else:
915 continue
917 with graph.inserting_before(output_node):
918 new_returned[i] = self._build_grad_reduction_nodes(
919 graph, grad_node, replicated=replicated, pad_info=pad_info
920 )
922 output_node.args = (type(returned)(new_returned),) + tuple(output_node.args[1:])
924 return graph_module
926 def _build_grad_reduction_nodes(
927 self,
928 graph: fx.Graph,
929 grad_node: fx.Node,
930 replicated: bool = False,
931 pad_info: Optional[_PaddedParam] = None,
932 ) -> fx.Node:
933 """Build and return the waited gradient-reduction node for ``grad_node``.
935 The returned (wait) node replaces ``grad_node`` in the graph output.
936 Emits the mode-specific chain (see ``_insert_grad_reduction``). When
937 ``pad_info`` is given the parameter is pad-sharded: the real ``[n, ...]``
938 gradient is zero-padded on dim 0 to ``[n_pad, ...]`` before the
939 reduce_scatter, so the scattered local gradient matches the padded live
940 shard (mirrors eager ``fully_shard``'s padded grad buffer). When
941 ``replicated`` is True the parameter is a 0-dim scalar that could not be
942 sharded, so its gradient is AVERAGED over the shard group with
943 ``all_reduce`` instead of reduce-scattered — the graph and the live
944 model both keep the parameter full, so a reduce_scatter would change the
945 shape and desync them.
947 Every collective is waited on before its result is consumed. A chained
948 collective runs on a different process-group stream (HSDP's
949 reduce_scatter is on ``fsdp`` while the all_reduce is on
950 ``dp_replicate``), so without the intermediate wait the all_reduce could
951 read a partially written shard. This mirrors eager ``fully_shard``,
952 which orders the all_reduce stream after the reduce_scatter stream.
953 """
954 reduced = grad_node
955 if pad_info is not None:
956 # constant_pad_nd's pad tuple runs last-dim-first: pad dim 0 on the
957 # right by ``dim0_pad`` and leave every other dim untouched.
958 pad_tuple = tuple([0, 0] * (pad_info.ndim - 1) + [0, pad_info.dim0_pad])
959 reduced = graph.call_function(
960 torch.constant_pad_nd, args=(reduced, pad_tuple)
961 )
962 reduced.meta["comm_type"] = "fsdp_grad_pad"
963 if replicated:
964 # Replicated parameter (a 0-dim scalar, or a partial-plan leftover
965 # that fell back to DDP): average the full gradient over the shard
966 # group so every replica takes the same optimizer step. DDP mode has
967 # no shard group — the replicate all_reduce below handles it.
968 if self._dp_mode in ("fsdp", "hsdp"):
969 ar_node = graph.call_function(
970 _c10d_functional.all_reduce,
971 args=(reduced, _GRAD_REDUCE_OP, self._fsdp_group_name),
972 )
973 ar_node.meta["comm_type"] = "fsdp_all_reduce"
974 ar_node.meta["comm_group"] = self._fsdp_group_name
975 ar_node.meta["fsdp_degree"] = self._fsdp_degree
976 reduced = ar_node
977 # HSDP chains the replicate all_reduce (a different process-group
978 # stream) after the shard all_reduce, so wait first.
979 if self._dp_mode == "hsdp":
980 reduced = self._insert_wait(graph, reduced)
981 elif self._dp_mode in ("fsdp", "hsdp"):
982 rs_node = graph.call_function(
983 _c10d_functional.reduce_scatter_tensor,
984 args=(
985 reduced,
986 _GRAD_REDUCE_OP,
987 self._fsdp_degree,
988 self._fsdp_group_name,
989 ),
990 )
991 rs_node.meta["comm_type"] = "fsdp_reduce_scatter"
992 rs_node.meta["comm_group"] = self._fsdp_group_name
993 rs_node.meta["fsdp_degree"] = self._fsdp_degree
994 reduced = rs_node
995 # HSDP chains all_reduce (a different process-group stream) after
996 # the reduce_scatter, so wait first; plain FSDP waits on the final
997 # node only.
998 if self._dp_mode == "hsdp":
999 reduced = self._insert_wait(graph, reduced)
1001 if self._dp_mode in ("ddp", "hsdp"):
1002 ar_node = graph.call_function(
1003 _c10d_functional.all_reduce,
1004 args=(reduced, _GRAD_REDUCE_OP, self._dp_replicate_group_name),
1005 )
1006 ar_node.meta["comm_type"] = "dp_all_reduce"
1007 ar_node.meta["comm_group"] = self._dp_replicate_group_name
1008 ar_node.meta["dp_replicate_degree"] = self._dp_replicate_degree
1009 reduced = ar_node
1011 return self._insert_wait(graph, reduced)
1013 @staticmethod
1014 def _insert_wait(graph: fx.Graph, tensor_node: fx.Node) -> fx.Node:
1015 """Append ``wait_tensor(tensor_node)`` and return the waited node."""
1016 wait_node = graph.call_function(
1017 _c10d_functional.wait_tensor,
1018 args=(tensor_node,),
1019 )
1020 wait_node.meta["wait_for"] = tensor_node.name
1021 return wait_node
1023 def _build_trainable_state_indices(
1024 self,
1025 state_fqns: List[str],
1026 num_state_inputs: int,
1027 state_is_param: Optional[List[bool]],
1028 model: Optional[nn.Module],
1029 ) -> List[int]:
1030 """
1031 Return state indices of trainable parameters, in ``state_fqns`` order.
1033 Mirrors the tracer's ``params`` list construction (skip buffers via
1034 ``state_is_param``, skip frozen params via ``requires_grad``) so
1035 gradient ``i`` maps to ``trainable_state_indices[i]``. Reproducing this
1036 filter here is what keeps reduce_scatter aligned for models with
1037 buffers or frozen params.
1038 """
1039 if model is None:
1040 # Without the model we cannot inspect requires_grad; fall back to
1041 # the all-trainable assumption. The grad-count check in the caller
1042 # raises if that assumption is wrong, so misalignment is not silent.
1043 return list(range(num_state_inputs))
1045 param_lookup = dict(model.named_parameters(remove_duplicate=False))
1046 trainable: List[int] = []
1047 for idx in range(num_state_inputs):
1048 if state_is_param is not None and not state_is_param[idx]:
1049 continue
1050 param = param_lookup.get(state_fqns[idx])
1051 if param is not None and param.requires_grad:
1052 trainable.append(idx)
1053 return trainable
1055 def _insert_reshard_logic(
1056 self,
1057 graph_module: fx.GraphModule,
1058 param_nodes: List[fx.Node],
1059 pass_config: PassConfig,
1060 ) -> fx.GraphModule:
1061 """Reshard each FSDP parameter after forward (FSDP release).
1063 For every FSDP parameter whose all_gather was inserted by
1064 ``_insert_all_gather_for_params``:
1066 1. Free the replicated full parameter once its last forward reader —
1067 including forward views aliasing it — has run.
1068 2. If the backward still needs the parameter (directly or through a
1069 saved forward view), insert a fresh all_gather just before the
1070 first backward consumer and rematerialize the saved view chain, so
1071 the backward computes the exact same values from the re-gathered
1072 (unsharded) parameter.
1073 3. Free that backward re-gather once its last backward reader runs.
1075 Steps 1-3 keep only the forward working set resident during forward and
1076 only the active layer's parameter resident during backward, instead of
1077 every replicated parameter for the whole joint graph.
1079 No-op when ``fsdp_reshard_after_forward`` is disabled.
1080 """
1081 if not pass_config.fsdp_reshard_after_forward:
1082 return graph_module
1084 for param_node in param_nodes:
1085 record = self._unsharded_params.get(param_node.name)
1086 if record is None:
1087 continue
1088 self._reshard_param(graph_module.graph, record)
1089 return graph_module
1091 def _reshard_param(self, graph: fx.Graph, record: _UnshardedParam) -> None:
1092 """Reshard one FSDP parameter (free + backward re-gather/remat)."""
1093 wait = record.wait_node
1094 alias_nodes = self._alias_descendants(wait)
1095 order = {n: i for i, n in enumerate(graph.nodes)}
1097 # Alias (view) descendants the backward reads must be rematerialized
1098 # from the re-gathered parameter; the unsharded tensor itself may be
1099 # consumed directly by the backward in some models.
1100 needed = [
1101 n for n in [wait, *alias_nodes] if self._feeds_backward(n, alias_nodes)
1102 ]
1104 # Remat plan first: the forward free is only safe once every backward
1105 # consumer has been rewired onto fresh (re-gathered) storage. If the
1106 # view chain cannot be cloned, keep the previous resident-for-the-graph
1107 # behavior — a missed free costs memory, a misordered one is unsound.
1108 if self._rematerialize_for_backward(graph, record, needed, order):
1109 # 1. Free after the last forward data read of the unsharded storage.
1110 self._insert_reshard_free(graph, wait, record, order, backward=False)
1112 def _rematerialize_for_backward(
1113 self,
1114 graph: fx.Graph,
1115 record: _UnshardedParam,
1116 needed: List[fx.Node],
1117 order: Dict[fx.Node, int],
1118 ) -> bool:
1119 """Re-gather ahead of the first backward consumer and remat views.
1121 The saved forward view chain is cloned against the re-gathered
1122 parameter so the backward reads identical values from a fresh,
1123 correctly-sized tensor.
1125 Returns:
1126 ``True`` when the free may proceed — no backward consumer exists,
1127 or every one of them was rewired onto rematerialized storage.
1128 ``False`` when a needed view cannot be cloned; the caller must
1129 keep the unsharded parameter resident.
1130 """
1131 backward_users = [
1132 u for n in needed for u in n.users if u.meta.get("autograd_backward", False)
1133 ]
1134 if not backward_users:
1135 return True
1137 wait = record.wait_node
1138 uncloneable = [n for n in needed if n is not wait and not self._cloneable(n)]
1139 if uncloneable:
1140 _LOG.warning(
1141 "FSDP reshard: cannot rematerialize %s for '%s'; keeping the "
1142 "unsharded parameter resident",
1143 [n.format_node() for n in uncloneable],
1144 record.param_node.meta.get("param_name"),
1145 )
1146 return False
1148 anchor = min(backward_users, key=lambda u: order.get(u, len(order)))
1149 wait2 = self._insert_backward_gather(graph, record, anchor)
1151 recreate: Dict[fx.Node, fx.Node] = {wait: wait2}
1152 last_inserted = wait2
1153 for node in sorted(
1154 (n for n in needed if n is not wait), key=lambda n: order[n]
1155 ):
1156 last_inserted = self._clone_with_remap(
1157 graph, node, recreate, wait, wait2, last_inserted
1158 )
1160 self._rewire_backward_users(recreate)
1161 self._insert_reshard_free(graph, wait2, record, order, backward=True)
1162 return True
1164 @staticmethod
1165 def _cloneable(node: fx.Node) -> bool:
1166 """Whether ``node`` can be cloned as a rematerialized view op."""
1167 return node.op == "call_function" and callable(node.target)
1169 def _insert_backward_gather(
1170 self, graph: fx.Graph, record: _UnshardedParam, anchor: fx.Node
1171 ) -> fx.Node:
1172 """Insert the backward all_gather + wait pair before ``anchor``."""
1173 with graph.inserting_before(anchor):
1174 ag_node = graph.call_function(
1175 _c10d_functional.all_gather_into_tensor,
1176 args=(record.param_node, self._fsdp_degree, self._fsdp_group_name),
1177 )
1178 ag_node.meta["comm_type"] = "fsdp_all_gather_backward"
1179 ag_node.meta["comm_group"] = self._fsdp_group_name
1180 ag_node.meta["param_node"] = record.param_node.name
1181 ag_node.meta["param_name"] = record.param_node.meta.get("param_name")
1182 ag_node.meta["fsdp_degree"] = self._fsdp_degree
1183 # This gather lives in the backward half; tag it so downstream
1184 # passes (e.g. PpPass phase classification) do not mistake it and
1185 # its remat clones for forward nodes — their only arg is a forward
1186 # state placeholder, which would otherwise imply the forward phase.
1187 ag_node.meta["autograd_backward"] = True
1188 with graph.inserting_before(anchor):
1189 wait_node = graph.call_function(
1190 _c10d_functional.wait_tensor, args=(ag_node,)
1191 )
1192 wait_node.meta["wait_for"] = ag_node.name
1193 wait_node.meta["autograd_backward"] = True
1194 return wait_node
1196 def _clone_with_remap(
1197 self,
1198 graph: fx.Graph,
1199 node: fx.Node,
1200 recreate: Dict[fx.Node, fx.Node],
1201 wait: fx.Node,
1202 wait2: fx.Node,
1203 insert_after: fx.Node,
1204 ) -> fx.Node:
1205 """Clone ``node`` after ``insert_after``, remapping args to re-gathers."""
1206 args = tuple(self._remap_arg(a, recreate, wait, wait2) for a in node.args)
1207 kwargs = {
1208 k: self._remap_arg(v, recreate, wait, wait2) for k, v in node.kwargs.items()
1209 }
1210 with graph.inserting_after(insert_after):
1211 clone = graph.call_function(node.target, args=args, kwargs=kwargs)
1212 self._copy_recreated_meta(node, clone)
1213 # Rematerialized views serve the backward half, so they carry the
1214 # backward phase (see ``_insert_backward_gather``).
1215 clone.meta["autograd_backward"] = True
1216 recreate[node] = clone
1217 return clone
1219 @staticmethod
1220 def _rewire_backward_users(recreate: Dict[fx.Node, fx.Node]) -> None:
1221 """Point every backward consumer of an original node at its clone."""
1222 for original, clone in recreate.items():
1223 for user in list(original.users):
1224 if user.meta.get("autograd_backward", False):
1225 user.replace_input_with(original, clone)
1227 def _insert_reshard_free(
1228 self,
1229 graph: fx.Graph,
1230 tensor: fx.Node,
1231 record: _UnshardedParam,
1232 order: Dict[fx.Node, int],
1233 backward: bool,
1234 ) -> None:
1235 """Free ``tensor``'s storage after its last reader in the phase."""
1236 last_reader = self._last_data_reader(tensor, order, backward=backward)
1237 if last_reader is None:
1238 return
1239 comm_type = "fsdp_reshard_free_backward" if backward else "fsdp_reshard_free"
1240 with graph.inserting_after(last_reader):
1241 free_node = graph.call_function(_free_tensor_storage, args=(tensor,))
1242 free_node.meta["comm_type"] = comm_type
1243 free_node.meta["param_name"] = record.param_node.meta.get("param_name")
1244 if backward:
1245 free_node.meta["autograd_backward"] = True
1247 def _alias_descendants(self, root: fx.Node) -> Set[fx.Node]:
1248 """Forward view descendants of ``root`` (nodes aliasing its storage).
1250 Only aliasing ops propagate: an op that allocates a new tensor
1251 (``mm``, ``addmm``, ``relu``, ...) copies the data out, so freeing
1252 ``root``'s storage does not invalidate it and it need not be
1253 rematerialized. Backward nodes terminate the walk.
1254 """
1255 alias: Set[fx.Node] = set()
1256 stack = [root]
1257 while stack:
1258 current = stack.pop()
1259 for user in current.users:
1260 if user is root or user in alias:
1261 continue
1262 if user.meta.get("autograd_backward", False):
1263 continue
1264 if current not in set(_iter_arg_nodes(user)):
1265 continue
1266 if not self._propagates_alias(user):
1267 continue
1268 alias.add(user)
1269 stack.append(user)
1270 return alias
1272 @staticmethod
1273 def _propagates_alias(node: fx.Node) -> bool:
1274 """Whether ``node``'s output aliases its already-aliasing input.
1276 Two forms propagate. Schema-declared view returns (``aten.t``,
1277 ``aten.slice``, ...), and ``operator.getitem`` picking one output of a
1278 multi-output view op such as ``aten.chunk`` / ``aten.split`` /
1279 ``aten.unbind`` — getitem itself has no schema, so the alias must be
1280 inherited from the indexed producer (which the caller has already
1281 established is in the alias closure).
1282 """
1283 if node.op != "call_function":
1284 return False
1285 if node.target is _GETITEM:
1286 return True
1287 schema = getattr(node.target, "_schema", None)
1288 if schema is None:
1289 return False
1290 return any(ret.alias_info is not None for ret in schema.returns)
1292 def _feeds_backward(self, node: fx.Node, alias_nodes: Set[fx.Node]) -> bool:
1293 """Whether ``node`` (or an aliasing descendant) is read by backward."""
1294 memo: Dict[fx.Node, bool] = {}
1296 def visit(current: fx.Node) -> bool:
1297 """Whether ``current`` or an aliasing descendant reaches backward.
1299 Args:
1300 current: Node to test.
1302 Returns:
1303 True when a backward node consumes the node or its aliases.
1304 """
1305 if current in memo:
1306 return memo[current]
1307 memo[current] = False # cycle guard (aliasing is a DAG)
1308 result = False
1309 for user in current.users:
1310 if user.meta.get("autograd_backward", False):
1311 result = True
1312 break
1313 if user in alias_nodes and visit(user):
1314 result = True
1315 break
1316 memo[current] = result
1317 return result
1319 return visit(node)
1321 def _last_data_reader(
1322 self, root: fx.Node, order: Dict[fx.Node, int], backward: bool
1323 ) -> Optional[fx.Node]:
1324 """Latest, in graph order, op that reads ``root``'s storage data.
1326 Walks the alias (view) chain so a free is not placed after a view
1327 *creation* but before the non-view op that actually reads it: freeing
1328 between ``t_2 = t(w)`` and ``mm(grad, t_2)`` would invalidate ``t_2``.
1329 Multi-output views reach their consumers through ``getitem``, so the
1330 walk continues through those picks instead of mistaking them for the
1331 final reader. ``backward=False`` picks the last forward read (free
1332 point for the unsharded parameter); ``backward=True`` the last
1333 backward read (free point for the backward re-gather).
1334 """
1335 last: Optional[fx.Node] = None
1336 last_pos = -1
1337 seen: Set[fx.Node] = set()
1338 stack = [root]
1339 while stack:
1340 current = stack.pop()
1341 if current in seen:
1342 continue
1343 seen.add(current)
1344 for user in current.users:
1345 if user.op == "output":
1346 continue
1347 if self._propagates_alias(user) and current in set(
1348 _iter_arg_nodes(user)
1349 ):
1350 stack.append(user)
1351 continue
1352 if bool(user.meta.get("autograd_backward", False)) != backward:
1353 continue
1354 pos = order.get(user, -1)
1355 if pos > last_pos:
1356 last_pos = pos
1357 last = user
1358 return last
1360 @classmethod
1361 def _remap_arg(
1362 cls,
1363 arg: Any,
1364 recreate: Dict[fx.Node, fx.Node],
1365 wait: fx.Node,
1366 wait2: fx.Node,
1367 ) -> Any:
1368 """Rewrite an op arg for a rematerialized clone (nested containers)."""
1369 if arg is wait:
1370 return wait2
1371 if isinstance(arg, fx.Node):
1372 return recreate.get(arg, arg)
1373 if isinstance(arg, tuple):
1374 return tuple(cls._remap_arg(a, recreate, wait, wait2) for a in arg)
1375 if isinstance(arg, list):
1376 return [cls._remap_arg(a, recreate, wait, wait2) for a in arg]
1377 if isinstance(arg, dict):
1378 return {k: cls._remap_arg(v, recreate, wait, wait2) for k, v in arg.items()}
1379 return arg
1381 @staticmethod
1382 def _copy_recreated_meta(original: fx.Node, clone: fx.Node) -> None:
1383 """Copy the module/stack/shape metadata a remat clone needs."""
1384 for key in (
1385 "nn_module_stack",
1386 "custom",
1387 "stack_trace",
1388 "tensor_meta",
1389 "val",
1390 ):
1391 value = original.meta.get(key)
1392 if value is not None:
1393 clone.meta[key] = value