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