Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / tracer / graph_tracer.py: 93%
263 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# Tracer reaches into torch internals (torch._dynamo, torch.nn.utils.stateless,
16# torch.utils._pytree) for joint-graph capture; these have no public API.
17# pylint: disable=E1136,E1129,W0212
18"""
19Graph Tracer - Model Graph Tracer (make_fx + autograd.grad)
21Traces a model into a joint Forward + Backward FX graph, following
22the design of torchtitan's ``experiments/graph_trainer/minimal_fx_tracer``:
241. Module parameters/buffers are extracted from the live model into flat
25 tensors and threaded through the graph as **static inputs** (leading
26 placeholders) instead of ``get_attr`` nodes, via
27 ``torch.nn.utils.stateless._reparametrize_module``. The graph therefore
28 contains no parameter ``get_attr`` nodes, and passes can split the graph by
29 reshaping the static-input placeholders.
312. The forward pass runs ``torch.autograd.grad(loss, params)`` *inside* the
32 traced function body. For that to be traceable the engine backward must be
33 patched (``torch.compiler._patch_engine_backward()``) and backward must run
34 on the calling thread (``torch.autograd.set_multithreading_enabled(False)``),
35 otherwise the C++ autograd engine dispatches backward to a worker thread
36 with a fresh ``contextvars.Context`` and the tracing context is lost.
383. The whole trace runs under ``FakeTensorMode(allow_non_fake_inputs=True)``
39 plus a non-strict tracing context, ``preserve_node_meta()`` and
40 ``_skip_nested_compile()``, mirroring the torchtitan tracer.
42The joint graph is executed with ``torch.no_grad()`` at runtime: it already
43contains the explicit backward ops, and a redundant autograd graph on top
44would keep every forward intermediate alive via ``grad_fn`` references.
45"""
47from __future__ import annotations
49__all__ = [
50 "JointGraph",
51 "trace_model_graph",
52 "run_traced_graph",
53 "extract_module_state",
54]
56import contextlib
57import copy
58import inspect
59import warnings
60from collections.abc import Callable, Generator
61from collections import deque
62from contextlib import contextmanager
63from dataclasses import dataclass
64from typing import Any, Dict, List
66import torch
67from torch import nn
68from torch._decomp import get_decompositions
69from torch._guards import tracing, TracingContext
70from torch._subclasses import FakeTensorMode
71from torch.fx.experimental.proxy_tensor import make_fx
72from torch.fx.traceback import preserve_node_meta
73from torch.nn.utils import stateless
75from .dynamic_shapes import InputGuards, build_symbolic_inputs
77# Older torch's make_fx predates record_stack_traces / record_module_stack
78# and rejects unknown kwargs; forward only the kwargs this torch supports.
79# The joint graph still captures without record_stack_traces -- the fwd->bw
80# metadata copy skips a missing stack_trace instead of failing.
81try:
82 _MAKE_FX_PARAMS = inspect.signature(make_fx).parameters
83except (ValueError, TypeError):
84 _MAKE_FX_PARAMS = {}
85_MAKE_FX_KWARGS: Dict[str, Any] = {
86 kw: val
87 for kw, val in (("record_stack_traces", True), ("record_module_stack", True))
88 if kw in _MAKE_FX_PARAMS
89}
91# Tensors and make_fx-safe primitives are allowed as graph leaves.
92# Everything else (callables, custom objects) must be captured in the
93# train_fn closure or registered as pytree constants.
94_ALLOWED_LEAF_TYPES = (torch.Tensor, int, float, bool, str, type(None))
97@contextmanager
98def _skip_nested_compile() -> Generator[None, None, None]:
99 """Tell dynamo to inline ``torch.compile``'d functions during make_fx tracing.
101 make_fx cannot trace through torch.compile'd functions (e.g. compiled
102 attention kernels). Setting ``error_on_nested_fx_trace`` to False makes
103 dynamo silently inline the wrapped function instead of raising.
104 """
105 prev = torch._dynamo.config.error_on_nested_fx_trace
106 torch._dynamo.config.error_on_nested_fx_trace = False
107 try:
108 yield
109 finally:
110 torch._dynamo.config.error_on_nested_fx_trace = prev
113@contextmanager
114def _patch_engine_backward() -> Generator[None, None, None]:
115 """No-op fallback for torch versions without the patched autograd engine.
117 ``torch.compiler._patch_engine_backward`` makes ``torch.autograd.grad``
118 inside a make_fx-traced function emit backward ops as FX nodes instead of
119 running a real autograd backward. Present in torchtitan's patched torch;
120 absent in stock torch, where joint-graph capture of ``autograd.grad`` is
121 not supported at all.
122 """
123 patcher = getattr(torch.compiler, "_patch_engine_backward", None)
124 if patcher is None:
125 warnings.warn(
126 "torch.compiler._patch_engine_backward not available. "
127 "Joint-graph capture of autograd.grad may not work correctly. "
128 "Consider using torchtitan's patched torch or a newer torch version.",
129 RuntimeWarning,
130 stacklevel=2,
131 )
132 yield
133 else:
134 with patcher():
135 yield
138@contextmanager
139def _non_strict_tracing_context() -> Generator[None, None, None]:
140 """Mark this make_fx pass as the non-strict tracing flow.
142 Required by ``_patch_engine_backward``; no-op fallback on stock torch.
143 """
144 ctx = getattr(torch.compiler, "_non_strict_tracing_context", None)
145 if ctx is None:
146 yield
147 else:
148 with ctx():
149 yield
152@dataclass
153class JointGraph:
154 """
155 Joint forward-backward computation graph
157 Features:
158 - No parallel information
159 - Parameters, activations, gradients are all complete tensors
160 - Forward + Loss + Backward complete joint graph
161 - Parameters/buffers are static graph inputs (leading placeholders),
162 not ``get_attr`` nodes
163 """
165 graph_module: torch.fx.GraphModule
166 # Storage-free descriptors of the tracing-time user inputs (shape/dtype
167 # for tensors) via ``_input_meta``. Real sample tensors are deliberately
168 # not kept: they would pin their storage for the whole training run. The
169 # FakeTensor counterparts live in ``example_inputs``.
170 inputs: List[Any]
171 outputs: List[Any]
172 param_names: List[str]
173 param_shapes: Dict[str, tuple]
174 num_layers: int
175 # Trace-time parameter/buffer values (flat, in graph placeholder order)
176 # and their FQNs. ``state_fqns`` is the authoritative order for feeding
177 # runtime parameter state into the graph.
178 state_fqns: List[str]
179 example_inputs: tuple
180 input_guards: InputGuards | None = None
183def extract_module_state(mod: nn.Module) -> Dict[str, torch.Tensor]:
184 """Return a merged dict of the module's named parameters and buffers.
186 Tied tensors (e.g. an ``lm_head.weight`` sharing storage with
187 ``embed_tokens.weight``) are DEDUPLICATED to a single entry: the graph
188 threads each unique tensor as one static input, so emitting one entry per
189 alias would double the placeholder count, duplicate every FSDP collective,
190 and make ``torch.autograd.grad`` return (and the compiler add) the same
191 gradient once per alias. The missing aliases are re-materialized only for
192 tracing, via ``_reparametrize_train_state``'s ``tie_weights=True``.
193 """
194 return {
195 **dict(mod.named_parameters(remove_duplicate=True)),
196 **dict(mod.named_buffers(remove_duplicate=True)),
197 }
200@contextlib.contextmanager
201def _reparametrize_train_state(
202 module: nn.Module,
203 model_state: Dict[str, torch.Tensor],
204):
205 """Rebind ``module``'s parameters/buffers to explicit trace-time tensors.
207 Inside the traced function this swaps in the fake/static tensors so every
208 parameter access in the model goes through the graph's static inputs.
210 ``tie_weights=True`` re-adds, for the duration of the trace, the alias
211 keys that ``extract_module_state`` deduplicated — mapping each alias to
212 the SAME static tensor. Both tied uses then reference one ``fx.Node``, so
213 the traced graph keeps a single placeholder (and a single gradient) for
214 the shared tensor instead of a baked-in ``get_attr`` constant.
215 """
216 with contextlib.ExitStack() as stack:
217 stack.enter_context(
218 stateless._reparametrize_module(module, model_state, tie_weights=True)
219 )
220 yield
223def _iter_node_args(node: torch.fx.Node):
224 """Yield every ``fx.Node`` among ``node.args`` (unwrapping nested lists)."""
225 for a in node.args:
226 if isinstance(a, torch.fx.Node):
227 yield a
228 elif isinstance(a, (tuple, list)):
229 for x in a:
230 if isinstance(x, torch.fx.Node):
231 yield x
234def _output_args(fx_g: torch.fx.GraphModule) -> List[torch.fx.Node]:
235 """Return the flattened forward arg list of the graph's ``output`` node."""
236 out_node = next((n for n in fx_g.graph.nodes if n.op == "output"), None)
237 if out_node is None:
238 return []
239 args: List[torch.fx.Node] = []
240 for a in out_node.args:
241 if isinstance(a, torch.fx.Node):
242 args.append(a)
243 elif isinstance(a, (tuple, list)):
244 for x in a:
245 if isinstance(x, torch.fx.Node):
246 args.append(x)
247 return args
250def _ancestors_of(node: torch.fx.Node) -> set:
251 """Return the transitive producers of ``node`` (excluding ``node``)."""
252 ancestors = set()
253 dq = deque([node])
254 while dq:
255 cur = dq.popleft()
256 for a in _iter_node_args(cur):
257 if a.op == "placeholder" or a in ancestors:
258 continue
259 ancestors.add(a)
260 dq.append(a)
261 return ancestors
264def _has_module_stack(node: torch.fx.Node) -> bool:
265 """True if the node carries any ``nn_module_stack`` entry.
267 make_fx leaves stack-less nodes (e.g. autograd engine output) with an
268 *empty* ``OrderedDict`` rather than ``None``, so a truthiness check is
269 the robust test.
270 """
271 return bool(node.meta.get("nn_module_stack"))
274def _annotate_autograd_backward(fx_g: torch.fx.GraphModule) -> None:
275 """Tag backward nodes with ``meta['autograd_backward'] = True``.
277 Stock torch's make_fx (without the torchtitan ``_patch_engine_backward``
278 hook) never attaches the autograd engine's ``autograd_backward`` tag, so
279 passes that split the joint graph into fwd/bwd halves get nothing to key
280 on. This reconstructs the tag structurally.
282 The joint graph's ``output`` returns ``[loss, grad_0, ... grad_n]``. The
283 loss and the nodes feeding it (logits, loss reduction, ``getitem`` of the
284 loss value) are forward even though they live outside any module and so
285 carry no ``nn_module_stack``. Every other output is a parameter gradient
286 produced by ``torch.autograd.grad``.
288 Two pre-computed sets separate the halves:
290 - ``fwd_loss_nodes``: the transitive producers of ``loss`` — the logits,
291 the shift/reshape that makes the loss target, and the loss reduction
292 itself. These are forward (they are the loss *source*); backward is
293 computed *from* them.
294 - ``bwd_nodes``: ``call_function`` nodes reachable from the parameter
295 gradients (``output[1:]``) that carry no ``nn_module_stack`` and are
296 not in ``fwd_loss_nodes``. The module stack is only recorded while a
297 traced module ``forward`` runs, so autograd-engine nodes have none.
298 Walking backward and stopping at forward boundaries (stack-bearing
299 nodes or the loss source) spans exactly the backward half.
300 """
301 output = _output_args(fx_g)
302 if not output:
303 return
304 loss_node = output[0]
305 grad_seeds = output[1:]
307 # Forward loss chain: everything that flows INTO the loss. Backward is
308 # computed from the loss, so these are forward regardless of whether any
309 # module recorded them (they may sit in ``train_fn``, outside modules).
310 fwd_loss_nodes = _ancestors_of(loss_node)
311 fwd_loss_nodes.add(loss_node)
313 # BFS backward over dataflow from the gradient outputs, marking every
314 # stack-less, non-loss-source node as backward. Stack-bearing nodes are
315 # forward (module boundaries); loss-source nodes are forward (the loss is
316 # the backward's *source*, not part of it).
317 bwd_nodes = set()
318 dq = deque()
319 for seed in grad_seeds:
320 if (
321 not _has_module_stack(seed)
322 and seed.op != "placeholder"
323 and seed not in fwd_loss_nodes
324 ):
325 bwd_nodes.add(seed)
326 dq.append(seed)
327 while dq:
328 n = dq.popleft()
329 for a in _iter_node_args(n):
330 if a.op == "placeholder" or a in bwd_nodes:
331 continue
332 if a in fwd_loss_nodes:
333 continue # forward loss source; do not expand
334 if _has_module_stack(a):
335 continue # forward module boundary; do not expand
336 bwd_nodes.add(a)
337 dq.append(a)
339 for node in bwd_nodes:
340 node.meta["autograd_backward"] = True
343def _copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
344 """Copy forward metadata to backward nodes across all nested FX subgraphs.
346 Uses a two-pass approach over all submodule graphs (including HOP
347 subgraphs). Pass 1 collects forward nodes by ``seq_nr``; pass 2 copies
348 ``custom``/``nn_module_stack``/``stack_trace`` from the matching forward
349 node to each backward node. Backward nodes are identified by the autograd
350 engine's ``autograd_backward`` tag on ``node.meta``.
351 """
352 seq_nr_to_fwd_node: Dict[int, torch.fx.Node] = {}
353 for node in _iter_seq_nr_nodes(fx_g, backward=False):
354 seq_nr_to_fwd_node.setdefault(node.meta["seq_nr"], node)
356 for node in _iter_seq_nr_nodes(fx_g, backward=True):
357 fwd_node = seq_nr_to_fwd_node.get(node.meta["seq_nr"])
358 if fwd_node is None or fwd_node is node:
359 continue
360 _copy_fwd_meta(fwd_node, node)
363def _iter_seq_nr_nodes(fx_g: torch.fx.GraphModule, backward: bool):
364 """Yield ``seq_nr``-bearing call/get_attr nodes across all subgraphs.
366 ``backward=False`` yields forward nodes, ``backward=True`` backward ones
367 (tagged by the autograd engine's ``autograd_backward`` meta).
368 """
369 for submod in fx_g.modules():
370 if not isinstance(submod, torch.fx.GraphModule):
371 continue
372 for node in submod.graph.nodes:
373 if (
374 node.op not in ("call_function", "get_attr")
375 or "seq_nr" not in node.meta
376 or node.meta.get("autograd_backward", False) != backward
377 ):
378 continue
379 yield node
382def _copy_fwd_meta(fwd_node: torch.fx.Node, node: torch.fx.Node) -> None:
383 """Copy ``custom``/``nn_module_stack``/``stack_trace`` from fwd to bwd."""
384 custom = fwd_node.meta.get("custom")
385 if custom:
386 node.meta.setdefault("custom", {}).update(copy.deepcopy(custom))
387 nn_module_stack = fwd_node.meta.get("nn_module_stack")
388 if nn_module_stack is not None:
389 node.meta["nn_module_stack"] = nn_module_stack.copy()
390 stack_trace = fwd_node.meta.get("stack_trace")
391 if stack_trace is not None:
392 node.meta["stack_trace"] = stack_trace
395def _fakeify_input(fake_mode: FakeTensorMode, x: Any) -> Any:
396 """Convert a real tensor input into its fake counterpart."""
397 if not isinstance(x, torch.Tensor):
398 return x
399 return fake_mode.from_tensor(x, static_shapes=False)
402def _input_meta(x: Any) -> Any:
403 """Return a storage-free descriptor of a traced user input.
405 Tensors become shape/dtype metadata; primitives pass through; pytree
406 containers (list/tuple/dict) are converted recursively. Used for
407 ``JointGraph.inputs`` so the joint graph never pins the real sample
408 tensors' storage for the lifetime of training.
409 """
410 if isinstance(x, torch.Tensor):
411 return {"shape": tuple(x.shape), "dtype": str(x.dtype)}
412 if isinstance(x, tuple):
413 return tuple(_input_meta(v) for v in x)
414 if isinstance(x, list):
415 return [_input_meta(v) for v in x]
416 if isinstance(x, dict):
417 return {k: _input_meta(v) for k, v in x.items()}
418 return x
421def _validate_trace_leaves(leaves: list[Any]) -> None:
422 """Reject state or input objects that cannot become explicit FX inputs."""
423 for leaf in leaves:
424 if isinstance(leaf, nn.Module):
425 raise ValueError(
426 "trace_model_graph requires explicit tensor state, not "
427 "nn.Module instances. Capture nn.Modules in train_fn's closure."
428 )
429 if not isinstance(leaf, _ALLOWED_LEAF_TYPES):
430 raise ValueError(
431 "trace_model_graph requires all pytree leaves in "
432 f"state/args to be tensors or primitives (int/float/bool/str), "
433 f"got {type(leaf).__name__}."
434 )
437def _make_fwd_bwd_fn(
438 model: nn.Module,
439 train_fn: Callable,
440 state_fqns: list[str],
441 state_is_param: list[bool],
442 state_spec: Any,
443 user_inputs_spec: Any,
444) -> tuple[Callable, list[str]]:
445 """Build the joint callable and its named-loss output metadata."""
446 num_state_inputs = len(state_fqns)
447 # ``_fwd_bwd_fn`` communicates the traced ``loss_dict`` key order to the
448 # enclosing scope through this box (plain closure writes cannot escape
449 # the nested ``def``); the keys are attached to the GraphModule below.
450 loss_dict_keys_box: list[str] = []
452 def _fwd_bwd_fn(*plain_args):
453 state_wrapped = plain_args[:num_state_inputs]
454 user_wrapped = plain_args[num_state_inputs:]
455 state_t = torch.utils._pytree.tree_unflatten(list(state_wrapped), state_spec)
456 user_args = torch.utils._pytree.tree_unflatten(
457 list(user_wrapped), user_inputs_spec
458 )
460 with (
461 _reparametrize_train_state(model, state_t["model"]),
462 _patch_engine_backward(),
463 ):
464 out = train_fn(model, **user_args)
466 # ``train_fn`` may return either a bare loss tensor or the
467 # ``(loss, loss_dict)`` pair produced by the trainer's
468 # ``postforward``. The named loss_dict values are emitted as
469 # graph outputs AFTER the gradients so passes that assume
470 # ``output[0] == loss`` and ``output[1:] == grads`` keep working
471 # once they trim the trailing loss_dict outputs.
472 if (
473 isinstance(out, tuple)
474 and len(out) == 2
475 and isinstance(out[1], dict)
476 ):
477 loss, loss_dict = out
478 loss_dict_keys = list(loss_dict.keys())
479 loss_dict_keys_box.extend(loss_dict_keys)
480 loss_dict_values = [loss_dict[key] for key in loss_dict_keys]
481 else:
482 loss, loss_dict_keys, loss_dict_values = out, [], []
484 # Read params from state_t["model"] so the exact trace-time
485 # tensors (not the live module's) drive autograd.grad.
486 params = []
487 for idx, fqn in enumerate(state_fqns):
488 if state_is_param[idx]:
489 param = state_t["model"][fqn]
490 if param.requires_grad:
491 params.append(param)
493 grads = torch.autograd.grad(loss, params, allow_unused=True)
495 # Replace None grads with zeros so the output arity is stable.
496 processed_grads = []
497 for grad, param in zip(grads, params):
498 if grad is None:
499 processed_grads.append(torch.zeros_like(param))
500 else:
501 processed_grads.append(grad)
503 return [loss] + processed_grads + loss_dict_values
505 return _fwd_bwd_fn, loss_dict_keys_box
508def trace_model_graph( # pylint: disable=too-many-locals
509 model: torch.nn.Module,
510 train_fn: Callable,
511 inputs: Dict[str, Any],
512 *,
513 dynamic: bool = False,
514 dynamic_arg_dims: dict[str, int | list[int]] | None = None,
515) -> JointGraph:
516 """
517 Trace model to generate complete forward + backward graph
519 Args:
520 model: Model (no parallel wrapping)
521 train_fn: Training function signature:
522 ``train_fn(model, **inputs) -> loss``
523 inputs: Model inputs, forwarded to ``train_fn`` as keyword arguments
524 dynamic: Symbolize all user tensor dimensions when no explicit mapping is supplied.
525 dynamic_arg_dims: Dotted input paths to dynamic dimensions (supports negative indices).
526 Providing a mapping enables selective symbolic tracing, even when dynamic is False.
527 Parameters and buffers always remain static.
529 Returns:
530 JointGraph: Joint forward-backward computation graph
532 Note:
533 Static tracing remains the default. Dynamic traces retain shape guards;
534 incompatible shapes or shape-dependent branches raise before execution.
535 """
536 # Extract module state (parameters/buffers) into flat tensors threaded
537 # through the graph as static inputs (leading placeholders).
538 model_state = extract_module_state(model)
539 state_fqns = list(model_state.keys())
540 # Which leading state inputs are parameters (vs buffers). FSDP shards
541 # parameters only: buffers -- e.g. the non-persistent RoPE ``cache`` --
542 # are full-rank by nature and must flow through as plain inputs, never
543 # all-gathered. ``state_is_param`` is indexed by state position, aligned
544 # with ``state_fqns``.
545 param_fqns = set(dict(model.named_parameters(remove_duplicate=False)).keys())
546 state_is_param = [fqn in param_fqns for fqn in state_fqns]
547 state_flat, state_spec = torch.utils._pytree.tree_flatten({"model": model_state})
549 # user_inputs is a plain dict so the traced closure unpacks it back into
550 # train_fn's keyword arguments.
551 user_inputs = inputs
552 user_inputs_flat, user_inputs_spec = torch.utils._pytree.tree_flatten(user_inputs)
554 _validate_trace_leaves([*state_flat, *user_inputs_flat])
556 full_args = list(state_flat) + list(user_inputs_flat)
558 if not isinstance(dynamic, bool):
559 raise ValueError("dynamic must be a bool")
560 symbolic = dynamic or dynamic_arg_dims is not None
561 if symbolic:
562 fake_mode, fake_args = build_symbolic_inputs(state_flat, inputs, dynamic_arg_dims)
563 else:
564 fake_mode = FakeTensorMode(allow_non_fake_inputs=True)
565 fake_args = tuple(_fakeify_input(fake_mode, value) for value in full_args)
566 num_state_inputs = len(state_flat)
568 _fwd_bwd_fn, loss_dict_keys_box = _make_fwd_bwd_fn( # pylint: disable=C0103
569 model, train_fn, state_fqns, state_is_param, state_spec, user_inputs_spec
570 )
572 # make_fx only records ``nn_module_stack`` when the traced callable
573 # carries ``_orig_mod``: its ``_init_modes_from_inputs`` then installs a
574 # ``_ModuleStackTracer`` that captures which module's ``forward`` was
575 # executing as each node is created. Forward nodes thereby gain the
576 # module-stack metadata ``PpPass._stage_from_stack`` needs to map a node
577 # to its pipeline stage. Backward nodes are annotated separately (see
578 # ``_annotate_autograd_backward``), after tracing.
579 _fwd_bwd_fn._orig_mod = model # type: ignore[attr-defined]
581 ctx = TracingContext(fake_mode)
582 with (
583 fake_mode,
584 tracing(ctx),
585 preserve_node_meta(),
586 _skip_nested_compile(),
587 torch.autograd.set_multithreading_enabled(False),
588 _non_strict_tracing_context(),
589 ):
590 # Native constant padding can specialize SymInt sizes (notably on
591 # NPU). Its upstream decomposition preserves causal-LM label lengths.
592 decompositions = get_decompositions([torch.ops.aten.constant_pad_nd.default]) if symbolic else None
593 traced_graph = make_fx(
594 _fwd_bwd_fn, decomposition_table=decompositions, **_MAKE_FX_KWARGS
595 )(*fake_args)
597 # Stock torch's make_fx does not tag backward nodes (no torchtitan
598 # ``_patch_engine_backward`` hook), so the joint graph splits into
599 # fwd/bwd halves with no ``autograd_backward`` meta to key on.
600 # ``_annotate_autograd_backward`` reconstructs the tag structurally so
601 # ``PpPass`` can classify both halves.
602 _annotate_autograd_backward(traced_graph)
604 # Copy forward metadata (nn_module_stack/stack_trace) to backward nodes so
605 # passes can match nodes by module FQN on both halves of the joint graph.
606 # On torchtitan's patched torch this is the primary source (the engine
607 # already tagged backward nodes); on stock torch it is a no-op because
608 # backward nodes carry no ``seq_nr`` that collides with a forward one.
609 _copy_fwd_metadata_to_bw_nodes(traced_graph)
611 # Expose the state layout on the GraphModule so downstream passes can
612 # identify which leading placeholders are parameters/buffers, and the
613 # trainer can feed runtime state in the same order.
614 traced_graph.state_fqns = state_fqns
615 traced_graph.state_is_param = state_is_param
616 traced_graph.num_state_inputs = num_state_inputs
617 # Kept for ``run_traced_graph`` to validate the runtime inputs' pytree
618 # structure against the traced one.
619 traced_graph.user_inputs_spec = user_inputs_spec
620 # Key order of the traced ``loss_dict``: its values are emitted as
621 # graph outputs AFTER the gradients, so ``run_traced_graph`` can split
622 # (loss, grads, loss_dict) and passes can trim the trailing loss_dict
623 # outputs off the gradient range.
624 traced_graph.loss_dict_keys = list(loss_dict_keys_box)
626 param_names = [name for name, _ in model.named_parameters() if _.requires_grad]
627 param_shapes = {
628 name: tuple(param.shape) for name, param in model.named_parameters()
629 }
631 num_layers = 0
632 if hasattr(model, "layers"):
633 num_layers = len(model.layers)
635 return JointGraph(
636 graph_module=traced_graph,
637 inputs=[_input_meta(v) for v in inputs.values()],
638 outputs=[],
639 param_names=param_names,
640 param_shapes=param_shapes,
641 num_layers=num_layers,
642 state_fqns=state_fqns,
643 example_inputs=fake_args,
644 input_guards=(
645 InputGuards(user_inputs_flat, fake_args[num_state_inputs:], fake_mode.shape_env)
646 if symbolic else None
647 ),
648 )
651def run_traced_graph(
652 joint_graph: JointGraph,
653 model: torch.nn.Module,
654 inputs: Dict[str, Any],
655) -> tuple:
656 """
657 Execute a traced joint graph against the live model state.
659 Parameters/buffers are sampled from ``model`` at call time and fed to the
660 graph as static inputs (in ``joint_graph.state_fqns`` order). ``inputs``
661 is flattened with the same pytree structure used at trace time (checked
662 against the stored spec) and appended after the state. Runs under
663 ``torch.no_grad()`` because the graph already contains the explicit
664 backward ops traced by ``torch.autograd.grad``.
666 FSDP is invisible here: FSDPPass shards ``model``'s parameters in place,
667 so ``model.parameters()`` yields the shards the graph expects as its
668 leading static inputs (the graph re-gathers them via AllGather each step).
670 Returns:
671 tuple: (loss, grads, loss_dict) where ``grads`` aligns with the
672 model's trainable parameter list and ``loss_dict`` maps the traced
673 loss names to their graph outputs (empty when ``train_fn`` returned
674 a bare loss tensor).
675 """
676 model_state = extract_module_state(model)
677 if list(model_state.keys()) != joint_graph.state_fqns:
678 raise ValueError(
679 "model has different parameter/buffer names than during tracing.\n"
680 f" Traced: {joint_graph.state_fqns}\n"
681 f" Got: {list(model_state.keys())}"
682 )
684 user_flat, user_spec = torch.utils._pytree.tree_flatten(inputs)
685 traced_spec = getattr(joint_graph.graph_module, "user_inputs_spec", None)
686 if traced_spec is not None and user_spec != traced_spec:
687 raise ValueError(
688 "model inputs have a different pytree structure than during "
689 "tracing (keys, nesting, or leaf types changed).\n"
690 f" Traced spec: {traced_spec}\n"
691 f" Got spec: {user_spec}"
692 )
694 state_flat, _ = torch.utils._pytree.tree_flatten({"model": model_state})
696 if joint_graph.input_guards is not None:
697 joint_graph.input_guards.validate(user_flat)
699 flat_inputs = list(state_flat) + list(user_flat)
701 with torch.no_grad():
702 outputs = joint_graph.graph_module(*flat_inputs)
704 loss_dict_keys = getattr(joint_graph.graph_module, "loss_dict_keys", [])
705 num_loss_outputs = len(loss_dict_keys)
706 if isinstance(outputs, (list, tuple)):
707 loss = outputs[0]
708 loss_dict = (
709 dict(zip(loss_dict_keys, outputs[len(outputs) - num_loss_outputs:]))
710 if num_loss_outputs
711 else {}
712 )
713 grads = (
714 list(outputs[1:len(outputs) - num_loss_outputs])
715 if num_loss_outputs
716 else list(outputs[1:])
717 )
718 else:
719 loss, grads, loss_dict = outputs, [], {}
720 return loss, grads, loss_dict