From ba572d739ce966cc0251810b8362e12998af84e5 Mon Sep 17 00:00:00 2001 From: Nurlan Nazaraliyev Date: Tue, 7 Jul 2026 14:21:56 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../graph_trainer/memory_estimator.py | 386 +++++++++++ .../graph_trainer/naive_autoAC_outer.py | 473 ++++++++++++++ .../graph_trainer/runtime_estimator.py | 611 ++++++++++++++++++ .../graph_trainer/transfertime_estimator.py | 276 ++++++++ 4 files changed, 1746 insertions(+) create mode 100644 torchtitan/experiments/graph_trainer/memory_estimator.py create mode 100644 torchtitan/experiments/graph_trainer/naive_autoAC_outer.py create mode 100644 torchtitan/experiments/graph_trainer/runtime_estimator.py create mode 100644 torchtitan/experiments/graph_trainer/transfertime_estimator.py diff --git a/torchtitan/experiments/graph_trainer/memory_estimator.py b/torchtitan/experiments/graph_trainer/memory_estimator.py new file mode 100644 index 0000000000..48afa85478 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/memory_estimator.py @@ -0,0 +1,386 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Static peak-memory estimator for a joint fwd+loss+bwd FX graph. + +``estimate_peak_memory`` sweeps the nodes of the joint graph produced by +``minimal_fx_tracer``, tracks per-storage liveness (birth/death over the node +schedule), and reports the peak as the maximum simultaneous live bytes -- plus a +per-category breakdown (parameter / activation / gradient / temporary) of what is +live at that peak. + +Liveness is storage-keyed: views/aliases that share an ``untyped_storage`` are +counted once. Parameters and buffers enter as the leading "state" placeholders; +they are pinned live for the whole step because the real allocator keeps them +resident even though the functionalized graph references each only once (a single +bf16 ``_to_copy`` cast). Optimizer state is not traced into this graph and is +accounted separately by ``optimizer_state_bytes``. +""" + +import math +from collections import Counter, defaultdict +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import torch.utils._pytree as pytree +from torch.fx.node import map_arg +from torchtitan.experiments.graph_trainer.common_utils import _is_backward_node +from torchtitan.experiments.graph_trainer.cpu_offload import _is_view +from torchtitan.experiments.graph_trainer.runtime_estimator import _is_costable_op + +if TYPE_CHECKING: + from torchtitan.components.optimizer import OptimizersContainer + +ROUNDING = 512 # CUDA caching allocator rounds small allocations to 512B. + +# Base optimizer states kept per parameter, by optimizer name (Adam/AdamW keep +# exp_avg + exp_avg_sq; amsgrad adds a third, max_exp_avg_sq -- handled below). +STATES_PER_PARAM = {"Adam": 2, "AdamW": 2} + +# ---- storage categories ---- +PARAM = "parameter" +GRAD = "gradient" +OPT = "optimizer_state" +ACT = "activation" # forward intermediate that survives into backward +TEMP = "temporary" # intermediate freed within its own (fwd or bwd) region +INPUT = "input" +BUFFER = "buffer" + + +@dataclass +class MemoryEstimatorResult: + peak_bytes: int + fwd_peak_bytes: int + bwd_peak_bytes: int + peak_node_index: int + peak_node_name: str + # category -> bytes live at the global peak point (sums to peak_bytes). + per_category_at_peak: dict + # category -> each category's own lifetime maximum (sum != peak_bytes). + per_category_independent_peak: dict + # category -> sum of bytes ever attributed to it (for sanity / debugging). + category_totals: dict + # storage key -> category (for debugging). + storage_category: dict + # live tensor bytes across the whole graph at each node index key: index, val: bytes + live_bytes: dict + # (size, birth, last_forward_use, death) for ilp solver + all_tensors: dict + # this is a dict for each time index -> live tensors + live_bytes_per_cat: dict + + def summary(self) -> str: + """Human-readable report: peak memory, the schedule point (node index + + name) where it occurs, and the per-category breakdown live at that point.""" + lines = [ + f"peak memory: {self.peak_bytes / 1e9:.3f} GB", + f"schedule point: node {self.peak_node_index} ({self.peak_node_name})", + "per-category at peak:", + ] + for cat, nbytes in sorted( + self.per_category_at_peak.items(), key=lambda kv: -kv[1] + ): + lines.append(f" {cat:16s} {nbytes / 1e9:7.3f} GB") + return "\n".join(lines) + + +def find_meaningful_last_fwd_use(node: torch.fx.Node, indeces_of_nodes): + memo = {} + lmf = _find_meaningful_last_fwd_use(node, indeces_of_nodes, memo) + return lmf if lmf >= 0 else indeces_of_nodes[node] + + +def _find_meaningful_last_fwd_use(node: torch.fx.Node, indeces_of_nodes, memo) -> None: + """Find the last fwd node that meaningfully uses given node's result. + Meaningful: not a view, not a no-op cast, not a no-op copy, transpose, etc. + One way to find this is to use `_is_costable_op` from runtime_estimator.py + """ + + def is_meaningful(node: torch.fx.Node) -> bool: + # A node "meaningfully" uses its input only if it is a real compute + # consumer. Views/reshapes/transposes (aliasing, no new allocation) and + # non-costable bookkeeping ops (getitem, the IGNORE_OPS set) do not pin + # the input, so we recurse through them. Casts (_to_copy, to.dtype, ...) + # allocate new storage and are treated as meaningful -- not recursing + # through them only overshoots the last forward use, which is safe + # (keeps the activation live slightly longer / offload gap smaller). + if not _is_costable_op(node) or _is_view(node): + return False + return True + + if node in memo: + return memo[node] + best = -1 # indeces_of_nodes[node] + for user in node.users: + if _is_backward_node(user): + continue + elif not is_meaningful(user): + best = max( + best, _find_meaningful_last_fwd_use(user, indeces_of_nodes, memo) + ) + else: # meaningful user + best = max(best, indeces_of_nodes[user]) # + memo[node] = best + return best + + +def estimate_peak_memory( + gm: torch.fx.GraphModule, + *, + num_state_inputs: int, + verbose: bool = False, +) -> MemoryEstimatorResult: + """Estimate the peak memory of a joint fwd+loss+bwd FX graph. + + Algorithm: for each storage (keyed by ``untyped_storage()._cdata``) compute a + birth (first node that produces it) and death (last node that uses it), then + sweep the schedule summing live bytes; the peak is the maximum. + + Categorization (by producer + lifetime): + - placeholder/get_attr pinned to the end (state) -> ``parameter`` + - backward producer surviving to a graph output -> ``gradient`` + - backward producer dying within backward -> ``temporary`` + - forward producer whose last use is backward -> ``activation`` + - forward producer whose last use is forward -> ``temporary`` + + ``num_state_inputs`` (``TracedResult.num_static_inputs``) is how many leading + placeholders are persistent state (parameters/buffers); they are pinned live + for the whole step. + """ + nodes = list(gm.graph.nodes) + index = {n: i for i, n in enumerate(nodes)} + val_of = lambda n: n.meta.get("val", None) # noqa: E731 + end = len(nodes) # "live to the end" sentinel for resident/returned storages + + def get_size(t: torch.Tensor) -> int: + # Count the underlying STORAGE, not the (possibly view) tensor's logical + # size: a narrow/slice/as_strided view has small numel but can reference a + # much larger storage. untyped_storage().nbytes() is the real allocation. + return math.ceil(t.untyped_storage().nbytes() / ROUNDING) * ROUNDING + + # Tensors returned to the caller (loss, grads) must live to the end. + output_inputs = set() + for node in nodes: + if node.op == "output": + output_inputs.update(node.all_input_nodes) + + # Parameters/buffers are the leading num_state_inputs placeholders; they stay + # resident on the GPU for the whole step even though the graph references each + # only once, so pin them live for the entire graph. + placeholders = [n for n in nodes if n.op == "placeholder"] + persistent_state = set(placeholders[:num_state_inputs]) + + # ---- storage-keyed birth / death / size ---- + # Key by (storage_id, birth_index) so a storage id that is freed and later + # reused starts a fresh interval instead of merging into one long lifetime. + birth, death, size, producer_of = {}, {}, {}, {} + live_key = {} + last_fwd_use, first_bwd_use = {}, {} + for i, node in enumerate(nodes): + # node.meta["val"] holds plain tensors -- minimal_fx_tracer unwraps tensor + # subclasses (e.g. DTensor) for tracing, so untyped_storage() is valid. For + # a subclass-carrying graph, use get_untyped_storages from + # torch.distributed._tools.common_utils instead. + for t in pytree.tree_leaves(node.meta.get("val")): + if not isinstance(t, torch.Tensor): + continue + # This estimates the GPU peak; CPU-resident tensors (e.g. the pinned + # host copies produced by ao.offload after apply_cpu_offload_pass) live + # in host DRAM, not GPU memory. Counting them would inflate the GPU peak + # by the offloaded byte total, which is exactly the offload savings. + if t.device.type != "cuda": + continue + sid = t.untyped_storage()._cdata + + key = live_key.get(sid) + if key is None or death.get(key, -1) < i: # new sid, or it died already + key = (sid, i) + live_key[sid] = key + birth[key] = i + size[key] = get_size(t) + producer_of[key] = node + + # Last use: the largest index among users that actually read THIS + # storage (a multi-output user may consume only some outputs). + d = i + for u in node.users: + u_inputs = pytree.tree_leaves( + (map_arg(u.args, val_of), map_arg(u.kwargs, val_of)) + ) + # for u_t in u_inputs: + # if ( + # isinstance(u_t, torch.Tensor) + # and u_t.untyped_storage()._cdata == sid + # ): + # d = max(d, index[u]) + reads_sid = any( # sid & key both in scope + isinstance(u_t, torch.Tensor) + and u_t.untyped_storage()._cdata == sid + for u_t in u_inputs + ) + if not reads_sid: + continue + ui = index[u] + d = max(d, ui) + if _is_backward_node(u): + first_bwd_use[key] = min(first_bwd_use.get(key, ui), ui) + else: + last_fwd_use[key] = max(last_fwd_use.get(key, ui), ui) + if node in output_inputs or node in persistent_state: + d = end # resident params/buffers and returned tensors live to end + death[key] = max(death.get(key, i), d) + + # ---- categorize each storage from its producer + lifetime ---- + category = {} + + for key in size: + (sid, i) = key + prod = producer_of[key] + # last_forward_use = max( + # [index[user] for user in prod.users if not _is_backward_node(user)] + # ) + # node_details[node] = (size[key], i, last_forward_use, death[key]) + if prod.op in ("placeholder", "get_attr"): + if prod in persistent_state or prod.op == "get_attr": + category[key] = PARAM + elif "tangent" in prod.name: + category[key] = GRAD # gradient seed + else: + category[key] = INPUT + elif _is_backward_node(prod): + category[key] = GRAD if death[key] >= end else TEMP + else: # forward-produced compute + last = nodes[min(death[key], end - 1)] + category[key] = ACT if _is_backward_node(last) else TEMP + + # (size, birth, last_forward_use, death) for ilp solver + all_tensors = {} + _lfu_memo: dict = {} + + def _refined_lfu(prod: torch.fx.Node) -> int: + lmf = _find_meaningful_last_fwd_use( + prod, index, _lfu_memo + ) # internal fn + shared memo + return lmf if lmf >= 0 else index[prod] # birth fallback here + + for key in size: + # if key in first_bwd_use and not _is_backward_node( + # producer_of[key] + # ): # fwd op read in bwd + v = producer_of[key] + all_tensors.setdefault(v, []).append( + { + "sid": key[0], + "category": category[key], + "size": size[key], + "birth": birth[key], # the first index produced this sid + "last_fwd_use": last_fwd_use.get( + key, birth[key] + ), # _refined_lfu(v), + "first_bwd_use": first_bwd_use.get(key, None), + # gap = [last_fwd_use, first_bwd_use) + "death": death[key], # the last index used this sid + "producer": v, # the node that produced this sid + } + ) + + # ---- O(n) sweep: peak, fwd/bwd peak, at-peak and per-category maxima ---- + add_at, free_at = defaultdict(list), defaultdict(list) + for key in size: + add_at[birth[key]].append(key) + free_at[death[key]].append(key) + + per_cat_live = Counter() + per_cat_at_peak, per_cat_independent = {}, {} + current = peak = peak_idx = fwd_peak = bwd_peak = 0 + # live bytes at each node index + live_bytes = {} + live_bytes_per_cat_curent = {} + live_bytes_per_cat = {} + for i in range(end + 1): # +1 so end-of-graph frees are processed + # key mean (sid, i) + + for key in add_at.get(i, ()): + # live_bytes_dict_of_nodes_curent[cat] += size[key] + current += size[key] + per_cat_live[category[key]] += size[key] + if i < end: + if _is_backward_node(nodes[i]): + bwd_peak = max(bwd_peak, current) + else: + fwd_peak = max(fwd_peak, current) + if current > peak: + peak, peak_idx = current, i + per_cat_at_peak = dict(per_cat_live) + + live_bytes[i] = current + live_bytes_per_cat[i] = dict(per_cat_live) + + for cat, live in per_cat_live.items(): + per_cat_independent[cat] = max(per_cat_independent.get(cat, 0), live) + for key in free_at.get(i, ()): + current -= size[key] + per_cat_live[category[key]] -= size[key] + + category_totals = Counter() + for key in size: + category_totals[category[key]] += size[key] + + result = MemoryEstimatorResult( + peak_bytes=peak, + fwd_peak_bytes=fwd_peak, + bwd_peak_bytes=bwd_peak, + peak_node_index=peak_idx, + peak_node_name=nodes[peak_idx].name, + per_category_at_peak=per_cat_at_peak, + per_category_independent_peak=per_cat_independent, + category_totals=dict(category_totals), + storage_category=category, + live_bytes=live_bytes, + all_tensors=all_tensors, + live_bytes_per_cat=live_bytes_per_cat, + ) + + if verbose: + print(result.summary()) + return result + + +def optimizer_state_bytes( + opt_config: "OptimizersContainer.Config | None", model: torch.nn.Module +) -> int: + """Persistent optimizer-state bytes, analytically from the optimizer config + (no optimizer instance / no training step needed). Adam/AdamW keep 2 state + tensors per param; states are fp32 unless the config requests bf16. + Returns 0 when there is no optimizer (e.g. the fwd/bwd-only test path).""" + if opt_config is None or not opt_config.param_groups: + return 0 + + # States-per-param per group: 2 for Adam/AdamW, +1 if amsgrad keeps + # max_exp_avg_sq. Take the max across groups -- exact for a uniform config; + # a mix of different per-group state counts would need per-group param + # attribution (precise FQN->group matching), which we don't do here. + def _states_per_param(pg) -> int: + n = STATES_PER_PARAM.get(pg.optimizer_name, 0) + if n and (getattr(pg, "optimizer_kwargs", None) or {}).get("amsgrad", False): + n += 1 + return n + + n_states = max((_states_per_param(pg) for pg in opt_config.param_groups), default=0) + if n_states == 0: + return 0 + + # bf16 optimizer states only under the fused_opt_states_bf16 implementation + # (Adam/AdamW momentum+variance in bf16); otherwise states are fp32. + state_dtype = ( + torch.bfloat16 + if opt_config.implementation == "fused_opt_states_bf16" + else torch.float32 + ) + elt_bytes = torch.finfo(state_dtype).bits // 8 + num_params = sum(p.numel() for p in model.parameters()) + return num_params * n_states * elt_bytes diff --git a/torchtitan/experiments/graph_trainer/naive_autoAC_outer.py b/torchtitan/experiments/graph_trainer/naive_autoAC_outer.py new file mode 100644 index 0000000000..2fc80e669d --- /dev/null +++ b/torchtitan/experiments/graph_trainer/naive_autoAC_outer.py @@ -0,0 +1,473 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Two-level per-tensor keep/recompute/offload ILP (work in progress). + +Plan (mirrors torch's sac_milp two-level decomposition, extended with offload): + + Outer ILP -- budget allocation across transformer blocks. One variable per + block (a GPU-keep budget), coupled by the global peak-memory constraint. + Inner ILP -- per-block three-way keep/recompute/offload. Given its allocated + budget, each block solves a small independent ILP. + +This file currently implements the pieces the outer ILP builds on: + * Step 1: group graph nodes by transformer block (block_of_node / + group_nodes_by_block). + * Step 2: the inner ILP (solve_inner_block) and the per-block tradeoff CURVE + (build_block_curve) -- added_time as a convex piecewise-linear function of + kept activation bytes. Structurally identical blocks are solved ONCE and + reused (build_all_block_curves). + +The outer budget-allocation ILP and graph tagging are TODO. +""" + +import copy +from collections import defaultdict +from dataclasses import dataclass, field + +import torch +from pulp import ( + LpBinary, + LpMinimize, + LpProblem, + LpStatus, + lpSum, + LpVariable, + PULP_CBC_CMD, +) +from torch.utils.checkpoint import CheckpointPolicy +from torchtitan.experiments.graph_trainer.common_utils import ( + _is_backward_node, + _MODULE_FQN, +) +from torchtitan.experiments.graph_trainer.cpu_offload import apply_cpu_offload_pass +from torchtitan.experiments.graph_trainer.make_fx_tracer import TracedResult +from torchtitan.experiments.graph_trainer.memory_estimator import ( + ACT, + estimate_peak_memory, + optimizer_state_bytes, +) +from torchtitan.experiments.graph_trainer.runtime_estimator import ( + COST_MODEL, + INTERPRETER, + RuntimeEstimator, + RuntimeEstimatorResult, +) +from torchtitan.experiments.graph_trainer.selective_activation_remat import ( + selective_activation_remat_pass, +) +from torchtitan.experiments.graph_trainer.transfertime_estimator import ( + _transfer_ms, + get_transfer_bw, +) +from torchtitan.tools.logging import logger + +# meta["recompute"] tag for each per-tensor policy decision. +_POLICY_TAG = { + "keep": CheckpointPolicy.MUST_SAVE, + "recompute": CheckpointPolicy.MUST_RECOMPUTE, + "offload": CheckpointPolicy.MUST_CPU_OFFLOAD, +} + +# Scale all memory terms to GiB inside the ILP (sac_milp's MEM_MULTIPLIER). Raw +# byte coefficients (~1e10) mixed with runtimes (~1e2) give CBC a coefficient +# spread that breaks its tolerances and yields constraint-violating "Optimal" +# solutions; GiB units keep every coefficient O(1-30). +MEM_MULTIPLIER = 1 << 30 + +# Offload only tensors at least this large: fewer, larger transfers hold less +# per-tensor overhead at the backward peak (matches the sac_and_offload default). +OFFLOAD_MIN_BYTES = 1 << 20 # 1 MiB + + +def _is_rng_op(node: torch.fx.Node) -> bool: + """RNG ops cannot be replayed by the remat pass, so they must never be + recomputed (they may still be kept or offloaded).""" + return torch.Tag.nondeterministic_seeded in getattr(node.target, "tags", set()) + + +# --------------------------------------------------------------------------- +# Step 1: group nodes by transformer block +# --------------------------------------------------------------------------- +def block_of_node(node: torch.fx.Node) -> str | None: + """Return the transformer-block FQN a node belongs to, e.g. "layers.3", + or None for nodes outside a block (embeddings, norm, lm_head, loss, ...). + + The block is the "layers." prefix of the node's module FQN tag. Nodes + without that prefix are not block-owned and are handled as forced-keep + (they contribute to the fixed baseline, not to any block's decisions). + """ + fqn = node.meta.get("custom", {}).get(_MODULE_FQN, "") + parts = fqn.split(".") + if len(parts) >= 2 and parts[0] == "layers" and parts[1].isdigit(): + return f"layers.{parts[1]}" + return None + + +def group_nodes_by_block(gm: torch.fx.GraphModule) -> dict[str, list[torch.fx.Node]]: + """Map each transformer-block FQN to its forward call_function nodes.""" + blocks: dict[str, list[torch.fx.Node]] = defaultdict(list) + for node in gm.graph.nodes: + if node.op != "call_function" or _is_backward_node(node): + continue + block = block_of_node(node) + if block is not None: + blocks[block].append(node) + return dict(blocks) + + +# --------------------------------------------------------------------------- +# outer ilp over all layers + inner ilp per layer +# --------------------------------------------------------------------------- +def plan_and_tag( + trace: TracedResult, + memory_budget: int, + optimizer, + model_parts: list[torch.nn.Module], + runtime_estimation_mode: str = COST_MODEL, + cpu_offload_budget_gb: float = 100.0, + interp_ctx: tuple | None = None, # (model, *run_args) for INTERPRETER mode +) -> torch.fx.GraphModule | None: + """Two-level keep/recompute/offload solver. Tags ``gm`` in place and returns + it (the caller runs apply_cpu_offload_pass + selective_activation_remat_pass). + + Steps: build per-block tradeoff curves -> outer budget allocation -> tag -> + measure the REAL post-pass peak on a clone. Because the modeled peak is only + a proxy (recompute working set omitted; peak-node approximation), we sweep the + outer cap over ``cap_grid`` multipliers of eff_budget, measure each plan's + real peak, and pick the one whose real peak is CLOSEST TO the budget from + below -- i.e. keep the most / free the least -> least added runtime. Returns + the untouched ``gm`` when the budget already fits, or None when no block + activations exist / no feasible plan is found. + """ + gm = trace.gm + mem_est = estimate_peak_memory(gm, num_state_inputs=trace.num_static_inputs) + opt_bytes = optimizer_state_bytes(optimizer, model_parts[0]) + estimated = mem_est.peak_bytes + opt_bytes + if memory_budget > estimated: + logger.info( + "new-autoAC: budget %.2f GB >= estimated peak %.2f GB; nothing to do", + memory_budget / 1e9, + estimated / 1e9, + ) + return gm + + import os as _os + + runtime_estimation_mode = _os.environ.get( + "AUTOAC_RT_MODE", runtime_estimation_mode + ) # diagnostic override: cost-model | benchmark | interpreter + logger.info("new-autoAC: runtime estimation mode = %s", runtime_estimation_mode) + if runtime_estimation_mode == INTERPRETER: + if interp_ctx is None: + raise ValueError( + "INTERPRETER runtime mode needs interp_ctx=(model, *run_args); " + "pass it from the trainer (see AUTOAC_MODE=scratch)." + ) + runtime = RuntimeEstimator()(INTERPRETER).estimate(trace, *interp_ctx) + else: + runtime = RuntimeEstimator()(runtime_estimation_mode).estimate(trace) + runtime_per_node = runtime.node_runtimes_ms + + blocks = defaultdict(list) + nodes = list(gm.graph.nodes) + for node in nodes: + b = block_of_node(node) + if b is not None: + blocks[b].append(node) + + # the outer ILP should make decision per module for shared resources: + # such as the GPU peak memory, CPU memory, PCIe bw + # simply, how much each layer should keep + + block_names = list(blocks) # <- index i <-> block_names[i] + num_blocks = len(block_names) + + _block_act = block_activation_bytes(mem_est) # freeable activation bytes per block + block_act = { + b: _block_act.get(b, 0) for b in block_names + } # ensure every block present + bw = get_transfer_bw() + bw_d2h, bw_h2d = bw["d2h"] * 1e6, bw["h2d"] * 1e6 # GB/s -> bytes/ms + + # per-block forward (recompute cost) and backward (window) runtimes + fwd_rt, bwd_rt = defaultdict(float), defaultdict(float) + for block, blk_nodes in blocks.items(): + for n in blk_nodes: + if _is_backward_node(n): + bwd_rt[block] += runtime_per_node.get(n.name, 0.0) + else: + fwd_rt[block] += runtime_per_node.get(n.name, 0.0) + + order = sorted( + block_names, key=lambda b: int(b.split(".")[1]) + ) # layers.0,1,...,L-1 + d2h_window, h2d_window = {}, {} + + suf_fb, suf_b = 0.0, 0.0 + for j in reversed(order): # last layer -> first + d2h_window[j] = suf_fb # blocks AFTER j (fwd only) # (fwd+bwd) + h2d_window[j] = suf_b # blocks AFTER j (bwd only) + suf_fb += fwd_rt[j] # + bwd_rt[j] + suf_b += bwd_rt[j] + + blocks_fwd_end_bwd_start = {j: [0, 0] for j in order} + time_f = 0.0 + for j in order: # layers.0, 1, ..., L-1 + time_f += fwd_rt[j] + blocks_fwd_end_bwd_start[j][0] = suf_fb + + time_b = time_f + for j in reversed(order): # layers.L-1, ..., 1, 0 + blocks_fwd_end_bwd_start[j][1] = suf_b + time_b += bwd_rt[j] + + # --- outer three-way ILP: per-block keep/recompute/offload fractions --- + # Only ACTIVATION bytes are freeable; params/grads/temp/non-block activations + # form the fixed resident baseline at the peak. + fixed = mem_est.peak_bytes - sum(block_act.values()) + eff_budget = memory_budget - opt_bytes + t_fwd_total = sum(fwd_rt[b] for b in block_names) # D2H overlap capacity (ms) + t_bwd_total = sum(bwd_rt[b] for b in block_names) # H2D overlap capacity (ms) + + alloc = solve_outer_three_way( + mem_est, + block_act, + {b: fwd_rt[b] for b in block_names}, + d2h_window, + h2d_window, + blocks_fwd_end_bwd_start, + order, + time_f, + time_b, + bw_d2h, + bw_h2d, + t_fwd_total, + t_bwd_total, + fixed, + eff_budget, + cpu_budget_bytes=100, + ) + if alloc is None: + logger.warning( + "new-autoAC: outer ILP infeasible -- fixed baseline %.2f GiB exceeds " + "eff_budget %.2f GiB (freeing all activation still cannot fit)", + fixed / (1 << 30), + eff_budget / (1 << 30), + ) + return None + + # --- report the per-block allocation --- + GiB = 1 << 30 + total_act = sum(block_act.values()) + keep_b = sum(k * block_act[b] for b, (k, _r, _o) in alloc.items()) + rec_b = sum(r * block_act[b] for b, (_k, r, _o) in alloc.items()) + off_b = sum(o * block_act[b] for b, (_k, _r, o) in alloc.items()) + logger.info( + "new-autoAC outer (3-way): %d blocks | budget=%.2f eff=%.2f fixed=%.2f " + "opt=%.2f GiB | act=%.2f GiB -> keep=%.2f recompute=%.2f offload=%.2f GiB | " + "modeled peak=%.2f GiB", + num_blocks, + memory_budget / GiB, + eff_budget / GiB, + fixed / GiB, + opt_bytes / GiB, + total_act / GiB, + keep_b / GiB, + rec_b / GiB, + off_b / GiB, + (fixed + keep_b) / GiB, + ) + logger.info( + "new-autoAC per-block (keep/recompute/offload frac):\n%s", + "\n".join( + f" {b}: k={alloc[b][0]:.2f} r={alloc[b][1]:.2f} o={alloc[b][2]:.2f} " + f"(act={block_act[b] / GiB:.3f} GiB, d2h_win={d2h_window[b]:.1f}ms " + f"h2d_win={h2d_window[b]:.1f}ms)" + for b in order + ), + ) + + # Outer ILP only: this returns the per-block fractions via the log above. + # Materialization (inner ILP -> node.meta["recompute"] tags) is a separate step. + # Inner ILP is to be implemented later + return None + + +# --------------------------------------------------------------------------- +# Outer three-way ILP: per-block keep/recompute/offload FRACTIONS +# --------------------------------------------------------------------------- +def solve_outer_three_way( + mem_est, + act_bytes: dict[str, int], # block -> freeable ACTIVATION bytes (activation only) + fwd_rt: dict[str, float], # block -> forward runtime ms (recompute cost) + d2h_window: dict[str, float], # block -> D2H overlap window ms (evict gap) + h2d_window: dict[str, float], # block -> H2D overlap window ms (reload gap) + blocks_fwd_end_bwd_start: dict[str, float], # block -> fwd end -> bwd start ms + ordered_blocks: list[str], # order of blocks: layers.0,1,...,L-1 + time_fwd, # total forward runtime ms + time_bwd, # total backward runtime ms + bw_d2h: float, # D2H bandwidth, bytes/ms + bw_h2d: float, # H2D bandwidth, bytes/ms + t_fwd_total: float, # total forward compute ms (D2H overlap capacity) + t_bwd_total: float, # total backward compute ms (H2D overlap capacity) + fixed_bytes: float, # non-freeable resident bytes at the peak + eff_budget_bytes: float, # memory_budget - opt_bytes + cpu_budget_bytes: float = 100, + *, + keep_eps: float = 1e-3, # ms/GiB tie-break: prefer keep > offload > recompute + time_limit: int = 120, +) -> dict[str, tuple[float, float, float]] | None: + """Allocate a per-block keep/recompute/offload split (fractions summing to 1) + minimizing added runtime under the global peak budget. + + Decision per block b: r_b (recompute frac), o_b (offload frac); keep is the + residual k_b = 1 - r_b - o_b. Both recompute and offload evict the activation + from GPU during its gap, so only KEPT activation is resident at the peak. All + memory terms are in GiB, all time in ms. + + (C1) peak: fixed + sum_b k_b*act_b <= eff_budget (only kept resident) + (C2) window: all act accumulated <= d2h window (evict fits its gap) + same for h2d window (reload fits its gap) + -> some penalty for unhidden transfer with stall: + (C3) shared: d2h_stall >= (sum_b o_b*act_b)/bw_d2h - t_fwd (contention: all + h2d_stall >= (sum_b o_b*act_b)/bw_h2d - t_bwd offloads share one + D2H and one H2D engine; bytes past the compute window stall). + objective: min sum_b r_b*fwd_rt[b] + d2h_stall + h2d_stall + + keep_eps*sum_b (r_b+o_b)*act_b + + Offload that fits both its window (C2) and the shared budget (C3) is modeled as + fully hidden (zero added time), so the solver prefers keep (if the budget fits), + then offload, then recompute. The keep_eps term breaks the free-offload tie + toward keep so a block that already fits the budget is not needlessly freed. + + Returns {block: (keep, recompute, offload)} fractions, or None if infeasible + (the fixed baseline alone exceeds the budget -- freeing all activation can't fit). + """ + M = MEM_MULTIPLIER + blocks = list(act_bytes) + prob = LpProblem("outer_three_way", LpMinimize) + + def san(b: str) -> str: + return b.replace(".", "_") + + r = {b: LpVariable(f"r_{san(b)}", 0.0, 1.0) for b in blocks} + o = {b: LpVariable(f"o_{san(b)}", 0.0, 1.0) for b in blocks} + for b in blocks: + prob += r[b] + o[b] <= 1.0, f"split_{san(b)}" # keep = 1 - r - o >= 0 + + act_gib = {b: act_bytes[b] / M for b in blocks} + bw_d2h_gib = bw_d2h / M # GiB/ms + bw_h2d_gib = bw_h2d / M + + # (C1') per-position peak: evaluate memory at each block b's backward step. + # Block b's activation is ALWAYS present there (kept / recomputed-now / reloaded- + # now); earlier blocks k window = fwd time after j. + # H2D: they must arrive before block j's bwd -> window = bwd time after j. + # One constraint family, shared o-variables = engine contention, linearly. + # same but written differently and correctly + for idx, j in enumerate(ordered_blocks): + suffix_off = lpSum(o[i] * act_gib[i] for i in ordered_blocks[idx:]) + prob += ( + # here we allow offloading go beyond the fwd-bwd boundary with stall but the memory peak can also increase? + # suffix_off <= bw_d2h_gib * (d2h_window[j] + d2h_stall), + suffix_off + <= bw_d2h_gib + * ( + d2h_window[j] + ), # here we don't allow offloading go beyond the fwd-bwd boundary + f"d2h_suffix_{san(j)}", + ) + prob += ( + suffix_off <= bw_h2d_gib * h2d_window[j] + bw_h2d_gib * h2d_stall, + f"h2d_suffix_{san(j)}", + ) + + prob += ( + lpSum(r[b] * fwd_rt[b] for b in blocks) + + d2h_stall + + h2d_stall + + keep_eps * lpSum((r[b] + o[b]) * act_gib[b] for b in blocks) + ), "added_runtime" + + prob += (lpSum(o[b] * act_gib[b] for b in blocks) <= cpu_budget_bytes), "cpu_budget" + + status = prob.solve(PULP_CBC_CMD(gapRel=0.02, timeLimit=time_limit, msg=0)) + if LpStatus[status] != "Optimal": + return None + + alloc: dict[str, tuple[float, float, float]] = {} + for b in blocks: + rv = r[b].value() or 0.0 + ov = o[b].value() or 0.0 + alloc[b] = (max(0.0, 1.0 - rv - ov), rv, ov) # (keep, recompute, offload) + return alloc + + +def block_activation_bytes(mem_est): + + totals, seen = defaultdict(int), set() + + for producer, entries in mem_est.all_tensors.items(): + b = block_of_node(producer) + if b is None: + continue + for e in entries: + if e["category"] != ACT or e["first_bwd_use"] is None: + continue # not a saved-for-backward activation + if e["sid"] in seen: # dedup: one variable per storage id + continue + seen.add(e["sid"]) + totals[b] += e["size"] # bytes (shape x dtype, computed by the estimator) + return dict(totals) + + +def block_fixed_bytes(mem_est) -> dict[str, int]: + totals, seen = defaultdict(int), set() + for producer, entries in mem_est.all_tensors.items(): + b = block_of_node(producer) + if b is None: + continue + for e in entries: + if e["category"] == ACT: + continue # a saved-for-backward activation + if e["sid"] in seen: # dedup: one variable per storage id + continue + seen.add(e["sid"]) + totals[b] += e["size"] # bytes (shape x dtype, computed by the estimator) + return dict(totals) diff --git a/torchtitan/experiments/graph_trainer/runtime_estimator.py b/torchtitan/experiments/graph_trainer/runtime_estimator.py new file mode 100644 index 0000000000..d6aa8c7441 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/runtime_estimator.py @@ -0,0 +1,611 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Static / measured runtime estimator for a joint fwd+loss+bwd FX graph. + +This mirrors the API/structure of +``torch.distributed._tools.runtime_estimator.RuntimeEstimator`` -- the mode is +selected via ``__call__``: + + - ``operator-level-cost-model``: a roofline estimate, ``max(compute, HBM transfer)`` + using the same ``torch.utils._runtime_estimation`` helpers the upstream + estimator uses. Static (no GPU execution). + - ``operator-level-benchmark``: runs and times each op on real randomly-initialized + tensors (warmup + timed iters, CUDA events) -- the same per-op micro-benchmark idea + as upstream's ``_maybe_run_and_benchmark_fallback_kernel``. Each op is benchmarked in + isolation, so it works even for a model too big to run end-to-end. + - ``operator-level-interpreter``: executes the *whole* real traced graph via + ``run_traced`` and times it -- the real end-to-end total (overlap-aware) plus + a per-node CUDA-event breakdown from an ``fx.Interpreter``. Needs the live + module and real run inputs, and cudagraph disabled. + +It also aggregates per-module forward/backward runtimes and records module +execution order, exposed via ``display_modulewise_stats``. + +Difference from upstream: upstream is a ``TorchDispatchMode`` run under +``FakeTensorMode`` that intercepts aten ops as an eager model executes. The +cost-model/benchmark modes here walk the joint graph statically (per FX node, +the granularity our save/recompute/offload solver needs), reading each node's +fake output (``node.meta['val']``) -- exactly the inputs ``__torch_dispatch__`` +would see -- and reuse the same cost-model helpers. + +Limitations (shared with the upstream cost model): + 1. Communication/collectives are not modeled (compute-only estimate). + 2. Like upstream, we cost every op and eliminate only _IGNORE_OPS (plus + getitem/structural nodes via _is_costable_op). An op the flop registry + doesn't know is NOT dropped to zero -- its transfer_time is always counted; + only its compute_time falls back to 0. So a fused/opaque op (e.g. + flex_attention after regional_inductor) still contributes its tensor I/O, + just not its hidden compute. + 3. Roofline ignores launch overhead, occupancy, and overlap. Benchmark and + interpreter modes capture real kernel time; the interpreter total is + overlap-aware, but its per-node breakdown is serialized. + +Note: The interpreter mode is the only one that can capture the real end-to-end. +""" + +import operator +from collections import defaultdict +from dataclasses import dataclass, field +from statistics import median + +import torch +import torch.utils._pytree as pytree +from torch.fx.node import map_arg +from torch.utils._runtime_estimation import ( + _FLOAT_TYPES, + _IGNORE_OPS, + _VIEW_OPS, + get_compute_time, + get_transfer_time, +) +from typing_extensions import Self + +from torchtitan.experiments.graph_trainer.common_utils import ( + _is_backward_node, + _MODULE_FQN, +) +from torchtitan.experiments.graph_trainer.make_fx_tracer import run_traced, TracedResult + +# Estimate modes, named to match the upstream RuntimeEstimator. +COST_MODEL = "operator-level-cost-model" +BENCHMARK = "operator-level-benchmark" +INTERPRETER = "operator-level-interpreter" + + +@dataclass +class RuntimeEstimatorResult: + """Per-node + per-module runtime of the graph as written.""" + + total_runtime_ms: float + fwd_runtime_ms: float + bwd_runtime_ms: float + estimate_mode_type: str = COST_MODEL + # node name -> estimated/measured time in ms + node_runtimes_ms: dict = field(default_factory=dict) + # module fqn -> {"fw": ms, "bw": ms} + mod_runtimes_ms: dict = field(default_factory=dict) + + def summary(self, top_k: int = 10) -> str: + lines = [ + f"runtime ({self.estimate_mode_type}): {self.total_runtime_ms:.3f} ms " + f"(fwd {self.fwd_runtime_ms:.3f} ms, bwd {self.bwd_runtime_ms:.3f} ms)", + # Attention (flex_attention HOP / aten SDPA) IS costed. Collectives + # are not, so treat this as a compute-only estimate (no comm/overlap). + "NOTE: collectives are not costed (compute-only estimate).", + f"top {top_k} ops by time:", + ] + for name, t in sorted(self.node_runtimes_ms.items(), key=lambda kv: -kv[1])[ + :top_k + ]: + lines.append(f" {t:8.4f} ms {name}") + return "\n".join(lines) + + +class _TimingInterpreter(torch.fx.Interpreter): + """Times each compute node with CUDA events (serialized; per-op cost table).""" + + def __init__(self, gm: torch.fx.GraphModule): + super().__init__(gm) + self.events: dict = {} + # self._ignored_time = 0.0 + + def run_node(self, n): + if n.op in ("placeholder", "output", "get_attr"): + return super().run_node(n) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + out = super().run_node(n) + # torch.cuda.synchronize() + end.record() + self.events[n.name] = (n, start, end) + return out + + def timings_ms(self) -> dict: + # One sync, THEN read every elapsed_time (events must be completed first). + torch.cuda.synchronize() + + # for name, (_, s, e) in self.events.items(): + # if ( + # "squeeze" in name + # or "transpose" in name + # or "expand" in name + # or "t_" in name + # or "split" in name + # or "zeros_like" in name + # or "ones_like" in name + # ): + # self._ignored_time += s.elapsed_time(e) + # print(f"ignored time: {self._ignored_time} ms") + return {name: s.elapsed_time(e) for name, (_, s, e) in self.events.items()} + + +def _is_costable_op(node: torch.fx.Node) -> bool: + """Whether to apply the roofline to this node. + + Mirrors the upstream RuntimeEstimator: cost *every* op, eliminating only + ``_IGNORE_OPS`` (and ``getitem``, which just unpacks multi-output ops). We do + NOT drop an op to zero merely because it is "uncostable": ``_roofline_estimate`` + always counts ``transfer_time`` from the node's tensors, and ``get_compute_time`` + falls back to 0 for targets the flop registry doesn't know. So an unknown op + (a non-registry HOP, a fused inductor kernel, a custom op) still contributes + its memory-transfer cost rather than nothing. Known compute ops (aten GEMMs, + the ``flex_attention`` HOP pre-fusion, etc.) additionally get ``compute_time``. + """ + if node.op != "call_function": + return False + target = node.target + # operator.getitem just indexes a multi-output op's result; no kernel cost. + if target is operator.getitem: + return False + if isinstance(target, torch._ops.OpOverload): + return target._overloadpacket not in _IGNORE_OPS + # HOPs / other call_functions: cost their transfer time (compute falls back + # to 0 inside get_compute_time if the flop registry doesn't know them). + return True + + +def _module_fqn(node: torch.fx.Node) -> str: + return node.meta.get("custom", {}).get(_MODULE_FQN, "") + + +def _capture(traced: TracedResult, module, args, build): + """Run the traced graph with a custom interpreter built by ``build(gm)`` and + return that interpreter instance (run_traced doesn't expose it otherwise).""" + holder: dict = {} + + def factory(gm): + inst = build(gm) + holder["inst"] = inst + return inst + + run_traced(traced, module=module, interpreter_cls=factory)(*args) + return holder["inst"] + + +class RuntimeEstimator: + """FX-graph-based runtime estimator mirroring the upstream + ``torch.distributed._tools.runtime_estimator.RuntimeEstimator`` API. + + Static modes (cost-model / benchmark) need only the graph:: + + est = RuntimeEstimator()(estimate_mode_type="operator-level-cost-model") + result = est.estimate(gm) # gm or TracedResult + + The interpreter mode runs the real graph, so it needs the live module and + real inputs (and cudagraph disabled):: + + est = RuntimeEstimator()(estimate_mode_type="operator-level-interpreter") + result = est.estimate(traced, module, *run_args) + + Then ``est.display_modulewise_stats(depth=2)``. + """ + + def __init__(self, warmup_iters: int = 2, bench_iters: int = 3) -> None: + self._estimate = self._roofline_estimate + self._estimate_mode_type = COST_MODEL + self._warmup_iters = warmup_iters + self._bench_iters = bench_iters + # module fqn -> {"fw": ms, "bw": ms} + self.mod_runtimes: dict[str, dict[str, float]] = defaultdict( + lambda: defaultdict(float) + ) + self.mod_fw_order: list[str] = [] + self.mod_bw_order: list[str] = [] + self.total_runtime: float = 0.0 + self.fwd_runtime: float = 0.0 + self.bwd_runtime: float = 0.0 + + # ---- mode selection (mirrors upstream __call__) ---- + def __call__(self, estimate_mode_type: str) -> Self: + if estimate_mode_type == COST_MODEL: + self._estimate = self._roofline_estimate + elif estimate_mode_type == BENCHMARK: + self._estimate = self._benchmark_estimate + elif estimate_mode_type == INTERPRETER: + # Interpreter times the whole graph at once, not per node, so there + # is no per-node estimate function for it. + self._estimate = None + else: + raise NotImplementedError( + f"estimate_mode_type {estimate_mode_type} not supported. Supported: " + f"{COST_MODEL}, {BENCHMARK}, {INTERPRETER}" + ) + self._estimate_mode_type = estimate_mode_type + return self + + # Adapted from: https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tools/runtime_estimator.py + # ---- per-node estimate methods (static modes) ---- + def _roofline_estimate(self, node: torch.fx.Node) -> float: + """Roofline time for a single node, in ms: ``max(compute, transfer)``. + + Reconstructs (args, kwargs, out) the cost model expects by substituting + each input ``Node`` with its fake value (``meta['val']``).""" + + if not torch.accelerator.is_available(): + raise AssertionError( + "Roofline estimation needs to access CUDA capabilities to make estimations" + ) + + val_of = lambda n: n.meta.get("val", None) # noqa: E731 + args = map_arg(node.args, val_of) + kwargs = map_arg(node.kwargs, val_of) + out = node.meta.get("val", None) + + flat_args_kwargs = pytree.tree_leaves((args, kwargs)) + flat_outs = pytree.tree_leaves(out) + transfer_time = get_transfer_time(flat_args_kwargs, flat_outs) + + target = node.target + func_packet = ( + target._overloadpacket + if isinstance(target, torch._ops.OpOverload) + else target + ) + + if func_packet in _IGNORE_OPS: + # print(f"ignoring {node.name} {func_packet}") + return 0.0 + + # get_compute_time asserts exactly one output dtype; multi-output flop ops + # (SDPA family, flex_attention return out=bf16 + logsumexp=fp32) would trip + # it, so cost against the primary (first float) output dtype. + primary_dtype = next( + ( + t.dtype + for t in flat_outs + if isinstance(t, torch.Tensor) and t.dtype in _FLOAT_TYPES + ), + None, + ) + compute_time = 0.0 + if primary_dtype is not None: + try: + compute_time = get_compute_time( + func_packet, args, kwargs, out, {primary_dtype}, node_meta=node.meta + ) + except Exception: + compute_time = 0.0 + return max(transfer_time, compute_time) / 1e6 + + def _benchmark_estimate(self, node: torch.fx.Node) -> float: + """Real per-op time (ms): materialize the node's inputs as real random + tensors and time the kernel (warmup + timed iters, CUDA events). + + Mirrors upstream's ``_maybe_run_and_benchmark_fallback_kernel``. Falls + back to the roofline for HOPs, view/inplace-view ops, on CPU, or when the + op errors on random inputs.""" + target = node.target + # HOPs (flex_attention) cannot be micro-benchmarked in isolation. + if not isinstance(target, torch._ops.OpOverload): + return self._roofline_estimate(node) + # Views/create ops are ~free and would not run cleanly on random inputs. + if target._overloadpacket in _VIEW_OPS: + return 0.0 + if torch.Tag.inplace_view in getattr(target, "tags", ()): + return 0.0 + if not torch.cuda.is_available(): + return self._roofline_estimate(node) + try: + real_args = map_arg(node.args, self._real_leaf) + real_kwargs = map_arg(node.kwargs, self._real_leaf) + target(*real_args, **(real_kwargs or {})) # correctness + allocate + for _ in range(self._warmup_iters): + target(*real_args, **(real_kwargs or {})) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(self._bench_iters): + target(*real_args, **(real_kwargs or {})) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / self._bench_iters + except Exception: + # Random inputs can violate an op's constraints; fall back analytically. + return self._roofline_estimate(node) + + def _real_leaf(self, n: torch.fx.Node): + """map_arg callback: real random tensor from an input node's fake value. + rand for floats, ones otherwise (matches upstream ``to_real_tensor``).""" + return self._to_real(n.meta.get("val", None)) + + @staticmethod + def _to_real(v): + if isinstance(v, torch.Tensor): + shape = tuple(v.shape) + if v.dtype in _FLOAT_TYPES: + return torch.rand(shape, dtype=v.dtype, device=v.device) + return torch.ones(shape, dtype=v.dtype, device=v.device) + return v + + # ---- entry point: dispatch on the selected mode ---- + def estimate( + self, + traced, + module=None, + *args, + warmup: int = 1, + reps: int = 3, + ) -> RuntimeEstimatorResult: + """Run the selected mode. + + Static modes (cost-model / benchmark) accept a ``GraphModule`` or a + ``TracedResult`` and ignore ``module``/``args``. The interpreter mode + requires a ``TracedResult``, the live ``module``, and the real run + ``args`` (the same inputs ``run_traced`` takes). + """ + if self._estimate_mode_type in (COST_MODEL, BENCHMARK): + gm = traced.gm if isinstance(traced, TracedResult) else traced + return self._estimate_static(gm) + elif self._estimate_mode_type == INTERPRETER: + if not isinstance(traced, TracedResult) or module is None: + raise ValueError( + "interpreter mode needs a TracedResult and the live module: " + "estimate(traced, module, *run_args). Run with cudagraph " + "disabled." + ) + return self._estimate_interpreter( + traced, module, *args, warmup=warmup, reps=reps + ) + else: + raise NotImplementedError( + f"estimate_mode_type {self._estimate_mode_type} not supported." + ) + + # ---- static run (mirrors upstream __torch_dispatch__ accumulation) ---- + def _estimate_static(self, gm: torch.fx.GraphModule) -> RuntimeEstimatorResult: + """Walk the graph, time each costable node via the selected per-node + estimate, and aggregate total / fwd / bwd / per-module runtimes.""" + self.total_runtime = self.fwd_runtime = self.bwd_runtime = 0.0 + self.mod_runtimes = defaultdict(lambda: defaultdict(float)) + self.mod_fw_order, self.mod_bw_order = [], [] + node_runtimes: dict = {} + seen_fw: set = set() + seen_bw: set = set() + + for node in gm.graph.nodes: + if not _is_costable_op(node): + continue + t = self._estimate(node) + if t <= 0.0: + continue + node_runtimes[node.name] = t + self._accumulate(node, t, seen_fw, seen_bw) + + return self._build_result(node_runtimes) + + # ---- interpreter run (real execution of the whole graph) ---- + def _estimate_interpreter( + self, + traced: TracedResult, + module, + *args, + warmup: int = 1, + reps: int = 3, + ) -> RuntimeEstimatorResult: + """Execute the real traced graph via ``run_traced`` and measure runtime. + + ``total_runtime_ms`` is the real end-to-end step time (whole-graph timed + run, overlap-aware, median over reps). The per-node times come from a + serialized ``fx.Interpreter`` and so do NOT sum to the total. Memory is + out of scope here -- use ``measure_traced`` for peak memory. Run with + cudagraph disabled (the per-node interpreter cannot replay a cudagraph). + """ + gm = traced.gm + + # warmup (tracing/autotune/allocator settling) + for _ in range(warmup): + run_traced(traced, module=module)(*args) + torch.cuda.synchronize() + + # total runtime: time the whole traced run once per rep (kernels overlap + # as in a real step), median over reps. We do NOT sum per-op times here. + totals_ms = [] + for _ in range(reps): + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run_traced(traced, module=module)(*args) + end.record() + torch.cuda.synchronize() + totals_ms.append(start.elapsed_time(end)) + total_runtime_ms = median(totals_ms) + + # per-node times via the timing interpreter (serialized cost table). + timing = _capture(traced, module, args, _TimingInterpreter) + per_node = timing.timings_ms() + + self.total_runtime = self.fwd_runtime = self.bwd_runtime = 0.0 + self.mod_runtimes = defaultdict(lambda: defaultdict(float)) + self.mod_fw_order, self.mod_bw_order = [], [] + name_to_node = {n.name: n for n in gm.graph.nodes} + seen_fw: set = set() + seen_bw: set = set() + for nm, t in per_node.items(): + node = name_to_node.get(nm) + if node is not None: + self._accumulate(node, t, seen_fw, seen_bw) + + # The total is the real overlap-aware end-to-end time, not the per-op sum. + self.total_runtime = total_runtime_ms + return self._build_result(per_node) + + # ---- shared accumulation / result helpers ---- + def _accumulate( + self, node: torch.fx.Node, t: float, seen_fw: set, seen_bw: set + ) -> None: + self.total_runtime += t + fqn = _module_fqn(node) + if _is_backward_node(node): + self.bwd_runtime += t + self.mod_runtimes[fqn]["bw"] += t + if fqn and fqn not in seen_bw: + seen_bw.add(fqn) + self.mod_bw_order.append(fqn) + else: + self.fwd_runtime += t + self.mod_runtimes[fqn]["fw"] += t + if fqn and fqn not in seen_fw: + seen_fw.add(fqn) + self.mod_fw_order.append(fqn) + + def _build_result(self, node_runtimes: dict) -> RuntimeEstimatorResult: + return RuntimeEstimatorResult( + total_runtime_ms=self.total_runtime, + fwd_runtime_ms=self.fwd_runtime, + bwd_runtime_ms=self.bwd_runtime, + estimate_mode_type=self._estimate_mode_type, + node_runtimes_ms=node_runtimes, + mod_runtimes_ms={k: dict(v) for k, v in self.mod_runtimes.items()}, + ) + + def display_modulewise_stats(self, depth: int = 2) -> None: + """Print module forward/backward execution orders and per-module + runtimes, up to ``depth`` levels (mirrors upstream).""" + print("Forward Execution Order:") + for fqn in self.mod_fw_order: + if fqn.count(".") + 1 <= depth: + print(f" {fqn}") + print("Backward Execution Order:") + for fqn in self.mod_bw_order: + if fqn.count(".") + 1 <= depth: + print(f" {fqn}") + for fqn, rt in self.mod_runtimes.items(): + if fqn and fqn.count(".") + 1 <= depth: + print( + f"{fqn} fw: {rt.get('fw', 0.0):.3f} ms bw: {rt.get('bw', 0.0):.3f} ms" + ) + + +def estimate_runtime( + gm: torch.fx.GraphModule, + *, + mode: str = COST_MODEL, +) -> RuntimeEstimatorResult: + """Convenience wrapper for the static modes: build a ``RuntimeEstimator`` in + ``mode`` and run it on ``gm``. + + Defaults to the roofline cost model (no GPU execution). Pass + ``mode="operator-level-benchmark"`` to micro-benchmark each op on real + tensors. The interpreter mode needs the live module + run inputs, so use + ``RuntimeEstimator()(INTERPRETER).estimate(traced, module, *args)`` directly. + """ + if mode == INTERPRETER: + raise ValueError( + "interpreter mode needs the live module and run inputs; use " + "RuntimeEstimator()(INTERPRETER).estimate(traced, module, *args)." + ) + return RuntimeEstimator()(mode).estimate(gm) + + +def _op_label(node: torch.fx.Node | None) -> str: + """Short op-type label for a node (overload packet / function name).""" + if node is None: + return "" + t = node.target + if isinstance(t, torch._ops.OpOverload): + return str(t._overloadpacket) + return getattr(t, "__name__", str(t)) + + +def compare_node_breakdown( + cm_result: RuntimeEstimatorResult, + bm_result: RuntimeEstimatorResult, + it_result: RuntimeEstimatorResult, + gm: torch.fx.GraphModule, + top_k: int = 15, +) -> None: + """compares the cost-model, benchmark, and interpreter. + + Joins the three modes' per-node times by node name and reports: + 1. the biggest individual missed nodes (and their op type); + 2. misses grouped by op type; + 3. per-op-type interpreter-vs-cost-model totals + cm/it ratio (where the + roofline under-counts even ops it does cost). + + The post-regional graph: the interpreter reflects the real fused execution. + The pre-regional: the interpreter runs eager flex -> not representative. + """ + name_to_node = {n.name: n for n in gm.graph.nodes} + cmn = cm_result.node_runtimes_ms # roofline per-node (costable, >0) + bmn = bm_result.node_runtimes_ms # benchmark per-node + itn = it_result.node_runtimes_ms # interpreter per-node (ALL compute ops) + + print( + "\n[Per-OP-Comparison] === why cost-model / benchmark underestimate interpreter ===" + ) + print( + f"[Per-OP-Comparison] node counts: interpreter={len(itn)} " + f"cost-model={len(cmn)} benchmark={len(bmn)}" + ) + + # 2. biggest individual nodes the interpreter timed but cost-model missed. + misses = sorted( + ((nm, t) for nm, t in itn.items() if cmn.get(nm, 0.0) <= 0.0), + key=lambda kv: -kv[1], + ) + print( + f"[Per-OP-Comparison] top {top_k} nodes timed by interpreter but ~0 in cost-model:" + ) + for nm, t in misses[:top_k]: + print(f" {t:8.3f} ms {nm:40s} {_op_label(name_to_node.get(nm))}") + + # 3. misses grouped by op type. + miss_by_op: dict = {} + for nm, t in misses: + a = miss_by_op.setdefault(_op_label(name_to_node.get(nm)), [0, 0.0]) + a[0] += 1 + a[1] += t + print("[Per-OP-Comparison] misses grouped by op type (by total interpreter ms):") + for lbl, (cnt, tot) in sorted(miss_by_op.items(), key=lambda kv: -kv[1][1])[:top_k]: + print(f" {tot:8.3f} ms x{cnt:<5} {lbl}") + + # 4. per-op-type interpreter vs cost-model vs benchmark (shows both the + # missed ops and the under-estimated ones). + by_op: dict = {} + for nm, t in itn.items(): + lbl = _op_label(name_to_node.get(nm)) + a = by_op.setdefault(lbl, [0, 0.0, 0.0, 0.0]) # cnt, it, cm, bm + a[0] += 1 + a[1] += t + a[2] += cmn.get(nm, 0.0) + a[3] += bmn.get(nm, 0.0) + print(f"[Per-OP-Comparison] per-op-type (top {top_k} by interpreter ms):") + print( + f" {'op':30s} {'cnt':>4} {'it_ms':>9} {'cm_ms':>9} {'bm_ms':>9} " + f"{'cm/it':>6} {'bm/it':>6}" + ) + for lbl, (cnt, itt, cmt, bmt) in sorted(by_op.items(), key=lambda kv: -kv[1][1])[ + :top_k + ]: + rc = cmt / itt if itt else float("nan") + rb = bmt / itt if itt else float("nan") + print( + f" {lbl:30s} {cnt:>4} {itt:>9.3f} {cmt:>9.3f} {bmt:>9.3f} " + f"{rc:>6.2f} {rb:>6.2f}" + ) diff --git a/torchtitan/experiments/graph_trainer/transfertime_estimator.py b/torchtitan/experiments/graph_trainer/transfertime_estimator.py new file mode 100644 index 0000000000..a2598066e9 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/transfertime_estimator.py @@ -0,0 +1,276 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Static D2H/H2D transfer-time estimator for a joint fwd+loss+bwd FX graph. + +The transfer-time analogue of ``memory_estimator.py``. Where that tracks storage +liveness, this one walks the +joint graph and, for every node, estimates the time to move its tensors across +the host<->device link: a node's *inputs* must be RELOADED (H2D) before it runs, +and its *outputs* would be OFFLOADED (D2H) after it produces them. It is the cost +model an activation-offload solver uses to decide what fits in the offload +bandwidth budget. + +Each ``TensorObject`` carries the underlying ``storage_id`` and the producer's +``consumer_nodes`` so the downstream solver can dedup: tensors that are +views/aliases share a storage and need to move only once, and a tensor with +multiple consumers is reloaded only once. Sizes use the underlying storage's +bytes (a view offloads its whole storage), not the logical tensor size. + +Bandwidth: ``estimate_transfertime`` accepts explicit ``bw_h2d`` / ``bw_d2h`` +(GB/s); when omitted it calls ``get_transfer_bw``, which benchmarks the +host<->device link once per device and caches the result for the process (the +link bandwidth -- GPU + PCIe/NVLink topology + NUMA binding -- does not change +during a run). Pass the bandwidths in (e.g. in tests) to skip the live +benchmark; call ``measure_transfer_bw`` directly to force a fresh measurement. +""" + +import functools +import time +from dataclasses import dataclass, field + +import torch +import torch.utils._pytree as pytree + +from torchtitan.experiments.graph_trainer.common_utils import _is_backward_node + + +def measure_transfer_bw( + nbytes: int = 512 * 1024 * 1024, + iters: int = 20, + pinned: bool = True, + dev: str = "cuda:0", +) -> dict: + """Measure H2D and D2H bandwidth (GB/s) with pinned host memory. + + Pageable host memory is much slower, so pinned is the default and the regime + an offload path would actually use. + + The default 512 MB saturates the link. A single tensor smaller than that may + not reach this bandwidth, so the roofline over-predicts throughput for tiny + offloads; in practice an offload solver should bundle small tensors to + saturate the link (or treat sub-link-saturating sizes as launch-bound). + """ + host = torch.empty(nbytes, dtype=torch.uint8, pin_memory=pinned) + gpu = torch.empty(nbytes, dtype=torch.uint8, device=dev) + + def time_copy(dst, src): + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + dst.copy_(src, non_blocking=True) + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters + + # warmup + gpu.copy_(host, non_blocking=True) + host.copy_(gpu, non_blocking=True) + torch.cuda.synchronize() + + h2d_s = time_copy(gpu, host) # host -> device + d2h_s = time_copy(host, gpu) # device -> host + return {"h2d": nbytes / h2d_s / 1e9, "d2h": nbytes / d2h_s / 1e9} + + +# Cache the immutable (h2d, d2h) pair keyed by the args that actually change the +# result: the device (its host link + NUMA binding) and the measurement regime +# (size / iters / pinned). All are fixed within a training process, so the +# benchmark runs once per distinct key. lru_cache caches a tuple (hashable + +# immutable) so a caller cannot mutate the shared result. +@functools.lru_cache(maxsize=None) +def _cached_transfer_bw( + dev: str, nbytes: int, iters: int, pinned: bool +) -> tuple[float, float]: + bw = measure_transfer_bw(nbytes=nbytes, iters=iters, pinned=pinned, dev=dev) + return bw["h2d"], bw["d2h"] + + +def get_transfer_bw( + dev: str = "cuda:0", + nbytes: int = 512 * 1024 * 1024, + iters: int = 20, + pinned: bool = True, +) -> dict: + """Measured ``{h2d, d2h}`` GB/s for ``dev``, benchmarked once per device and + cached for the process lifetime (see ``_cached_transfer_bw``). + + Returns a fresh dict each call so callers may mutate it safely. Use + ``measure_transfer_bw`` directly to bypass the cache and force a fresh run. + """ + h2d, d2h = _cached_transfer_bw(dev, nbytes, iters, pinned) + return {"h2d": h2d, "d2h": d2h} + + +def _transfer_ms(nbytes: int, bw_gbps: float) -> float: + """Transfer time in ms for ``nbytes`` at ``bw_gbps`` GB/s. + GB/s -> bytes/ms is a factor of 1e6 (1e9 bytes/s / 1e3 ms/s).""" + assert nbytes >= 0, f"negative bytes: {nbytes}" + return nbytes / (bw_gbps * 1e6) + + +def _storage_bytes_and_id(t: torch.Tensor) -> tuple[int, int]: + """Underlying storage size (bytes) and its id. The id (``_cdata``) identifies + views/aliases (they share a storage) but is only unique among *live* tensors: + the allocator reuses it after a storage frees, so dedup within a liveness + window, not globally.""" + st = t.untyped_storage() + return st.nbytes(), st._cdata + + +class TensorObject: + """One tensor moved across the host<->device link, with its transfer time. + + ``size`` is the underlying storage's bytes (a view offloads its whole + storage). ``storage_id`` lets a downstream solver merge views/aliases so a + storage moves once; ``consumer_nodes`` lists every reader of the producer so a + multi-consumer tensor is reloaded once. + """ + + def __init__(self, tensor_type, shape, dtype, device, size, storage_id, time_ms): + self.tensor_type = tensor_type # "forward" | "backward" (producer region) + self.shape = shape + self.dtype = dtype + self.device = device + self.size = size # underlying storage bytes + self.storage_id = storage_id + self.offload_time_ms = time_ms + self.consumer_nodes = [] + + def __repr__(self): + return ( + f"TensorObject(type={self.tensor_type}, shape={tuple(self.shape)}, " + f"dtype={self.dtype}, device={self.device}, size={self.size}, " + f"storage_id={self.storage_id})" + ) + + +def categorize_fwd_bwd(node: torch.fx.Node) -> str: + return "backward" if _is_backward_node(node) else "forward" + + +def _max_offload(d: dict) -> float: + return max((t.offload_time_ms for lst in d.values() for t in lst), default=0.0) + + +@dataclass +class TransferEstimatorResult: + # node name -> list[TensorObject]: + # inputs = tensors this node would RELOAD (H2D) + # outputs = tensors this node would OFFLOAD (D2H) + input_tensor_transfer_times_ms: dict = field(default_factory=dict) + output_tensor_transfer_times_ms: dict = field(default_factory=dict) + + def max_input_offload_time(self) -> float: + return _max_offload(self.input_tensor_transfer_times_ms) + + def max_output_offload_time(self) -> float: + return _max_offload(self.output_tensor_transfer_times_ms) + + def summary(self, top_k: int = 10) -> str: + if ( + not self.output_tensor_transfer_times_ms + and not self.input_tensor_transfer_times_ms + ): + return "No tensor data found" + lines = [ + f"max input reload (H2D) time: {self.max_input_offload_time():.3f} ms", + f"max output offload (D2H) time: {self.max_output_offload_time():.3f} ms", + ] + # rank nodes by their heaviest produced tensor + ranked = sorted( + self.output_tensor_transfer_times_ms.items(), + key=lambda kv: max((t.offload_time_ms for t in kv[1]), default=0.0), + reverse=True, + ) + for name, lst in ranked[:top_k]: + t = max(lst, key=lambda t: t.offload_time_ms, default=None) + if t is not None: + lines.append(f" {t.offload_time_ms:8.4f} ms {name} ({t.size} B)") + return "\n".join(lines) + + +def _tensor_offload_ms( + node: torch.fx.Node, bw_h2d: float, bw_d2h: float +) -> tuple[list[TensorObject], list[TensorObject]]: + """Build the input (reload, H2D) and output (offload, D2H) ``TensorObject`` + lists for ``node``. + + Inputs are the values of this node's predecessors -- if offloaded, they must + be reloaded (H2D) before this node runs -- tagged with the producer's real + consumer set (``prod.users``). Outputs are this node's own produced tensors, + which an offload path would move to host (D2H), tagged with ``node.users``. + Sizes/ids come from the underlying storage so views/aliases are dedupable + downstream by ``storage_id``. + """ + input_tensor_objects: list[TensorObject] = [] + output_tensor_objects: list[TensorObject] = [] + + for prod in node.all_input_nodes: + prod_consumers = [u.name for u in prod.users] + for t in pytree.tree_leaves(prod.meta.get("val", None)): + if isinstance(t, torch.Tensor): + nbytes, sid = _storage_bytes_and_id(t) + obj = TensorObject( + tensor_type=categorize_fwd_bwd(prod), + shape=t.shape, + dtype=t.dtype, + device=t.device, + size=nbytes, + storage_id=sid, + time_ms=_transfer_ms(nbytes, bw_h2d), # reload: H2D + ) + obj.consumer_nodes = prod_consumers + input_tensor_objects.append(obj) + + # Consumers of this node's outputs are node.users. For a multi-output op these + # are the getitem accessors; mapping each output index to its true downstream + output_consumers = [u.name for u in node.users] + for t in pytree.tree_leaves(node.meta.get("val", None)): + if isinstance(t, torch.Tensor): + nbytes, sid = _storage_bytes_and_id(t) + obj = TensorObject( + tensor_type=categorize_fwd_bwd(node), + shape=t.shape, + dtype=t.dtype, + device=t.device, + size=nbytes, + storage_id=sid, + time_ms=_transfer_ms(nbytes, bw_d2h), # offload: D2H + ) + obj.consumer_nodes = output_consumers + output_tensor_objects.append(obj) + + return input_tensor_objects, output_tensor_objects + + +def estimate_transfertime( + gm: torch.fx.GraphModule, + bw_h2d: float | None = None, + bw_d2h: float | None = None, + dev: str = "cuda:0", +) -> TransferEstimatorResult: + """Per-node D2H/H2D transfer times for the joint fwd+bwd graph. + + ``bw_h2d`` / ``bw_d2h`` (GB/s) default to the cached ``get_transfer_bw(dev)`` + measurement, so repeated calls in one process reuse a single benchmark. + """ + if bw_h2d is None or bw_d2h is None: + measured = get_transfer_bw(dev=dev) + bw_h2d = bw_h2d if bw_h2d is not None else measured["h2d"] + bw_d2h = bw_d2h if bw_d2h is not None else measured["d2h"] + + input_times: dict = {} + output_times: dict = {} + for node in gm.graph.nodes: + inp, out = _tensor_offload_ms(node, bw_h2d, bw_d2h) + input_times[node.name] = inp + output_times[node.name] = out + + result = TransferEstimatorResult( + input_tensor_transfer_times_ms=input_times, + output_tensor_transfer_times_ms=output_times, + ) + return result