Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
+191
@@ -0,0 +1,191 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from torch._logging import trace_structured
|
||||
from torch.fx import Graph, Node
|
||||
|
||||
|
||||
log: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_joint_graph_node_information(
|
||||
joint_graph: Graph,
|
||||
recomputable_node_info: dict[str, int],
|
||||
) -> dict[str, Any]:
|
||||
joint_graph_node_information: dict[str, Any] = {}
|
||||
|
||||
for i, joint_graph_node in enumerate(joint_graph.nodes):
|
||||
is_recomputable_candidate: bool = (
|
||||
joint_graph_node.name in recomputable_node_info
|
||||
)
|
||||
tensor_meta = joint_graph_node.meta.get("tensor_meta")
|
||||
# pyrefly: ignore [implicit-any]
|
||||
shape = getattr(tensor_meta, "shape", []) if tensor_meta else []
|
||||
|
||||
node_info: dict[str, Any] = {
|
||||
"index": i,
|
||||
"name": joint_graph_node.name,
|
||||
"is_recomputable_candidate": is_recomputable_candidate,
|
||||
"target": str(joint_graph_node.target),
|
||||
"shape": str(shape),
|
||||
"input_arguments": [inp.name for inp in joint_graph_node.all_input_nodes],
|
||||
"stack_trace": joint_graph_node.meta.get("stack_trace", ""),
|
||||
}
|
||||
|
||||
if is_recomputable_candidate:
|
||||
idx: int = recomputable_node_info[joint_graph_node.name]
|
||||
node_info["recomputable_candidate_info"] = {
|
||||
"recomputable_node_idx": idx,
|
||||
}
|
||||
|
||||
joint_graph_node_information[joint_graph_node.name] = node_info
|
||||
|
||||
return joint_graph_node_information
|
||||
|
||||
|
||||
def create_joint_graph_edges(joint_graph: Graph) -> list[tuple[str, str]]:
|
||||
joint_graph_edges: list[tuple[str, str]] = [
|
||||
(inp.name, node.name)
|
||||
for node in joint_graph.nodes
|
||||
for inp in node.all_input_nodes
|
||||
]
|
||||
return joint_graph_edges
|
||||
|
||||
|
||||
def create_activation_checkpointing_logging_structure_payload(
|
||||
joint_graph: Graph,
|
||||
joint_graph_node_information: dict[str, Any],
|
||||
joint_graph_edges: list[tuple[str, str]],
|
||||
all_recomputable_banned_nodes: list[Node],
|
||||
expected_runtime: float,
|
||||
saved_node_idxs: list[int],
|
||||
recomputable_node_idxs: list[int],
|
||||
memories_banned_nodes: list[int],
|
||||
normalized_memories_banned_nodes: list[float],
|
||||
runtimes_banned_nodes: list[float],
|
||||
min_cut_saved_values: list[Node],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Creates a structured payload for logging activation checkpointing information.
|
||||
|
||||
Args:
|
||||
joint_graph: The computational graph representing operations.
|
||||
joint_graph_node_information: Dictionary containing information about nodes in the joint graph.
|
||||
joint_graph_edges: List of edges in the joint graph represented as tuples of node names.
|
||||
all_recomputable_banned_nodes: List of nodes that are banned from recomputation.
|
||||
expected_runtime: Expected runtime of the computation.
|
||||
saved_node_idxs: Indices of nodes that are saved (not recomputed).
|
||||
recomputable_node_idxs: Indices of nodes that can be recomputed.
|
||||
memories_banned_nodes: Memory usage values (in absolute units) for banned nodes.
|
||||
normalized_memories_banned_nodes: Normalized memory usage values for banned nodes,
|
||||
used as input to the knapsack algorithm.
|
||||
runtimes_banned_nodes: Runtime values for banned nodes, used as input to the
|
||||
knapsack algorithm.
|
||||
min_cut_saved_values: List of nodes saved by the min-cut algorithm.
|
||||
|
||||
Returns:
|
||||
A dictionary containing structured logging information for activation checkpointing.
|
||||
"""
|
||||
activation_checkpointing_logging_structure_payload: dict[str, Any] = {
|
||||
"Joint Graph Size": len(joint_graph.nodes),
|
||||
"Joint Graph Edges": {
|
||||
"Total": len(joint_graph_edges),
|
||||
"Edges": joint_graph_edges,
|
||||
},
|
||||
"Joint Graph Node Information": joint_graph_node_information,
|
||||
"Recomputable Banned Nodes Order": [
|
||||
node.name for node in all_recomputable_banned_nodes
|
||||
],
|
||||
"Expected Runtime": expected_runtime,
|
||||
"Knapsack Saved Nodes": saved_node_idxs,
|
||||
"Knapsack Recomputed Nodes": recomputable_node_idxs,
|
||||
"Absolute Memories": memories_banned_nodes,
|
||||
"Knapsack Input Memories": normalized_memories_banned_nodes,
|
||||
"Knapsack Input Runtimes": runtimes_banned_nodes,
|
||||
"Min Cut Solution Saved Values": [node.name for node in min_cut_saved_values],
|
||||
}
|
||||
return activation_checkpointing_logging_structure_payload
|
||||
|
||||
|
||||
def create_structured_trace_for_min_cut_info(
|
||||
joint_graph: Graph,
|
||||
all_recomputable_banned_nodes: list[Node],
|
||||
saved_node_idxs: list[int],
|
||||
recomputable_node_idxs: list[int],
|
||||
expected_runtime: float,
|
||||
memories_banned_nodes: list[int],
|
||||
normalized_memories_banned_nodes: list[float],
|
||||
runtimes_banned_nodes: list[float],
|
||||
min_cut_saved_values: list[Node],
|
||||
) -> None:
|
||||
"""
|
||||
Creates a structured trace for minimum cut information in the graph.
|
||||
|
||||
Args:
|
||||
joint_graph: The computational graph representation.
|
||||
all_recomputable_banned_nodes: List of nodes that can be recomputed.
|
||||
saved_node_idxs: Indices of nodes that are saved in memory.
|
||||
recomputable_node_idxs: Indices of nodes that are recomputed.
|
||||
expected_runtime: Expected runtime for the computation.
|
||||
memories_banned_nodes: Memory requirements for each banned node in bytes.
|
||||
normalized_memories_banned_nodes: Normalized memory requirements for each banned node
|
||||
(typically scaled between 0 and 1 for relative comparison).
|
||||
runtimes_banned_nodes: Runtime costs associated with each banned node.
|
||||
min_cut_saved_values: Nodes that are saved as part of the minimum cut solution.
|
||||
"""
|
||||
# Create a dictionary to store recomputable node information
|
||||
recomputable_node_info: dict[str, int] = {
|
||||
node.name: idx for idx, node in enumerate(all_recomputable_banned_nodes)
|
||||
}
|
||||
|
||||
# Create joint graph node information
|
||||
joint_graph_node_information = create_joint_graph_node_information(
|
||||
joint_graph, recomputable_node_info
|
||||
)
|
||||
|
||||
# Update node information with recomputable candidate details
|
||||
for node_name, node_info in joint_graph_node_information.items():
|
||||
if node_info["is_recomputable_candidate"]:
|
||||
idx = recomputable_node_info[node_name]
|
||||
node_info["recomputable_candidate_info"]["memory"] = memories_banned_nodes[
|
||||
idx
|
||||
]
|
||||
node_info["recomputable_candidate_info"]["runtime"] = runtimes_banned_nodes[
|
||||
idx
|
||||
]
|
||||
node_info["recomputable_candidate_info"]["is_saved"] = (
|
||||
idx in saved_node_idxs
|
||||
)
|
||||
node_info["recomputable_candidate_info"]["is_recomputed"] = (
|
||||
idx in recomputable_node_idxs
|
||||
)
|
||||
|
||||
# Create joint graph edges
|
||||
joint_graph_edges = create_joint_graph_edges(joint_graph)
|
||||
|
||||
# Create activation checkpointing logging structure payload
|
||||
activation_checkpointing_logging_structure_payload = (
|
||||
create_activation_checkpointing_logging_structure_payload(
|
||||
joint_graph=joint_graph,
|
||||
joint_graph_node_information=joint_graph_node_information,
|
||||
joint_graph_edges=joint_graph_edges,
|
||||
all_recomputable_banned_nodes=all_recomputable_banned_nodes,
|
||||
expected_runtime=expected_runtime,
|
||||
saved_node_idxs=saved_node_idxs,
|
||||
recomputable_node_idxs=recomputable_node_idxs,
|
||||
memories_banned_nodes=memories_banned_nodes,
|
||||
normalized_memories_banned_nodes=normalized_memories_banned_nodes,
|
||||
runtimes_banned_nodes=runtimes_banned_nodes,
|
||||
min_cut_saved_values=min_cut_saved_values,
|
||||
)
|
||||
)
|
||||
|
||||
# Create structured trace
|
||||
trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {"name": "min_cut_information", "encoding": "json"},
|
||||
payload_fn=lambda: json.dumps(
|
||||
activation_checkpointing_logging_structure_payload
|
||||
),
|
||||
)
|
||||
+319
@@ -0,0 +1,319 @@
|
||||
from typing import Any
|
||||
|
||||
import networkx as nx
|
||||
|
||||
from torch.fx import Graph, Node
|
||||
|
||||
|
||||
class GraphInfoProvider:
|
||||
"""
|
||||
This class provides information about the graph, such as the nodes, edges, and their runtime and memory requirements.
|
||||
It also provides methods to create graphs from the information provided.
|
||||
"""
|
||||
|
||||
__RECOMPUTABLE_NODE_ONLY_GRAPH = "recomputable_node_only_graph"
|
||||
__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT = (
|
||||
"recomputable_node_only_graph_with_larger_graph_context"
|
||||
)
|
||||
__FULL_NX_JOINT_GRAPH = "full_nx_joint_graph"
|
||||
__SIMPLIFIED_FX_JOINT_GRAPH = "fx_joint_graph"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
graph_nodes_in_order: list[str],
|
||||
graph_edges: list[tuple[str, str]],
|
||||
all_recomputable_banned_nodes: list[str],
|
||||
all_node_runtimes: dict[str, float] | None = None,
|
||||
all_node_memories: dict[str, float] | None = None,
|
||||
recorded_knapsack_input_memories: list[float] | None = None,
|
||||
recorded_knapsack_input_runtimes: list[float] | None = None,
|
||||
joint_graph: Graph | None = None,
|
||||
) -> None:
|
||||
self.graph_nodes_in_order = graph_nodes_in_order
|
||||
self.graph_edges = graph_edges
|
||||
self.all_node_runtimes: dict[str, float] = dict()
|
||||
if all_node_runtimes is None:
|
||||
if recorded_knapsack_input_runtimes is None:
|
||||
raise ValueError(
|
||||
"Either all_node_runtimes or recorded_knapsack_input_runtimes must be provided."
|
||||
)
|
||||
self.all_node_runtimes = {
|
||||
node: recorded_knapsack_input_runtimes[i]
|
||||
for i, node in enumerate(all_recomputable_banned_nodes)
|
||||
}
|
||||
else:
|
||||
self.all_node_runtimes.update(all_node_runtimes)
|
||||
self.all_node_memories: dict[str, float] = dict()
|
||||
if all_node_memories is None:
|
||||
if recorded_knapsack_input_memories is None:
|
||||
raise ValueError(
|
||||
"Either all_node_memories or recorded_knapsack_input_memories must be provided."
|
||||
)
|
||||
self.all_node_memories = {
|
||||
node: recorded_knapsack_input_memories[i]
|
||||
for i, node in enumerate(all_recomputable_banned_nodes)
|
||||
}
|
||||
else:
|
||||
self.all_node_memories.update(all_node_memories)
|
||||
self.all_recomputable_banned_nodes = all_recomputable_banned_nodes
|
||||
self.all_recomputable_banned_nodes_set = set(all_recomputable_banned_nodes)
|
||||
self.recorded_knapsack_input_memories = recorded_knapsack_input_memories
|
||||
self.recorded_knapsack_input_runtimes = recorded_knapsack_input_runtimes
|
||||
self._lazily_initialized_graphs: dict[str, Any] = {
|
||||
self.__RECOMPUTABLE_NODE_ONLY_GRAPH: None,
|
||||
self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT: None,
|
||||
self.__FULL_NX_JOINT_GRAPH: None,
|
||||
self.__SIMPLIFIED_FX_JOINT_GRAPH: None,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def inialize_from_graph(
|
||||
cls,
|
||||
joint_graph: Graph,
|
||||
all_recomputable_banned_nodes: list[Node],
|
||||
recorded_knapsack_input_memories: list[float],
|
||||
recorded_knapsack_input_runtimes: list[float],
|
||||
) -> "GraphInfoProvider":
|
||||
"""
|
||||
Enables initialization from a joint graph.
|
||||
"""
|
||||
graph_nodes_in_order = [node.name for node in joint_graph.nodes]
|
||||
graph_edges = [
|
||||
(node.name, user.name) for node in joint_graph.nodes for user in node.users
|
||||
]
|
||||
all_recomputable_banned_node_names = [
|
||||
node.name for node in all_recomputable_banned_nodes
|
||||
]
|
||||
return cls(
|
||||
graph_nodes_in_order=graph_nodes_in_order,
|
||||
graph_edges=graph_edges,
|
||||
all_recomputable_banned_nodes=all_recomputable_banned_node_names,
|
||||
recorded_knapsack_input_memories=recorded_knapsack_input_memories,
|
||||
recorded_knapsack_input_runtimes=recorded_knapsack_input_runtimes,
|
||||
joint_graph=joint_graph,
|
||||
)
|
||||
|
||||
@property
|
||||
def recomputable_node_only_graph(self) -> nx.DiGraph:
|
||||
if self._lazily_initialized_graphs[self.__RECOMPUTABLE_NODE_ONLY_GRAPH] is None:
|
||||
self._lazily_initialized_graphs[self.__RECOMPUTABLE_NODE_ONLY_GRAPH] = (
|
||||
self._create_recomputable_node_only_graph()
|
||||
)
|
||||
return self._lazily_initialized_graphs[self.__RECOMPUTABLE_NODE_ONLY_GRAPH]
|
||||
|
||||
@property
|
||||
def recomputable_node_only_graph_with_larger_graph_context(self) -> nx.DiGraph:
|
||||
if (
|
||||
self._lazily_initialized_graphs[
|
||||
self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT
|
||||
]
|
||||
is None
|
||||
):
|
||||
self._lazily_initialized_graphs[
|
||||
self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT
|
||||
] = self._create_recomputable_node_only_graph_with_larger_graph_context()
|
||||
return self._lazily_initialized_graphs[
|
||||
self.__RECOMPUTABLE_NODE_ONLY_GRAPH_WITH_LARGER_GRAPH_CONTEXT
|
||||
]
|
||||
|
||||
@property
|
||||
def full_joint_nx_graph(self) -> nx.DiGraph:
|
||||
if self._lazily_initialized_graphs[self.__FULL_NX_JOINT_GRAPH] is None:
|
||||
self._lazily_initialized_graphs[self.__FULL_NX_JOINT_GRAPH] = (
|
||||
self._create_full_joint_graph()
|
||||
)
|
||||
return self._lazily_initialized_graphs[self.__FULL_NX_JOINT_GRAPH]
|
||||
|
||||
@property
|
||||
def simplified_fx_joint_graph(self) -> Graph:
|
||||
if self._lazily_initialized_graphs[self.__SIMPLIFIED_FX_JOINT_GRAPH] is None:
|
||||
self._lazily_initialized_graphs[self.__SIMPLIFIED_FX_JOINT_GRAPH] = (
|
||||
self._recreate_psuedo_joint_graph()
|
||||
)
|
||||
return self._lazily_initialized_graphs[self.__SIMPLIFIED_FX_JOINT_GRAPH]
|
||||
|
||||
def get_non_ac_peak_memory(self) -> float:
|
||||
return sum(
|
||||
self.all_node_memories[node_name]
|
||||
for node_name in self.all_recomputable_banned_nodes_set
|
||||
)
|
||||
|
||||
def get_theoretical_max_runtime(self) -> float:
|
||||
return sum(
|
||||
self.all_node_runtimes[node_name]
|
||||
for node_name in self.all_recomputable_banned_nodes_set
|
||||
)
|
||||
|
||||
def get_knapsack_memory_input(self) -> list[float]:
|
||||
return (
|
||||
self.recorded_knapsack_input_memories
|
||||
if self.recorded_knapsack_input_memories
|
||||
else [
|
||||
self.all_node_memories[node_name]
|
||||
for node_name in self.all_recomputable_banned_nodes
|
||||
]
|
||||
)
|
||||
|
||||
def get_knapsack_runtime_input(self) -> list[float]:
|
||||
return (
|
||||
self.recorded_knapsack_input_runtimes
|
||||
if self.recorded_knapsack_input_runtimes
|
||||
else [
|
||||
self.all_node_runtimes[node_name]
|
||||
for node_name in self.all_recomputable_banned_nodes
|
||||
]
|
||||
)
|
||||
|
||||
def _create_recomputable_node_only_graph(self) -> nx.DiGraph:
|
||||
graph = nx.DiGraph()
|
||||
for recomputable_node in self.all_recomputable_banned_nodes:
|
||||
graph.add_node(recomputable_node)
|
||||
|
||||
for a, b in self.graph_edges:
|
||||
if (
|
||||
a in self.all_recomputable_banned_nodes_set
|
||||
and b in self.all_recomputable_banned_nodes_set
|
||||
):
|
||||
graph.add_edge(a, b)
|
||||
return graph
|
||||
|
||||
def _create_recomputable_node_only_graph_with_larger_graph_context(
|
||||
self,
|
||||
) -> nx.DiGraph:
|
||||
# Create a dictionary to store the reachable nodes for each node
|
||||
all_recomputable_banned_nodes_set = set(self.all_recomputable_banned_nodes)
|
||||
|
||||
reachable_nodes = {}
|
||||
for node in all_recomputable_banned_nodes_set:
|
||||
# Use BFS to find all reachable nodes
|
||||
predecessors = dict(nx.bfs_predecessors(self.full_joint_nx_graph, node))
|
||||
reachable_recomputable_nodes = set(predecessors.keys()).intersection(
|
||||
all_recomputable_banned_nodes_set
|
||||
)
|
||||
reachable_nodes[node] = reachable_recomputable_nodes
|
||||
# Create the candidate graph
|
||||
candidate_graph = nx.DiGraph()
|
||||
candidate_graph.add_nodes_from(all_recomputable_banned_nodes_set)
|
||||
for node1 in all_recomputable_banned_nodes_set:
|
||||
for node2 in reachable_nodes[node1]:
|
||||
# Check if there is an overlapping path
|
||||
overlapping_path = False
|
||||
for intermediate_node in reachable_nodes[node1]:
|
||||
if (
|
||||
intermediate_node != node2
|
||||
and node2 in reachable_nodes[intermediate_node]
|
||||
):
|
||||
overlapping_path = True
|
||||
break
|
||||
if not overlapping_path:
|
||||
candidate_graph.add_edge(node1, node2)
|
||||
return candidate_graph
|
||||
|
||||
def _create_full_joint_graph(self) -> nx.DiGraph:
|
||||
graph = nx.DiGraph()
|
||||
for node in self.graph_nodes_in_order:
|
||||
if node == "output":
|
||||
continue
|
||||
graph.add_node(node)
|
||||
|
||||
for a, b in self.graph_edges:
|
||||
if a == "output" or b == "output":
|
||||
continue
|
||||
graph.add_edge(a, b)
|
||||
return graph
|
||||
|
||||
def _recreate_psuedo_joint_graph(self) -> Graph:
|
||||
# Create a dictionary to store the dependencies of each node
|
||||
node_dependencies: dict[str, list[str]] = {
|
||||
node: [] for node in self.graph_nodes_in_order
|
||||
}
|
||||
for a, b in self.graph_edges:
|
||||
if a not in node_dependencies or b not in node_dependencies:
|
||||
raise ValueError(f"Edge ({a}, {b}) references a non-existent node.")
|
||||
node_dependencies[b].append(a)
|
||||
|
||||
joint_graph = Graph()
|
||||
# Create nodes in the graph
|
||||
nodes: dict[str, Node] = {}
|
||||
for node_name in self.graph_nodes_in_order:
|
||||
input_nodes = [nodes[dep] for dep in node_dependencies[node_name]]
|
||||
if input_nodes:
|
||||
node = joint_graph.call_function(lambda *x: x, tuple(input_nodes))
|
||||
node.name = node_name
|
||||
else:
|
||||
node = joint_graph.placeholder(node_name)
|
||||
nodes[node_name] = node
|
||||
return joint_graph
|
||||
|
||||
def _visualize_recomputable_candidate_graph_with_larger_context(
|
||||
self,
|
||||
layout_k: float = 0.5,
|
||||
layout_iterations: int = 30,
|
||||
) -> None:
|
||||
"""
|
||||
Visualize the recomputable candidate graph with larger context.
|
||||
"""
|
||||
from matplotlib import cm, colors as mcolors, pyplot as plt
|
||||
|
||||
pos = nx.spring_layout(
|
||||
self.recomputable_node_only_graph_with_larger_graph_context,
|
||||
k=layout_k,
|
||||
iterations=layout_iterations,
|
||||
)
|
||||
# pos = nx.spectral_layout(graph_with_indirect_edges)
|
||||
plt.figure(figsize=(20, 15))
|
||||
|
||||
# Create a dictionary for node labels using the index
|
||||
labels = {
|
||||
node: self.recomputable_node_only_graph_with_larger_graph_context.nodes[
|
||||
node
|
||||
].get("index", node)
|
||||
for node in self.recomputable_node_only_graph_with_larger_graph_context.nodes
|
||||
}
|
||||
|
||||
# Extract memory values and normalize them
|
||||
norm = mcolors.Normalize(
|
||||
vmin=min(self.get_knapsack_memory_input()),
|
||||
vmax=max(self.get_knapsack_memory_input()),
|
||||
)
|
||||
cmap = cm.viridis # type: ignore[attr-defined]
|
||||
|
||||
# Assign colors based on memory
|
||||
node_colors = [
|
||||
cmap(
|
||||
norm(
|
||||
float(
|
||||
self.recomputable_node_only_graph_with_larger_graph_context.nodes[
|
||||
node
|
||||
]["memory"]
|
||||
)
|
||||
)
|
||||
)
|
||||
for node in self.recomputable_node_only_graph_with_larger_graph_context.nodes
|
||||
]
|
||||
|
||||
# Draw the graph with parsed nodes only
|
||||
nx.draw_networkx_nodes(
|
||||
self.recomputable_node_only_graph_with_larger_graph_context,
|
||||
pos,
|
||||
node_color=node_colors,
|
||||
node_size=300,
|
||||
label="Parsed Nodes",
|
||||
)
|
||||
nx.draw_networkx_edges(
|
||||
self.recomputable_node_only_graph_with_larger_graph_context,
|
||||
pos,
|
||||
arrows=True,
|
||||
arrowsize=10,
|
||||
)
|
||||
nx.draw_networkx_labels(
|
||||
self.recomputable_node_only_graph_with_larger_graph_context,
|
||||
pos,
|
||||
labels=labels,
|
||||
font_size=8,
|
||||
font_weight="bold",
|
||||
)
|
||||
|
||||
plt.title("Memory Colour Coded Dependency Graph for Recomputable Nodes")
|
||||
plt.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), label="Memory")
|
||||
plt.show()
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
import torch
|
||||
|
||||
|
||||
def greedy_knapsack(
|
||||
memory: list[float], runtimes: list[float], max_memory: float
|
||||
) -> tuple[float, list[int], list[int]]:
|
||||
n = len(runtimes)
|
||||
items = list(range(n))
|
||||
|
||||
# Sort items based on the ratio of runtime to memory in descending order
|
||||
items = sorted(items, key=lambda i: runtimes[i] / memory[i], reverse=True)
|
||||
|
||||
total_memory = 0.0
|
||||
total_runtime = 0.0
|
||||
items_to_save = []
|
||||
items_to_allow_recomputing = []
|
||||
|
||||
for i in items:
|
||||
if total_memory + memory[i] <= max_memory:
|
||||
total_memory += memory[i]
|
||||
total_runtime += runtimes[i]
|
||||
items_to_save.append(i)
|
||||
else:
|
||||
items_to_allow_recomputing.append(i)
|
||||
return total_runtime, items_to_save, items_to_allow_recomputing
|
||||
|
||||
|
||||
def ilp_knapsack(
|
||||
memory: list[float], runtimes: list[float], max_memory: float
|
||||
) -> tuple[float, list[int], list[int]]:
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from scipy.optimize import Bounds, LinearConstraint, milp
|
||||
except ImportError:
|
||||
raise RuntimeError(
|
||||
"To use the ILP for memory budget checkpointing you need to install scipy"
|
||||
) from None
|
||||
|
||||
np_memory = np.array(memory)
|
||||
np_runtimes = np.array(runtimes)
|
||||
c = -np_runtimes # type: ignore[operator]
|
||||
|
||||
memory_constraint = LinearConstraint(A=np_memory, ub=np.array(max_memory))
|
||||
constraints = [memory_constraint]
|
||||
|
||||
integrality = np.ones_like(c)
|
||||
res = milp(
|
||||
c=c, constraints=constraints, integrality=integrality, bounds=Bounds(0, 1)
|
||||
)
|
||||
if not res.success:
|
||||
raise RuntimeError("Somehow scipy solving failed")
|
||||
|
||||
items_to_save = []
|
||||
items_to_allow_recomputing = []
|
||||
for idx, i in enumerate(res.x):
|
||||
if i == 1:
|
||||
items_to_save.append(idx)
|
||||
else:
|
||||
items_to_allow_recomputing.append(idx)
|
||||
return -res.fun, items_to_save, items_to_allow_recomputing
|
||||
|
||||
|
||||
def dp_knapsack(
|
||||
memory: list[float], runtime: list[float], max_memory: float
|
||||
) -> tuple[float, list[int], list[int]]:
|
||||
# Scaling factor to convert floating point weights to integers
|
||||
S = 10000
|
||||
|
||||
# Quantize the memory weights
|
||||
quantized_memory = torch.tensor(
|
||||
[round(m * S) for m in memory], dtype=torch.long, device="cpu"
|
||||
)
|
||||
runtimes = torch.tensor(runtime, dtype=torch.float32, device="cpu")
|
||||
|
||||
# Quantized pseudopolynomial DP for 0-1 Knapsack
|
||||
quantized_max_memory = round(max_memory * S)
|
||||
|
||||
n = len(memory)
|
||||
|
||||
# Initialize the DP table
|
||||
# TODO(chilli): I think if needed, this memory can be optimized with sliding
|
||||
# window trick + Hirschberg trick:
|
||||
# https://codeforces.com/blog/entry/47247?#comment-316200
|
||||
dp = torch.zeros(
|
||||
(n + 1, quantized_max_memory + 1), dtype=torch.float32, device="cpu"
|
||||
)
|
||||
|
||||
for i in range(1, n + 1):
|
||||
current_memory = quantized_memory[i - 1]
|
||||
current_runtime = runtimes[i - 1]
|
||||
|
||||
# Copy the previous row
|
||||
dp[i, :] = dp[i - 1, :]
|
||||
|
||||
# Update dp[i, j] for all j >= current_memory
|
||||
if current_memory == 0:
|
||||
dp[i, :] = dp[i - 1, :] + current_runtime
|
||||
else:
|
||||
dp[i, current_memory:] = torch.maximum(
|
||||
dp[i - 1, current_memory:],
|
||||
dp[i - 1, :-current_memory] + current_runtime,
|
||||
)
|
||||
|
||||
# Backtrack to find the items included in the knapsack
|
||||
saved_items = []
|
||||
recomputable_items = []
|
||||
j: int = quantized_max_memory
|
||||
for i in range(n, 0, -1):
|
||||
if dp[i][j] != dp[i - 1][j]:
|
||||
saved_items.append(i - 1) # Include this item (indexing from 0)
|
||||
j -= int(quantized_memory[i - 1].item())
|
||||
else:
|
||||
recomputable_items.append(i - 1)
|
||||
|
||||
saved_items.reverse() # To get items in the order they were added
|
||||
|
||||
# The maximum runtime that can be achieved within the max_memory constraint
|
||||
max_runtime = dp[n][quantized_max_memory].item()
|
||||
|
||||
return max_runtime, saved_items, recomputable_items
|
||||
|
||||
|
||||
def dp_knapsack_sliding_hirschberg(
|
||||
memory: list[float], runtime: list[float], max_memory: float
|
||||
) -> tuple[float, list[int], list[int]]:
|
||||
# Scaling factor to convert floating point weights to integers
|
||||
S = 10000
|
||||
|
||||
# q_ prefix stands for quantized
|
||||
q_memory = [int(round(m * S)) for m in memory]
|
||||
runtimes = [float(v) for v in runtime]
|
||||
|
||||
q_max_memory = int(round(max_memory * S))
|
||||
|
||||
q_memory_length = len(q_memory)
|
||||
if q_memory_length == 0:
|
||||
return 0.0, [], []
|
||||
|
||||
item_indices = list(range(q_memory_length))
|
||||
dp_profile_size = q_max_memory + 1
|
||||
|
||||
# Current DP profile (row)
|
||||
dp_profile = torch.zeros(dp_profile_size, dtype=torch.float32, device="cpu")
|
||||
# Store a candidate for next dp_profile - current dp row + item
|
||||
candidate_profile = torch.empty(dp_profile_size, dtype=torch.float32, device="cpu")
|
||||
left_profile = torch.empty(dp_profile_size, dtype=torch.float32, device="cpu")
|
||||
right_profile = torch.empty(dp_profile_size, dtype=torch.float32, device="cpu")
|
||||
|
||||
saved_items: list[int] = []
|
||||
recomputable_items: list[int] = []
|
||||
|
||||
# Explicit stack to optimize memory and avoid recursion
|
||||
# Stack stores segments as (start index, end index, capacity for segment)
|
||||
stack: list[tuple[int, int, int]] = [(0, q_memory_length, q_max_memory)]
|
||||
|
||||
# LIFO
|
||||
while stack:
|
||||
start, end, capacity = stack.pop()
|
||||
length = end - start
|
||||
if length == 0:
|
||||
continue
|
||||
|
||||
# Leaf
|
||||
if length == 1:
|
||||
index = item_indices[start]
|
||||
memory_item = q_memory[index]
|
||||
runtime_item = runtimes[index]
|
||||
if memory_item <= capacity and runtime_item > 0.0:
|
||||
saved_items.append(index)
|
||||
else:
|
||||
recomputable_items.append(index)
|
||||
continue
|
||||
|
||||
# Split the segment into two halves
|
||||
middle = start + (length // 2)
|
||||
left_start, left_end = middle, end
|
||||
right_start, right_end = start, middle
|
||||
|
||||
# Assign items to both halves
|
||||
left_items = item_indices[left_start:left_end]
|
||||
right_items = item_indices[right_start:right_end]
|
||||
|
||||
# Working only on items allowed by segment's capacity
|
||||
capacity = capacity + 1
|
||||
dp_view = dp_profile[:capacity]
|
||||
candidate_view = candidate_profile[:capacity]
|
||||
left_dp_local = left_profile[:capacity]
|
||||
right_dp_local = right_profile[:capacity]
|
||||
|
||||
# Left part
|
||||
dp_view.zero_()
|
||||
for index in left_items:
|
||||
memory_item = q_memory[index]
|
||||
runtime_item = runtimes[index]
|
||||
|
||||
if memory_item == 0:
|
||||
# Weight is 0, so add it to all capacities; a "free lunch", essentially
|
||||
dp_view.add_(runtime_item)
|
||||
continue
|
||||
|
||||
# If item is too heavy, we skip it
|
||||
if memory_item >= capacity:
|
||||
continue
|
||||
|
||||
# Add the current item so we can then pick the highest value
|
||||
dp_view_candidate = candidate_view[: capacity - memory_item]
|
||||
torch.add(dp_view[:-memory_item], runtime_item, out=dp_view_candidate)
|
||||
# Take the highest - either previous (without current) or with current
|
||||
torch.maximum(
|
||||
dp_view[memory_item:], dp_view_candidate, out=dp_view[memory_item:]
|
||||
)
|
||||
|
||||
# Store the left profile
|
||||
left_dp_local.copy_(dp_view)
|
||||
|
||||
# Right part
|
||||
dp_view.zero_()
|
||||
for index in right_items:
|
||||
memory_item = q_memory[index]
|
||||
runtime_item = runtimes[index]
|
||||
|
||||
if memory_item == 0:
|
||||
dp_view.add_(runtime_item)
|
||||
continue
|
||||
|
||||
if memory_item >= capacity:
|
||||
continue
|
||||
|
||||
dp_view_candidate = candidate_view[: capacity - memory_item]
|
||||
torch.add(dp_view[:-memory_item], runtime_item, out=dp_view_candidate)
|
||||
torch.maximum(
|
||||
dp_view[memory_item:], dp_view_candidate, out=dp_view[memory_item:]
|
||||
)
|
||||
|
||||
# Store the reversed right profile
|
||||
right_dp_local.copy_(dp_view.flip(-1))
|
||||
|
||||
# In-place compute item-wise sum of left and right to pick the split point where the sum is highest
|
||||
left_dp_local.add_(right_dp_local)
|
||||
|
||||
# Pick the index of highest value of a pair, which we then use as a split point
|
||||
best_split = int(torch.argmax(left_dp_local).item())
|
||||
|
||||
left_capacity = best_split
|
||||
right_capacity = capacity - best_split
|
||||
|
||||
# Clamp (might be removed if we're 100% sure that there is no edge case that will mess up the indices math)
|
||||
if left_capacity < 0:
|
||||
left_capacity = 0
|
||||
if right_capacity < 0:
|
||||
right_capacity = 0
|
||||
if left_capacity > q_max_memory:
|
||||
left_capacity = q_max_memory
|
||||
if right_capacity > q_max_memory:
|
||||
right_capacity = q_max_memory
|
||||
|
||||
# Push right then left, so left is processed next
|
||||
stack.append((right_start, right_end, right_capacity))
|
||||
stack.append((left_start, left_end, left_capacity))
|
||||
|
||||
saved_items = sorted(saved_items)
|
||||
recomputable_items = sorted(recomputable_items)
|
||||
|
||||
max_runtime = sum(runtime[i] for i in saved_items)
|
||||
recomputable_items.reverse()
|
||||
return max_runtime, saved_items, recomputable_items
|
||||
+282
@@ -0,0 +1,282 @@
|
||||
import operator
|
||||
from collections import deque
|
||||
from collections.abc import Callable
|
||||
|
||||
import networkx as nx
|
||||
|
||||
from torch._functorch._activation_checkpointing.graph_info_provider import (
|
||||
GraphInfoProvider,
|
||||
)
|
||||
|
||||
|
||||
class KnapsackEvaluator:
|
||||
"""
|
||||
This class evaluates the theoretical runtime and peak memory usage of a given checkpointing strategy.
|
||||
It takes in a graph and a list of nodes that are saved and recomputed, and then simulates the
|
||||
backward pass to calculate the peak memory usage.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
graph_info_provider: GraphInfoProvider,
|
||||
) -> None:
|
||||
self._graph_info_provider = graph_info_provider
|
||||
|
||||
def _get_backward_memory_from_topologically_sorted_graph(
|
||||
self,
|
||||
node_graph: nx.DiGraph,
|
||||
node_memories: dict[str, float],
|
||||
saved_nodes_set: set[str],
|
||||
peak_memory_after_forward_pass: float,
|
||||
) -> list[tuple[float, str]]:
|
||||
"""
|
||||
Simulates the backward pass and keeps track of the peak memory usage.
|
||||
|
||||
High Level Steps:
|
||||
1. Set Initial Peak/Current Memory
|
||||
Allows you to set the peak memory after the forward pass, but typically this is
|
||||
the sum of the estimated memory of the saved nodes.
|
||||
2. Perform a reverse topological sort of the node_graph.
|
||||
If full graph is defined then will sort the full graph and only process the subset
|
||||
of nodes in the node_graph.
|
||||
3. Iterate through the sorted graph nodes.
|
||||
If the node is saved then just drop it's memory from current memory.
|
||||
If the node is not saved then add it's memory to current memory and then traverse it's
|
||||
predecessors to simulate recomuptation chain. Will check if new peak memory after all
|
||||
predecessors are processed.
|
||||
|
||||
Args:
|
||||
node_graph (nx.DiGraph): A directed graph representing the recomputable forward nodes.
|
||||
saved_nodes_set (Set[str]): A set of node names that are saved.
|
||||
peak_memory_after_forward_pass (float): The peak memory usage after the forward pass.
|
||||
"""
|
||||
current_memory = [
|
||||
(peak_memory_after_forward_pass, "Initial Peak/Current Memory")
|
||||
]
|
||||
already_computed = set()
|
||||
sorted_nodes = list(reversed(list(nx.topological_sort(node_graph))))
|
||||
dependencies_computed = set()
|
||||
|
||||
for node in sorted_nodes:
|
||||
if node in saved_nodes_set or node in already_computed:
|
||||
current_memory.append(
|
||||
(
|
||||
current_memory[-1][0] - node_memories[node],
|
||||
f"Dropping Node(already saved): {node}",
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
already_computed.add(node)
|
||||
current_memory.append(
|
||||
(
|
||||
current_memory[-1][0] + node_memories[node],
|
||||
f"Recomputing Node: {node}",
|
||||
)
|
||||
)
|
||||
# Create a queue of dependencies required for recomputation
|
||||
predecessor_queue = deque(
|
||||
[
|
||||
dependency
|
||||
# pyrefly: ignore [bad-unpacking]
|
||||
for dependency, v in node_graph.in_edges(node)
|
||||
if dependency not in already_computed
|
||||
]
|
||||
)
|
||||
while predecessor_queue:
|
||||
dep = predecessor_queue.popleft()
|
||||
already_computed.add(dep)
|
||||
dependencies_computed.add(dep)
|
||||
current_memory.append(
|
||||
(
|
||||
current_memory[-1][0] + node_memories[dep],
|
||||
f"Recomputing Predecessor of {node}: {dep}",
|
||||
)
|
||||
)
|
||||
# Add predecessors of the predecessor to the queue if they haven't been recomputed yet
|
||||
# pyrefly: ignore [bad-unpacking]
|
||||
for dependency_of_dependency, _ in node_graph.in_edges(dep):
|
||||
if (
|
||||
dependency_of_dependency in already_computed
|
||||
or dependency_of_dependency in saved_nodes_set
|
||||
or dependency_of_dependency in predecessor_queue
|
||||
):
|
||||
continue
|
||||
predecessor_queue.append(dependency_of_dependency)
|
||||
dependencies_computed.clear()
|
||||
current_memory.append(
|
||||
(current_memory[-1][0] - node_memories[node], f"Dropping Node: {node}")
|
||||
)
|
||||
return current_memory
|
||||
|
||||
def _validate_all_indexes_accounted_for_in_provided_output(
|
||||
self, saved_nodes_idxs: list[int], recomputable_node_idxs: list[int]
|
||||
) -> None:
|
||||
"""
|
||||
Validate that all indexes are accounted for in the provided output.
|
||||
This function checks that the union of saved nodes and recomputable nodes
|
||||
covers all candidate nodes without any overlaps.
|
||||
"""
|
||||
recomputable_node_idxs_set = set(recomputable_node_idxs)
|
||||
saved_nodes_idxs_set = set(saved_nodes_idxs)
|
||||
all_candidate_nodes_idxs = set(
|
||||
range(len(self._graph_info_provider.all_recomputable_banned_nodes))
|
||||
)
|
||||
# Check that there are no overlaps between saved nodes and recomputable nodes
|
||||
if len(recomputable_node_idxs_set.intersection(saved_nodes_idxs_set)) != 0:
|
||||
raise AssertionError(
|
||||
f"Saved nodes and recomputable nodes cannot have any overlaps, "
|
||||
f"but found overlap: {recomputable_node_idxs_set.intersection(saved_nodes_idxs_set)}"
|
||||
)
|
||||
# Check that all candidate nodes are accounted for
|
||||
if (
|
||||
recomputable_node_idxs_set.union(saved_nodes_idxs_set)
|
||||
!= all_candidate_nodes_idxs
|
||||
):
|
||||
raise AssertionError(
|
||||
f"All candidate nodes must be accounted for in the provided output, "
|
||||
f"got union={recomputable_node_idxs_set.union(saved_nodes_idxs_set)}, "
|
||||
f"expected={all_candidate_nodes_idxs}"
|
||||
)
|
||||
|
||||
def evaluate_knapsack_output(
|
||||
self,
|
||||
saved_nodes_idxs: list[int],
|
||||
recomputable_node_idxs: list[int],
|
||||
account_for_backward_pass: bool = False,
|
||||
) -> dict[str, float]:
|
||||
"""
|
||||
Evaluate the theoretical runtime and peak memory usage of a given checkpointing strategy.
|
||||
Args:
|
||||
- saved_nodes_idxs (List[int]): The indices of nodes that are saved.
|
||||
- recomputable_node_idxs (List[int]): The indices of nodes that need to be recomputed.
|
||||
"""
|
||||
self._validate_all_indexes_accounted_for_in_provided_output(
|
||||
saved_nodes_idxs, recomputable_node_idxs
|
||||
)
|
||||
recomputation_runtime = sum(
|
||||
self._graph_info_provider.all_node_runtimes[
|
||||
self._graph_info_provider.all_recomputable_banned_nodes[node]
|
||||
]
|
||||
for node in recomputable_node_idxs
|
||||
)
|
||||
if account_for_backward_pass:
|
||||
memory_list = self._get_backward_memory_from_topologically_sorted_graph(
|
||||
node_graph=self._graph_info_provider.recomputable_node_only_graph_with_larger_graph_context,
|
||||
saved_nodes_set={
|
||||
self._graph_info_provider.all_recomputable_banned_nodes[i]
|
||||
for i in saved_nodes_idxs
|
||||
},
|
||||
node_memories=self._graph_info_provider.all_node_memories,
|
||||
peak_memory_after_forward_pass=sum(
|
||||
self._graph_info_provider.all_node_memories[
|
||||
self._graph_info_provider.all_recomputable_banned_nodes[i]
|
||||
]
|
||||
for i in saved_nodes_idxs
|
||||
),
|
||||
)
|
||||
peak_memory = max(memory_list, key=operator.itemgetter(0))[0]
|
||||
else:
|
||||
peak_memory = sum(
|
||||
self._graph_info_provider.all_node_memories[
|
||||
self._graph_info_provider.all_recomputable_banned_nodes[node]
|
||||
]
|
||||
for node in saved_nodes_idxs
|
||||
)
|
||||
return {
|
||||
"peak_memory": peak_memory,
|
||||
"recomputation_runtime": recomputation_runtime,
|
||||
"non_ac_peak_memory": self._graph_info_provider.get_non_ac_peak_memory(),
|
||||
"theoretical_max_runtime": self._graph_info_provider.get_theoretical_max_runtime(),
|
||||
"percentage_of_theoretical_peak_memory": peak_memory
|
||||
/ self._graph_info_provider.get_non_ac_peak_memory(),
|
||||
"percentage_of_theoretical_peak_runtime": recomputation_runtime
|
||||
/ self._graph_info_provider.get_theoretical_max_runtime(),
|
||||
}
|
||||
|
||||
def evaluate_distribution_of_results_for_knapsack_algo(
|
||||
self,
|
||||
knapsack_algo: Callable[
|
||||
[list[float], list[float], float], tuple[float, list[int], list[int]]
|
||||
],
|
||||
memory_budget_values: list[float],
|
||||
) -> list[dict[str, float]]:
|
||||
"""
|
||||
Evaluates the distribution of results for a given knapsack algorithm.
|
||||
Args:
|
||||
knapsack_algo (Callable): The knapsack algorithm to use for evaluation.
|
||||
memory_budget_values (List[float]): A list of memory budgets to evaluate.
|
||||
"""
|
||||
results = []
|
||||
for memory_budget in memory_budget_values:
|
||||
_, saved_nodes, recomputed_nodes = knapsack_algo(
|
||||
self._graph_info_provider.get_knapsack_memory_input(),
|
||||
self._graph_info_provider.get_knapsack_runtime_input(),
|
||||
memory_budget,
|
||||
)
|
||||
result = self.evaluate_knapsack_output(
|
||||
saved_nodes_idxs=saved_nodes,
|
||||
recomputable_node_idxs=recomputed_nodes,
|
||||
)
|
||||
result["memory_budget"] = memory_budget
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def get_knee_point_memory_budget(
|
||||
self,
|
||||
knapsack_algo: Callable[
|
||||
[list[float], list[float], float], tuple[float, list[int], list[int]]
|
||||
],
|
||||
max_mem_budget: float = 0.1,
|
||||
min_mem_budget: float = 0.001,
|
||||
iterations: int = 100,
|
||||
) -> float:
|
||||
"""
|
||||
Finds the memory budget at the knee point in the Pareto frontier.
|
||||
|
||||
The knee point is defined as the point where the trade-off between
|
||||
runtime and memory usage is optimal.
|
||||
|
||||
Args:
|
||||
knapsack_algo (callable): Knapsack algorithm to use for evaluation.
|
||||
max_mem_budget (float, optional): Maximum memory budget. Defaults to 0.1.
|
||||
min_mem_budget (float, optional): Minimum memory budget. Defaults to 0.001.
|
||||
iterations (int, optional): Number of memory budgets to evaluate. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
float: Memory budget at the knee point.
|
||||
"""
|
||||
results = self.evaluate_distribution_of_results_for_knapsack_algo(
|
||||
knapsack_algo=knapsack_algo,
|
||||
memory_budget_values=[
|
||||
min_mem_budget
|
||||
+ i * (max_mem_budget - min_mem_budget) / (iterations - 1)
|
||||
for i in range(iterations)
|
||||
],
|
||||
)
|
||||
runtime_values = [
|
||||
result["percentage_of_theoretical_peak_runtime"] for result in results
|
||||
]
|
||||
memory_values = [
|
||||
result["percentage_of_theoretical_peak_memory"] for result in results
|
||||
]
|
||||
runtime_range = max(runtime_values) - min(runtime_values)
|
||||
memory_range = max(memory_values) - min(memory_values)
|
||||
if runtime_range == 0 or memory_range == 0:
|
||||
return max_mem_budget
|
||||
|
||||
# Normalize values
|
||||
runtime_min = min(runtime_values)
|
||||
memory_min = min(memory_values)
|
||||
runtime_norm = [
|
||||
(value - runtime_min) / runtime_range for value in runtime_values
|
||||
]
|
||||
memory_norm = [(value - memory_min) / memory_range for value in memory_values]
|
||||
# Calculate Euclidean distance
|
||||
distances = [
|
||||
(runtime_norm[i] ** 2 + memory_norm[i] ** 2) ** 0.5
|
||||
for i in range(len(runtime_norm))
|
||||
]
|
||||
# Find the knee point(shortest distance from the origin)
|
||||
knee_index = distances.index(min(distances))
|
||||
return results[knee_index]["memory_budget"]
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
"""AC rematerialize pass: Duplicates recompute nodes for backward, then DCE removes unused forward versions."""
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
from typing import Any, overload
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch._functorch.compile_utils import raise_getitems
|
||||
from torch._functorch.partitioners import (
|
||||
cleanup_recompute_tags,
|
||||
force_save_bw_mutation_src,
|
||||
has_recomputable_ops,
|
||||
has_recomputable_rng_ops,
|
||||
is_not_collective,
|
||||
must_recompute,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
_EMPTY_CUSTOM_META: dict[str, object] = {}
|
||||
|
||||
|
||||
def is_impure_node_for_dce(node: fx.Node) -> bool:
|
||||
# Check for special collectives that should be treated as pure
|
||||
if not is_not_collective(node):
|
||||
# It's a collective (wait_tensor, all_gather_into_tensor, etc.)
|
||||
# Treat as pure - can be eliminated if unused
|
||||
return False
|
||||
|
||||
# For everything else, fall back to the DEFAULT logic
|
||||
# This is what eliminate_dead_code() calls when is_impure_node=None
|
||||
impure_random = True
|
||||
if torch._guards.TracingContext.try_get():
|
||||
impure_random = torch._inductor.config.fallback_random
|
||||
return node.is_impure(impure_random)
|
||||
|
||||
|
||||
def _is_backward_node(node: fx.Node, use_phase: bool = False) -> bool:
|
||||
"""Check if node is in backward region.
|
||||
|
||||
If use_phase is True, only checks custom["phase"] == "backward"
|
||||
(user annotation). Otherwise falls back to node.meta["autograd_backward"],
|
||||
which Dynamo adds when tracing torch.autograd.grad.
|
||||
"""
|
||||
custom = node.meta.get("custom", _EMPTY_CUSTOM_META)
|
||||
if use_phase:
|
||||
return custom.get("phase") == "backward"
|
||||
return node.meta.get("autograd_backward", False)
|
||||
|
||||
|
||||
def _has_user_phase_annotation(gm: fx.GraphModule) -> bool:
|
||||
"""Check if any node has the user-level phase: backward annotation."""
|
||||
return any(
|
||||
node.meta.get("custom", _EMPTY_CUSTOM_META).get("phase") == "backward"
|
||||
for node in gm.graph.nodes
|
||||
)
|
||||
|
||||
|
||||
def _collect_backward_region_data(
|
||||
gm: fx.GraphModule,
|
||||
) -> tuple[bool, int | None, int | None, int]:
|
||||
use_phase = _has_user_phase_annotation(gm)
|
||||
bwd_start: int | None = None
|
||||
bwd_end: int | None = None
|
||||
num_regions = 0
|
||||
in_backward = False
|
||||
|
||||
for idx, node in enumerate(gm.graph.nodes):
|
||||
is_bwd = _is_backward_node(node, use_phase=use_phase)
|
||||
if is_bwd:
|
||||
if bwd_start is None:
|
||||
bwd_start = idx
|
||||
bwd_end = idx + 1
|
||||
if not in_backward:
|
||||
num_regions += 1
|
||||
in_backward = is_bwd
|
||||
|
||||
return use_phase, bwd_start, bwd_end, num_regions
|
||||
|
||||
|
||||
def remat_using_tags_for_fwd_loss_bwd_graph(gm: fx.GraphModule) -> fx.GraphModule:
|
||||
"""
|
||||
Duplicate recompute nodes for backward use. DCE removes unused forward versions.
|
||||
|
||||
Backward regions are identified by custom["phase"] == "backward" (user
|
||||
annotation) or node.meta["autograd_backward"] == True (set automatically when
|
||||
Dynamo traces torch.autograd.grad). When the user provides phase
|
||||
annotations, only those annotated regions are used.
|
||||
|
||||
Only a single contiguous backward region is supported. If multiple disjoint
|
||||
backward regions are detected, an error is raised. Consecutive backward
|
||||
operations without non-backward nodes between them are treated as a single
|
||||
backward region.
|
||||
"""
|
||||
if not has_recomputable_ops(gm):
|
||||
return gm
|
||||
|
||||
use_phase, bwd_start, bwd_end, num_regions = _collect_backward_region_data(gm)
|
||||
if num_regions > 1:
|
||||
if use_phase:
|
||||
raise RuntimeError(
|
||||
f"Detected {num_regions} disjoint backward regions annotated with "
|
||||
'phase: "backward" but remat only supports a single backward region. '
|
||||
"Please ensure only one contiguous region is annotated."
|
||||
)
|
||||
raise RuntimeError(
|
||||
f"Detected {num_regions} disjoint backward regions in the graph but remat only supports "
|
||||
"a single backward region. This can happen when non-backward computation appears "
|
||||
"between backward sections. Please annotate the real backward with "
|
||||
'torch.fx.traceback.annotate({"phase": "backward"}).'
|
||||
)
|
||||
|
||||
if bwd_start is None:
|
||||
return gm
|
||||
if bwd_end is None:
|
||||
raise AssertionError(
|
||||
"backward region should end somewhere when there was explicit backward region start."
|
||||
)
|
||||
|
||||
order = {node: idx for idx, node in enumerate(gm.graph.nodes)}
|
||||
|
||||
if has_recomputable_rng_ops(gm):
|
||||
raise RuntimeError(
|
||||
"Activation checkpoint rematerialization in `forward-loss-backward` graph does not support RNG ops "
|
||||
"in recompute regions. Please move RNG operations outside "
|
||||
"of recompute regions, or use joint graph mode (where partitioner handles RNG)."
|
||||
)
|
||||
|
||||
# Use partitioner pass to normalize AC node tags.
|
||||
gm = cleanup_recompute_tags(gm, is_default_partition=True)
|
||||
|
||||
force_save_bw_mutation_src(gm)
|
||||
|
||||
new_graph = fx.Graph()
|
||||
env: dict[fx.Node, fx.Node] = {}
|
||||
recomputed_nodes: dict[fx.Node, fx.Node] = {}
|
||||
|
||||
# Insert forward nodes
|
||||
for node in itertools.islice(gm.graph.nodes, 0, bwd_start):
|
||||
env[node] = new_graph.node_copy(node, lambda x: env[x])
|
||||
|
||||
@overload
|
||||
def remat_input(x: fx.Node) -> fx.Node: ...
|
||||
@overload
|
||||
def remat_input(x: Any) -> Any: ...
|
||||
|
||||
def remat_input(x: object) -> object:
|
||||
# fx.Node can have args that are primitive types (e.g. int, float, bool)
|
||||
if not isinstance(x, fx.Node):
|
||||
return x
|
||||
return recomputed_nodes.get(x, env[x])
|
||||
|
||||
def gather_recompute_deps(node: fx.Node) -> set[fx.Node]:
|
||||
deps: set[fx.Node] = set()
|
||||
|
||||
def _gather(n: fx.Node) -> None:
|
||||
if n in deps or n in recomputed_nodes or not must_recompute(n):
|
||||
return
|
||||
deps.add(n)
|
||||
for inp in n.all_input_nodes:
|
||||
_gather(inp)
|
||||
|
||||
# Can't call _gather(node) directly: node itself may not be must_recompute
|
||||
# (e.g. backward nodes), so _gather would return early without visiting inputs.
|
||||
for inp in node.all_input_nodes:
|
||||
_gather(inp)
|
||||
return deps
|
||||
|
||||
# Insert backward nodes
|
||||
for node in itertools.islice(gm.graph.nodes, bwd_start, bwd_end):
|
||||
# Gather all deps that need to be recomputed for this node
|
||||
deps = gather_recompute_deps(node)
|
||||
|
||||
# Insert deps in forward order (guaranteed disjoint from already-inserted)
|
||||
# This is not as inefficient as it looks, because we only add fresh dependencies
|
||||
# when they are not yet processed as recomputed nodes.
|
||||
new_deps = sorted(deps, key=lambda n: order[n])
|
||||
if new_deps:
|
||||
log.debug(
|
||||
"To compute backward node %s, recomputing [%s]",
|
||||
node.name,
|
||||
", ".join(dep.name for dep in new_deps),
|
||||
)
|
||||
for dep in new_deps:
|
||||
dup = new_graph.node_copy(dep, remat_input)
|
||||
dup.name = dep.name + "_recomputed"
|
||||
recomputed_nodes[dep] = dup
|
||||
|
||||
env[node] = new_graph.node_copy(node, remat_input)
|
||||
|
||||
for node in itertools.islice(gm.graph.nodes, bwd_end, None):
|
||||
env[node] = new_graph.node_copy(node, lambda x: env[x])
|
||||
|
||||
new_gm = torch.fx.GraphModule(gm, new_graph)
|
||||
|
||||
# DCE with custom is_impure_node (like default_partition)
|
||||
# Treats certain collectives as pure while delegating to default impurity logic
|
||||
new_gm.graph.eliminate_dead_code(is_impure_node=is_impure_node_for_dce)
|
||||
|
||||
# raise_getitems pass for better memory (like default_partition)
|
||||
new_gm = raise_getitems(new_gm)
|
||||
|
||||
new_gm.recompile()
|
||||
|
||||
return new_gm
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
+813
@@ -0,0 +1,813 @@
|
||||
"""Activation offloading for memory optimization during compilation.
|
||||
|
||||
This module provides functionality to offload activations to CPU during the forward
|
||||
pass and reload them during the backward pass, reducing GPU memory usage. It can be
|
||||
applied to graphs produced by both AOT Autograd partitioners and make_fx-based tracing.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import operator
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch._functorch._activation_offloading.offload_ops import ( # noqa: F401 -- registers ao::offload, ao::reload, ao::wait_tensor ops
|
||||
offload,
|
||||
reload,
|
||||
wait_tensor,
|
||||
)
|
||||
from torch._inductor.fx_passes.overlap_scheduling import benchmark_node, is_compute_node
|
||||
from torch._subclasses.fake_tensor import extract_tensor_metadata
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .. import config
|
||||
from ..partitioners import _size_of, get_default_op_list, OpTypes
|
||||
|
||||
|
||||
log: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Node name prefixes for offload/reload operations
|
||||
# NOTE: right now we are using these prefixes as identifiers for offload/reload
|
||||
CPU_OFFLOAD_PREFIX = "cpu_offload_"
|
||||
GPU_RELOAD_PREFIX = "gpu_reload_"
|
||||
|
||||
|
||||
def _find_all_effective_users(node: fx.Node, op_types: OpTypes) -> OrderedSet[fx.Node]:
|
||||
"""Find all effective users of a node, where view ops extend the lifetime
|
||||
of the original node. If a user is a view op, recursively find users of
|
||||
the view."""
|
||||
effective_users: OrderedSet[fx.Node] = OrderedSet()
|
||||
for user in node.users:
|
||||
if user.op == "output":
|
||||
continue
|
||||
effective_users.add(user)
|
||||
if op_types.is_view(user):
|
||||
effective_users.update(_find_all_effective_users(user, op_types))
|
||||
return effective_users
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReloadNodeInfo:
|
||||
"""
|
||||
Information about backward reload related nodes for each reload operation.
|
||||
|
||||
Pattern: ao.reload → ao.wait_tensor
|
||||
|
||||
- Reload group (ao.reload): Performs the actual asynchronous data transfer.
|
||||
Can be moved earlier in the graph to overlap with computation.
|
||||
- Wait node (ao.wait_tensor): Synchronization point that blocks until the data
|
||||
transfer completes. Must remain at the point where the data is first needed.
|
||||
"""
|
||||
|
||||
reload_group_nodes: list[fx.Node]
|
||||
wait_event_node: fx.Node
|
||||
transfer_size_bytes: int
|
||||
transfer_time_ms: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReloadQueueEntry:
|
||||
"""
|
||||
Entry in the reload queue for prefetch scheduling.
|
||||
|
||||
Attributes:
|
||||
pattern: The reload pattern information
|
||||
remaining_time_ms: Remaining overlap time needed in milliseconds
|
||||
"""
|
||||
|
||||
pattern: ReloadNodeInfo
|
||||
remaining_time_ms: float
|
||||
|
||||
|
||||
def offload_activation_fw(graph: fx.Graph) -> None:
|
||||
"""
|
||||
Insert CPU offload operations in the forward pass graph.
|
||||
|
||||
Offload operations are placed after the last effective use of each tensor marked
|
||||
for offloading. This ensures the tensor is no longer needed on the GPU before
|
||||
transferring it to CPU memory.
|
||||
|
||||
NOTE: An alternative approach would offload tensors immediately after generation
|
||||
to maximize compute-communication overlap. However, this requires additional
|
||||
synchronization to ensure tensor deletion (which occurs on the default stream)
|
||||
waits for the asynchronous offload operation to complete. This would necessitate
|
||||
more complex tracking to separate operation scheduling from memory cleanup.
|
||||
|
||||
Args:
|
||||
graph: The forward graph to modify
|
||||
"""
|
||||
|
||||
op_types: OpTypes = get_default_op_list()
|
||||
|
||||
output_node: fx.Node = graph.find_nodes(op="output")[0]
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
fwd_outputs: tuple[fx.Node, ...] = output_node.args[
|
||||
0
|
||||
] # pyrefly: ignore [bad-assignment]
|
||||
node_to_offload: dict[fx.Node, fx.Node] = dict()
|
||||
node_to_index: dict[fx.Node, int] = {
|
||||
node: idx for idx, node in enumerate(graph.nodes)
|
||||
}
|
||||
|
||||
for node in fwd_outputs:
|
||||
if node.meta.get("saved_for_offloading", False) is False:
|
||||
continue
|
||||
|
||||
# Find insertion point, which is the last use
|
||||
if all_effective_users := _find_all_effective_users(node, op_types):
|
||||
last_user = max(all_effective_users, key=lambda n: node_to_index[n])
|
||||
else:
|
||||
last_user: fx.Node = node
|
||||
|
||||
# Insert the CPU offload operation after the last user
|
||||
with graph.inserting_after(last_user):
|
||||
cpu_node: fx.Node = graph.call_function(
|
||||
torch.ops.prims.device_put.default,
|
||||
args=(node, torch.device("cpu")),
|
||||
kwargs={"non_blocking": True},
|
||||
name=CPU_OFFLOAD_PREFIX + str(node.name),
|
||||
)
|
||||
cpu_node.meta["val"] = node.meta["val"].to(torch.device("cpu"))
|
||||
cpu_node.meta["tensor_meta"] = extract_tensor_metadata(cpu_node.meta["val"])
|
||||
|
||||
node_to_offload[node] = cpu_node
|
||||
|
||||
# Update the return node args
|
||||
output_node.update_arg(
|
||||
0, tuple(node_to_offload.get(node, node) for node in fwd_outputs)
|
||||
)
|
||||
|
||||
|
||||
def reload_activation_bw(graph: fx.Graph) -> None:
|
||||
"""
|
||||
Insert GPU reload operations in the backward pass graph.
|
||||
|
||||
Reload operations are placed before the first use of each offloaded tensor,
|
||||
transferring it from CPU back to GPU memory before it's needed for computation.
|
||||
|
||||
Args:
|
||||
graph: The backward graph to modify
|
||||
"""
|
||||
|
||||
node_to_index: dict[fx.Node, int] = {
|
||||
node: idx for idx, node in enumerate(graph.nodes)
|
||||
}
|
||||
output_node: fx.Node = graph.find_nodes(op="output")[0]
|
||||
|
||||
for node in graph.find_nodes(op="placeholder"):
|
||||
if node.meta.get("saved_for_offloading", False) is False:
|
||||
continue
|
||||
|
||||
# Find insertion point, which is the first use or output node if no users
|
||||
# The later should not happen, but inserting before output node is safe
|
||||
insert_point: fx.Node = (
|
||||
min(node.users.keys(), key=lambda n: node_to_index[n])
|
||||
if node.users
|
||||
else output_node
|
||||
)
|
||||
|
||||
# Insert the GPU reload operation before the first user
|
||||
original_device: torch.Device = node.meta["original_device"]
|
||||
with graph.inserting_before(insert_point):
|
||||
gpu_node: fx.Node = graph.call_function(
|
||||
torch.ops.prims.device_put.default,
|
||||
args=(node, original_device),
|
||||
kwargs={"non_blocking": True},
|
||||
name=str(node.name).replace(CPU_OFFLOAD_PREFIX, GPU_RELOAD_PREFIX),
|
||||
)
|
||||
gpu_node.meta["val"] = node.meta["val"].to(original_device)
|
||||
gpu_node.meta["tensor_meta"] = extract_tensor_metadata(gpu_node.meta["val"])
|
||||
|
||||
# Replace all uses of the CPU tensor with the GPU tensor
|
||||
for user in list(node.users.keys()):
|
||||
if user != gpu_node:
|
||||
user.replace_input_with(node, gpu_node)
|
||||
|
||||
|
||||
def offload_activation_fw_async(graph: fx.Graph) -> None:
|
||||
"""Insert async CPU offload operations in the forward pass graph.
|
||||
|
||||
Uses ao.offload + ao.wait_tensor ops which encapsulate stream management
|
||||
internally, producing a clean 2-node IR per offloaded tensor.
|
||||
"""
|
||||
|
||||
op_types: OpTypes = get_default_op_list()
|
||||
|
||||
output_node: fx.Node = graph.find_nodes(op="output")[0]
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
fwd_outputs: tuple[fx.Node, ...] = output_node.args[
|
||||
0
|
||||
] # pyrefly: ignore [bad-assignment]
|
||||
node_to_offload: dict[fx.Node, fx.Node] = dict()
|
||||
node_to_index: dict[fx.Node, int] = {
|
||||
node: idx for idx, node in enumerate(graph.nodes)
|
||||
}
|
||||
|
||||
if not any(n.meta.get("saved_for_offloading", False) for n in fwd_outputs):
|
||||
return
|
||||
|
||||
for node in fwd_outputs:
|
||||
if node.meta.get("saved_for_offloading", False) is False:
|
||||
continue
|
||||
|
||||
if all_effective_users := _find_all_effective_users(node, op_types):
|
||||
last_user = max(all_effective_users, key=lambda n: node_to_index[n])
|
||||
else:
|
||||
last_user: fx.Node = node
|
||||
|
||||
with graph.inserting_after(last_user):
|
||||
offload_node: fx.Node = graph.call_function(
|
||||
torch.ops.ao.offload.default,
|
||||
args=(node,),
|
||||
name=f"async_{CPU_OFFLOAD_PREFIX}{node.name}",
|
||||
)
|
||||
offload_node.meta["val"] = node.meta["val"].to(torch.device("cpu"))
|
||||
offload_node.meta["tensor_meta"] = extract_tensor_metadata(
|
||||
offload_node.meta["val"]
|
||||
)
|
||||
# The keepalive=node arg extends the GPU tensor's lifetime in the
|
||||
# graph so the allocator doesn't reclaim it before the async D2H
|
||||
# copy completes.
|
||||
with graph.inserting_after(offload_node):
|
||||
wait_node: fx.Node = graph.call_function(
|
||||
torch.ops.ao.wait_tensor.default,
|
||||
args=(offload_node, node),
|
||||
name=CPU_OFFLOAD_PREFIX + str(node.name),
|
||||
)
|
||||
wait_node.meta["val"] = offload_node.meta["val"]
|
||||
wait_node.meta["tensor_meta"] = offload_node.meta["tensor_meta"]
|
||||
|
||||
node_to_offload[node] = wait_node
|
||||
|
||||
output_node.update_arg(
|
||||
0, tuple(node_to_offload.get(node, node) for node in fwd_outputs)
|
||||
)
|
||||
|
||||
|
||||
def reload_activation_bw_async(graph: fx.Graph) -> None:
|
||||
"""Insert async GPU reload operations in the backward pass graph.
|
||||
|
||||
Uses ao.reload + ao.wait_tensor ops which encapsulate stream management internally,
|
||||
producing a clean 2-node IR per reloaded tensor.
|
||||
"""
|
||||
|
||||
node_to_index: dict[fx.Node, int] = {
|
||||
node: idx for idx, node in enumerate(graph.nodes)
|
||||
}
|
||||
|
||||
nodes_to_reload = [
|
||||
n
|
||||
for n in graph.find_nodes(op="placeholder")
|
||||
if n.meta.get("saved_for_offloading", False)
|
||||
]
|
||||
if not nodes_to_reload:
|
||||
return
|
||||
|
||||
for node in nodes_to_reload:
|
||||
if not node.users:
|
||||
raise RuntimeError(
|
||||
f"Offloaded tensor {node.name} has no users in the backward graph"
|
||||
)
|
||||
insert_point: fx.Node = min(node.users.keys(), key=lambda n: node_to_index[n])
|
||||
|
||||
original_device: torch.device = node.meta["original_device"]
|
||||
with graph.inserting_before(insert_point):
|
||||
reload_node: fx.Node = graph.call_function(
|
||||
torch.ops.ao.reload.default,
|
||||
args=(node, original_device),
|
||||
name=f"async_{str(node.name).replace(CPU_OFFLOAD_PREFIX, GPU_RELOAD_PREFIX)}",
|
||||
)
|
||||
reload_node.meta["val"] = node.meta["val"].to(original_device)
|
||||
reload_node.meta["tensor_meta"] = extract_tensor_metadata(
|
||||
reload_node.meta["val"]
|
||||
)
|
||||
wait_node: fx.Node = graph.call_function(
|
||||
torch.ops.ao.wait_tensor.default,
|
||||
args=(reload_node,),
|
||||
name=str(node.name).replace(CPU_OFFLOAD_PREFIX, GPU_RELOAD_PREFIX),
|
||||
)
|
||||
wait_node.meta["val"] = reload_node.meta["val"]
|
||||
wait_node.meta["tensor_meta"] = reload_node.meta["tensor_meta"]
|
||||
|
||||
for user in list(node.users.keys()):
|
||||
if user != reload_node:
|
||||
user.replace_input_with(node, wait_node)
|
||||
|
||||
|
||||
def can_offload(
|
||||
node: fx.Node,
|
||||
fwd_outputs: OrderedSet[fx.Node],
|
||||
model_outputs: OrderedSet[fx.Node],
|
||||
static_lifetime_input_nodes: OrderedSet[fx.Node],
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if a node can be offloaded to CPU.
|
||||
|
||||
Args:
|
||||
node: The node to check
|
||||
fwd_outputs: Forward module outputs, including model outputs and activations
|
||||
model_outputs: Model outputs
|
||||
|
||||
NOTE: Additional context for the logic behind these offloading checks:
|
||||
|
||||
* fwd_outputs: Only saved intermediate tensors should be offloaded.
|
||||
|
||||
* model_outputs / static_lifetime_input_nodes: Tensors that may be accessed outside
|
||||
the compiled region (e.g., model outputs, static inputs) cannot be offloaded as
|
||||
they must remain accessible beyond the scope of the compiled graph.
|
||||
|
||||
* views / getitems: Offloading such nodes can lead to segmentation faults.
|
||||
|
||||
* contiguous: Offloading non-contiguous tensors causes CPU-side stride changes
|
||||
during both forward and backward passes when using the Inductor backend. While
|
||||
these stride changes cancel each other out, they introduce significant compute
|
||||
overhead. This is due to the contiguity check in ir.py (see link below).
|
||||
TODO: This restriction could potentially be bypassed in the future.
|
||||
Reference: https://github.com/pytorch/pytorch/blob/44ac69388a4a5eb463dbd2a13f00d1e3b924566c/torch/_inductor/ir.py#L3214
|
||||
|
||||
Additional criteria to consider for offloading optimization:
|
||||
|
||||
* Tensor size: Small tensors may not fully utilize available bandwidth, reducing the
|
||||
efficiency gains from offloading.
|
||||
|
||||
* Position in forward/backward graph: Activations generated near the end of the forward
|
||||
pass are typically consumed near the beginning of the backward pass. Offloading such
|
||||
tensors may be counterproductive since they are quickly reloaded, not having sufficient
|
||||
time to overlap the transfer with computation.
|
||||
"""
|
||||
|
||||
log.debug(f"Checking node {node.name} for offloading...") # noqa: G004
|
||||
|
||||
op_types: OpTypes = get_default_op_list()
|
||||
|
||||
if node not in fwd_outputs:
|
||||
log.debug("\tSkipped! Can only offload nodes in fwd_module_outputs.")
|
||||
return False
|
||||
if node in model_outputs:
|
||||
log.debug("\tSkipped! Cannot offload model outputs.")
|
||||
return False
|
||||
if node in static_lifetime_input_nodes:
|
||||
log.debug("\tSkipped! Cannot offload static input nodes.")
|
||||
return False
|
||||
if op_types.is_view(node):
|
||||
log.debug("\tSkipped! Cannot offload views.")
|
||||
return False
|
||||
if node.target == operator.getitem:
|
||||
log.debug("\tSkipped! Cannot offload getitems.")
|
||||
return False
|
||||
if hasattr(node, "meta") and "val" in node.meta:
|
||||
if (
|
||||
isinstance(val := node.meta["val"], torch.Tensor)
|
||||
and not val.is_contiguous()
|
||||
):
|
||||
log.debug("\tSkipped! Cannot offload non-contiguous tensors.")
|
||||
return False
|
||||
|
||||
log.debug("\tGood!")
|
||||
return True
|
||||
|
||||
|
||||
def choose_offload_sets(
|
||||
fwd_module: fx.GraphModule,
|
||||
num_fwd_outputs: int,
|
||||
static_lifetime_input_nodes: OrderedSet[fx.Node],
|
||||
) -> bool:
|
||||
"""
|
||||
Decide which nodes will be offloaded based on the marked nodes and feasibility.
|
||||
Marks nodes with "saved_for_offloading" if they should and can be offloaded.
|
||||
|
||||
Args:
|
||||
fwd_module: Forward graph module
|
||||
bwd_module: Backward graph module
|
||||
num_fwd_outputs: Number of forward outputs
|
||||
|
||||
Returns:
|
||||
bool: Whether activation offloading should be performed
|
||||
"""
|
||||
|
||||
fwd_outputs: OrderedSet[fx.Node] = OrderedSet(
|
||||
fwd_module.graph.find_nodes(op="output")[0].args[0]
|
||||
)
|
||||
model_outputs: OrderedSet[fx.Node] = OrderedSet(
|
||||
fwd_module.graph.find_nodes(op="output")[0].args[0][:num_fwd_outputs]
|
||||
)
|
||||
|
||||
should_perform_offloading = False
|
||||
for node in fwd_module.graph.nodes:
|
||||
if node.meta.get("should_offload", False) and can_offload(
|
||||
node, fwd_outputs, model_outputs, static_lifetime_input_nodes
|
||||
):
|
||||
node.meta["saved_for_offloading"] = True
|
||||
node.meta["original_device"] = node.meta["val"].device
|
||||
should_perform_offloading = True
|
||||
|
||||
return should_perform_offloading
|
||||
|
||||
|
||||
def offload_chosen_sets(
|
||||
fwd_module: fx.GraphModule,
|
||||
bwd_module: fx.GraphModule,
|
||||
) -> None:
|
||||
"""
|
||||
Add offload and reload nodes to the forward and backward graphs.
|
||||
This function adds device_put operations without any stream handling.
|
||||
|
||||
Args:
|
||||
fwd_module: Forward module graph
|
||||
bwd_module: Backward module graph
|
||||
"""
|
||||
|
||||
# Add offload nodes in forward graph
|
||||
offload_activation_fw(fwd_module.graph)
|
||||
|
||||
# Update backward graph inputs to be offloaded tensors
|
||||
bwd_inputs: dict[str, fx.Node] = {
|
||||
node.name: node for node in bwd_module.graph.find_nodes(op="placeholder")
|
||||
}
|
||||
for fwd_node in fwd_module.graph.find_nodes(op="output")[0].args[0]:
|
||||
if CPU_OFFLOAD_PREFIX not in fwd_node.name:
|
||||
continue
|
||||
|
||||
bwd_node: fx.Node = bwd_inputs[fwd_node.name.replace(CPU_OFFLOAD_PREFIX, "")]
|
||||
with bwd_module.graph.inserting_after(bwd_node):
|
||||
bwd_offload_node: fx.Node = bwd_module.graph.placeholder(name=fwd_node.name)
|
||||
|
||||
bwd_offload_node.meta.update(fwd_node.meta)
|
||||
bwd_offload_node.meta["saved_for_offloading"] = True
|
||||
bwd_offload_node.meta["original_device"] = bwd_node.meta["val"].device
|
||||
bwd_node.replace_all_uses_with(bwd_offload_node)
|
||||
bwd_module.graph.erase_node(bwd_node)
|
||||
|
||||
# Add reload nodes in backward graph
|
||||
reload_activation_bw(bwd_module.graph)
|
||||
|
||||
|
||||
def offload_chosen_sets_async(
|
||||
fwd_module: fx.GraphModule,
|
||||
bwd_module: fx.GraphModule,
|
||||
) -> None:
|
||||
"""
|
||||
Add async offload and reload nodes using ao ops.
|
||||
|
||||
Uses ao.offload/ao.reload + ao.wait_tensor which encapsulate stream management,
|
||||
instead of device_put + explicit stream operations. Can be applied to
|
||||
partitioned forward/backward graphs or to a joint graph produced by make_fx.
|
||||
"""
|
||||
|
||||
offload_activation_fw_async(fwd_module.graph)
|
||||
|
||||
# Replace backward graph placeholders with their offloaded (CPU) counterparts.
|
||||
# For each offloaded forward output, find the matching backward input and swap
|
||||
# it with a new placeholder carrying the CPU tensor's metadata, then mark it
|
||||
# for reloading.
|
||||
bwd_inputs: dict[str, fx.Node] = {
|
||||
node.name: node for node in bwd_module.graph.find_nodes(op="placeholder")
|
||||
}
|
||||
for fwd_node in fwd_module.graph.find_nodes(op="output")[0].args[0]:
|
||||
if CPU_OFFLOAD_PREFIX not in fwd_node.name:
|
||||
continue
|
||||
|
||||
bwd_node: fx.Node = bwd_inputs[fwd_node.name.replace(CPU_OFFLOAD_PREFIX, "")]
|
||||
with bwd_module.graph.inserting_after(bwd_node):
|
||||
bwd_offload_node: fx.Node = bwd_module.graph.placeholder(name=fwd_node.name)
|
||||
|
||||
bwd_offload_node.meta.update(fwd_node.meta)
|
||||
bwd_offload_node.meta["saved_for_offloading"] = True
|
||||
bwd_offload_node.meta["original_device"] = bwd_node.meta["val"].device
|
||||
bwd_node.replace_all_uses_with(bwd_offload_node)
|
||||
bwd_module.graph.erase_node(bwd_node)
|
||||
|
||||
reload_activation_bw_async(bwd_module.graph)
|
||||
|
||||
|
||||
def activation_offload_sink_wait_async(fwd_module: fx.GraphModule) -> None:
|
||||
"""Sink ao.wait_tensor operations for offload completion to the end of the graph.
|
||||
|
||||
This allows computation to overlap with offload operations.
|
||||
|
||||
NOTE: Sinking waits to the end delays GPU memory release of the source
|
||||
tensor (kept alive via the wait's keepalive arg) until the end of the
|
||||
compiled graph. For per-layer compile this is fine (one layer's worth of
|
||||
memory), but for full-model compile this means offloaded GPU tensors are
|
||||
not freed until the entire forward pass completes.
|
||||
"""
|
||||
graph: fx.Graph = fwd_module.graph
|
||||
output_node: fx.Node = graph.find_nodes(op="output")[0]
|
||||
|
||||
wait_nodes_to_sink: list[fx.Node] = [
|
||||
node
|
||||
for node in graph.nodes
|
||||
if node.op == "call_function"
|
||||
and node.target == torch.ops.ao.wait_tensor.default
|
||||
and isinstance(node.args[0], fx.Node)
|
||||
and node.args[0].op == "call_function"
|
||||
and node.args[0].target == torch.ops.ao.offload.default
|
||||
]
|
||||
|
||||
# prepend moves the node from its current position (no manual removal needed)
|
||||
for wait_node in wait_nodes_to_sink:
|
||||
output_node.prepend(wait_node)
|
||||
|
||||
|
||||
def activation_reload_prefetch_async(bwd_module: fx.GraphModule) -> None:
|
||||
"""
|
||||
Prefetch backward reload operations by moving ao.reload nodes earlier
|
||||
in the graph to overlap data transfer with computation, while keeping
|
||||
ao.wait_tensor at its original position.
|
||||
"""
|
||||
graph: fx.Graph = bwd_module.graph
|
||||
nodes_list: list[fx.Node] = list(graph.nodes)
|
||||
|
||||
# Identify reload + wait pairs
|
||||
reload_patterns: dict[fx.Node, ReloadNodeInfo] = {}
|
||||
for node in graph.nodes:
|
||||
if not (
|
||||
node.op == "call_function" and node.target == torch.ops.ao.reload.default
|
||||
):
|
||||
continue
|
||||
wait_node = next(
|
||||
(u for u in node.users if u.target == torch.ops.ao.wait_tensor.default),
|
||||
None,
|
||||
)
|
||||
if wait_node is None:
|
||||
continue
|
||||
transfer_size_bytes: int = _calculate_transfer_size(node)
|
||||
transfer_time_ms: float = _estimate_transfer_time_in_ms(transfer_size_bytes)
|
||||
reload_patterns[node] = ReloadNodeInfo(
|
||||
reload_group_nodes=[node],
|
||||
wait_event_node=wait_node,
|
||||
transfer_size_bytes=transfer_size_bytes,
|
||||
transfer_time_ms=transfer_time_ms,
|
||||
)
|
||||
|
||||
reorder_for_prefetch(nodes_list, reload_patterns)
|
||||
|
||||
|
||||
def _calculate_transfer_size(device_put_node: fx.Node) -> int:
|
||||
"""Calculate the size in bytes of data being transferred."""
|
||||
|
||||
# ao.offload(tensor) -> tensor at args[0]
|
||||
# ao.reload(tensor, device) -> tensor at args[0]
|
||||
if device_put_node.target in (
|
||||
torch.ops.ao.offload.default,
|
||||
torch.ops.ao.reload.default,
|
||||
):
|
||||
return _size_of(device_put_node.args[0]) # pyrefly: ignore [bad-argument-type]
|
||||
raise ValueError(f"Unexpected transfer op: {device_put_node.target}")
|
||||
|
||||
|
||||
def _estimate_transfer_time_in_ms(transfer_size_bytes: int) -> float:
|
||||
"""Estimate transfer time in milliseconds based on size and bandwidth.
|
||||
|
||||
Uses config.activation_offload_cpu_gpu_bw (GB/s) which should be set by
|
||||
the user to match their hardware.
|
||||
"""
|
||||
return (
|
||||
transfer_size_bytes / (1024**3) * 1_000 / config.activation_offload_cpu_gpu_bw
|
||||
)
|
||||
|
||||
|
||||
def identify_reload_patterns(
|
||||
graph: fx.Graph, nodes_list: list[fx.Node], node_to_idx: dict[fx.Node, int]
|
||||
) -> dict[fx.Node, ReloadNodeInfo]:
|
||||
"""
|
||||
Identify backward reload patterns in the graph.
|
||||
|
||||
Pattern: fork → wait_stream → device_put → record_event → join → wait_event
|
||||
|
||||
This uses position-based matching since these nodes are inserted together in
|
||||
add_backward_reload_stream_ops() in a specific order. Since stream operations
|
||||
do not have data dependencies between them, they are unsuitable for subgroup
|
||||
pattern matching type of checks.
|
||||
|
||||
Returns a dict mapping device_put node to ReloadNodeInfo containing:
|
||||
- reload_group_nodes: fork → wait_stream → device_put → record_event → join
|
||||
- wait_event_node: the wait_event node
|
||||
- transfer_size_bytes: size of data being transferred
|
||||
- transfer_time_ms: estimated transfer time in milliseconds
|
||||
"""
|
||||
patterns: dict[fx.Node, ReloadNodeInfo] = {}
|
||||
|
||||
# Find all GPU reload device_put nodes whose inputs are placeholder nodes
|
||||
reload_nodes: list[fx.Node] = [
|
||||
node
|
||||
for node in graph.find_nodes(
|
||||
op="call_function", target=torch.ops.prims.device_put.default
|
||||
)
|
||||
if GPU_RELOAD_PREFIX in node.name
|
||||
and (
|
||||
node.args
|
||||
and isinstance(node.args[0], fx.Node)
|
||||
and node.args[0].op == "placeholder"
|
||||
)
|
||||
]
|
||||
|
||||
# Extract patterns for each reload device_put node
|
||||
for reload_node in reload_nodes:
|
||||
reload_node_idx: int = node_to_idx[reload_node]
|
||||
|
||||
fork_node: fx.Node = nodes_list[reload_node_idx - 2]
|
||||
wait_stream_node: fx.Node = nodes_list[reload_node_idx - 1]
|
||||
record_event_node: fx.Node = nodes_list[reload_node_idx + 1]
|
||||
join_node: fx.Node = nodes_list[reload_node_idx + 2]
|
||||
wait_event_node: fx.Node = nodes_list[reload_node_idx + 3]
|
||||
|
||||
# Validate the nodes are what we expect
|
||||
# Removed in follow-up commit
|
||||
_validate_pattern_nodes( # noqa: F821 # pyrefly: ignore [unknown-name]
|
||||
fork_node,
|
||||
wait_stream_node,
|
||||
record_event_node,
|
||||
join_node,
|
||||
wait_event_node,
|
||||
)
|
||||
|
||||
# Calculate transfer size and time
|
||||
transfer_size_bytes: int = _calculate_transfer_size(reload_node)
|
||||
transfer_time_ms: float = _estimate_transfer_time_in_ms(transfer_size_bytes)
|
||||
|
||||
patterns[reload_node] = ReloadNodeInfo(
|
||||
reload_group_nodes=[
|
||||
fork_node,
|
||||
wait_stream_node,
|
||||
reload_node,
|
||||
record_event_node,
|
||||
join_node,
|
||||
],
|
||||
wait_event_node=wait_event_node,
|
||||
transfer_size_bytes=transfer_size_bytes,
|
||||
transfer_time_ms=transfer_time_ms,
|
||||
)
|
||||
|
||||
return patterns
|
||||
|
||||
|
||||
def reorder_for_prefetch(
|
||||
nodes_list: list[fx.Node],
|
||||
reload_patterns: dict[fx.Node, ReloadNodeInfo],
|
||||
) -> None:
|
||||
"""
|
||||
Reorder nodes to prefetch reload operations by directly manipulating the graph.
|
||||
|
||||
This follows the algorithm as follows:
|
||||
- Go through nodes in reverse order
|
||||
- When encountering a reload pattern, add it to a queue with its transfer time
|
||||
- When encountering a compute node, use its runtime to satisfy overlap requirements
|
||||
- Place reload patterns when their overlap requirement is satisfied
|
||||
- When encountering placeholder nodes, flush queue as reloads cannot move before inputs
|
||||
"""
|
||||
|
||||
# Build a set of all nodes in reload groups for quick lookup
|
||||
reload_group_nodes_set: set[fx.Node] = set()
|
||||
for pattern in reload_patterns.values():
|
||||
reload_group_nodes_set.update(pattern.reload_group_nodes)
|
||||
|
||||
# Queue to hold reload group nodes waiting to be placed (FIFO)
|
||||
reload_queue: list[ReloadQueueEntry] = []
|
||||
|
||||
# Loop through nodes in reverse
|
||||
for node in reversed(nodes_list):
|
||||
if node.op == "output":
|
||||
continue
|
||||
elif node.op == "placeholder":
|
||||
# Flush queue - place all remaining reloads after the last placeholder
|
||||
while reload_queue:
|
||||
entry: ReloadQueueEntry = reload_queue.pop(0)
|
||||
for reload_group_node in reversed(entry.pattern.reload_group_nodes):
|
||||
node.append(reload_group_node)
|
||||
break
|
||||
elif node in reload_patterns:
|
||||
pattern: ReloadNodeInfo = reload_patterns[node]
|
||||
reload_queue.append(
|
||||
ReloadQueueEntry(
|
||||
pattern=pattern, remaining_time_ms=pattern.transfer_time_ms
|
||||
)
|
||||
)
|
||||
elif node in reload_group_nodes_set:
|
||||
continue
|
||||
else:
|
||||
if not reload_queue:
|
||||
continue
|
||||
compute_runtime_ms: float = (
|
||||
benchmark_node(node) if is_compute_node(node) else 0
|
||||
)
|
||||
reload_queue[0].remaining_time_ms -= compute_runtime_ms
|
||||
|
||||
# Pop and place reload if its remaining time is satisfied (<= 0)
|
||||
if reload_queue[0].remaining_time_ms <= 0:
|
||||
entry: ReloadQueueEntry = reload_queue.pop(0)
|
||||
for reload_group_node in entry.pattern.reload_group_nodes:
|
||||
node.prepend(reload_group_node)
|
||||
|
||||
|
||||
def activation_offload_sink_wait(fwd_module: fx.GraphModule) -> None:
|
||||
"""
|
||||
Sink wait_event operations for offload completion to the end of the graph.
|
||||
|
||||
This function identifies wait_event nodes for offload completion and moves them
|
||||
to the end of the graph, allowing computation to overlap with offload operations.
|
||||
|
||||
Args:
|
||||
fwd_module: Forward module graph
|
||||
"""
|
||||
graph: fx.Graph = fwd_module.graph
|
||||
nodes_list: list[fx.Node] = list(graph.nodes)
|
||||
node_to_idx: dict[fx.Node, int] = {node: idx for idx, node in enumerate(nodes_list)}
|
||||
|
||||
# Find all CPU offload device_put nodes
|
||||
offload_nodes: list[fx.Node] = [
|
||||
node
|
||||
for node in graph.find_nodes(
|
||||
op="call_function", target=torch.ops.prims.device_put.default
|
||||
)
|
||||
if CPU_OFFLOAD_PREFIX in node.name
|
||||
]
|
||||
|
||||
# Collect all wait_event nodes that need to be moved
|
||||
wait_nodes_to_sink: list[fx.Node] = []
|
||||
for offload_node in offload_nodes:
|
||||
offload_idx: int = node_to_idx[offload_node]
|
||||
wait_event_node: fx.Node = nodes_list[offload_idx + 3]
|
||||
|
||||
# Validate it's actually a wait_event node
|
||||
if not (
|
||||
wait_event_node.op == "call_function"
|
||||
and wait_event_node.target == torch.ops.streams.wait_event.default
|
||||
):
|
||||
raise ValueError(
|
||||
f"Expected wait_event node three positions after {offload_node.name}"
|
||||
)
|
||||
|
||||
wait_nodes_to_sink.append(wait_event_node)
|
||||
|
||||
# Find the output node, and move all wait_event nodes to just before the output node
|
||||
output_node: fx.Node = graph.find_nodes(op="output")[0]
|
||||
for wait_node in wait_nodes_to_sink:
|
||||
output_node.prepend(wait_node)
|
||||
|
||||
|
||||
def activation_reload_prefetch(bwd_module: fx.GraphModule) -> None:
|
||||
"""
|
||||
Prefetch backward reload operations by moving them earlier in the graph
|
||||
to overlap communication with computation.
|
||||
|
||||
This function identifies backward reload patterns (fork → wait_stream → device_put →
|
||||
record_event → join) and moves them earlier in the execution order to overlap
|
||||
the data transfer with computation, while keeping the wait_event at its original
|
||||
position.
|
||||
|
||||
Args:
|
||||
bwd_module: Backward module graph
|
||||
"""
|
||||
graph: fx.Graph = bwd_module.graph
|
||||
nodes_list: list[fx.Node] = list(graph.nodes)
|
||||
node_to_idx: dict[fx.Node, int] = {node: idx for idx, node in enumerate(nodes_list)}
|
||||
|
||||
# Step 1: Identify reload patterns
|
||||
reload_patterns: dict[fx.Node, ReloadNodeInfo] = identify_reload_patterns(
|
||||
graph, nodes_list, node_to_idx
|
||||
)
|
||||
|
||||
# Step 2: Reorder nodes by directly manipulating the graph
|
||||
reorder_for_prefetch(nodes_list, reload_patterns)
|
||||
|
||||
|
||||
def enable_activation_offloading(
|
||||
fwd_module: fx.GraphModule,
|
||||
bwd_module: fx.GraphModule,
|
||||
num_fwd_outputs: int,
|
||||
static_lifetime_input_nodes: OrderedSet[fx.Node],
|
||||
) -> None:
|
||||
"""
|
||||
Main entry point for activation offloading.
|
||||
|
||||
Args:
|
||||
fwd_module: Forward module graph
|
||||
bwd_module: Backward module graph
|
||||
num_fwd_outputs: Number of forward outputs
|
||||
"""
|
||||
|
||||
# Step 1: Decide which nodes to offload and mark them
|
||||
should_perform_offloading: bool = choose_offload_sets(
|
||||
fwd_module,
|
||||
num_fwd_outputs,
|
||||
static_lifetime_input_nodes,
|
||||
)
|
||||
if not should_perform_offloading:
|
||||
return
|
||||
|
||||
# Step 2: Add offload and reload nodes to the graphs
|
||||
if config.activation_offload_separate_stream:
|
||||
# Use async ao ops (2 nodes each: offload/reload + wait_tensor)
|
||||
offload_chosen_sets_async(fwd_module, bwd_module)
|
||||
if config.activation_offload_sink_wait:
|
||||
activation_offload_sink_wait_async(fwd_module)
|
||||
if config.activation_reload_prefetch:
|
||||
activation_reload_prefetch_async(bwd_module)
|
||||
else:
|
||||
# Use synchronous device_put (1 node each)
|
||||
offload_chosen_sets(fwd_module, bwd_module)
|
||||
|
||||
fwd_module.graph.lint()
|
||||
bwd_module.graph.lint()
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
"""Custom ops for async activation offloading between GPU and CPU.
|
||||
|
||||
These ops encapsulate stream management internally, producing a clean 2-node
|
||||
IR pattern (offload/reload + wait_tensor) similar to c10d functional collectives.
|
||||
|
||||
A single dedicated transfer stream handles all D2H/H2D copies.
|
||||
Completion events are keyed by output tensor data_ptr() and stored in a
|
||||
module-level registry, so ``ao.wait_tensor`` takes only the tensor itself
|
||||
(plus an optional keepalive).
|
||||
|
||||
Offload pattern:
|
||||
cpu_tensor = ao.offload(gpu_tensor)
|
||||
cpu_tensor = ao.wait_tensor(cpu_tensor, gpu_tensor)
|
||||
(keepalive arg extends gpu_tensor lifetime past the async D2H copy)
|
||||
|
||||
Reload pattern:
|
||||
gpu_tensor = ao.reload(cpu_tensor, device)
|
||||
gpu_tensor = ao.wait_tensor(gpu_tensor)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch._library.custom_ops import custom_op
|
||||
from torch.fx import has_side_effect
|
||||
|
||||
|
||||
# --- Global transfer stream (one per device, lazily created) ---
|
||||
_transfer_streams: dict[torch.device, torch.Stream] = {}
|
||||
|
||||
|
||||
def _get_or_create_transfer_stream(device: torch.device) -> torch.Stream:
|
||||
if device not in _transfer_streams:
|
||||
_transfer_streams[device] = torch.Stream(device=device)
|
||||
return _transfer_streams[device]
|
||||
|
||||
|
||||
# --- Wait registry: maps data_ptr() -> (completion_event, device) ---
|
||||
# Created by ao.offload / ao.reload, consumed (popped) by ao.wait_tensor.
|
||||
# Not thread-safe — graph execution is single-threaded Python.
|
||||
_wait_registry: dict[int, tuple[torch.Event, torch.device]] = {}
|
||||
|
||||
|
||||
def _register_wait(tensor: torch.Tensor, device: torch.device) -> torch.Event:
|
||||
"""Create an event for an async transfer and register it for wait_tensor."""
|
||||
event = torch.Event()
|
||||
_wait_registry[tensor.data_ptr()] = (event, device)
|
||||
return event
|
||||
|
||||
|
||||
def _pop_wait(tensor: torch.Tensor) -> tuple[torch.Event, torch.device]:
|
||||
key = tensor.data_ptr()
|
||||
try:
|
||||
return _wait_registry.pop(key)
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"ao.wait_tensor: no pending transfer for tensor with data_ptr={key}. "
|
||||
"Every ao.wait_tensor must be paired with a preceding ao.offload or ao.reload."
|
||||
) from None
|
||||
|
||||
|
||||
def _clear_wait_registry() -> None:
|
||||
_wait_registry.clear()
|
||||
|
||||
|
||||
@custom_op("ao::offload", mutates_args=())
|
||||
def offload(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Async offload a GPU tensor to CPU on the dedicated transfer stream.
|
||||
|
||||
Callers MUST pair this with an ``ao.wait_tensor`` that passes the source GPU
|
||||
tensor as ``keepalive`` to extend its lifetime past the async D2H copy.
|
||||
Do NOT use ``record_stream`` — it causes memory fragmentation and
|
||||
unbounded memory growth.
|
||||
|
||||
Uses pinned-memory allocation + copy_ so the transfer is compatible
|
||||
with CUDA graph capture.
|
||||
"""
|
||||
device = tensor.device
|
||||
transfer_stream = _get_or_create_transfer_stream(device)
|
||||
current_stream = torch.accelerator.current_stream(device)
|
||||
|
||||
transfer_stream.wait_stream(current_stream)
|
||||
|
||||
torch.accelerator.set_stream(transfer_stream)
|
||||
result = torch.empty_like(tensor, device="cpu", pin_memory=True)
|
||||
completion_event = _register_wait(result, device)
|
||||
result.copy_(tensor, non_blocking=True)
|
||||
transfer_stream.record_event(completion_event)
|
||||
torch.accelerator.set_stream(current_stream)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@offload.register_fake
|
||||
def _(tensor: torch.Tensor) -> torch.Tensor:
|
||||
return torch.empty_like(tensor, device="cpu")
|
||||
|
||||
|
||||
@custom_op("ao::reload", mutates_args=())
|
||||
def reload(
|
||||
tensor: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Async reload a CPU tensor to GPU on the dedicated transfer stream.
|
||||
|
||||
The GPU tensor is allocated on the compute stream to avoid cross-stream
|
||||
allocator ownership issues. The H2D copy runs on the transfer stream.
|
||||
The completion event is keyed by the output tensor's data_ptr.
|
||||
"""
|
||||
transfer_stream = _get_or_create_transfer_stream(device)
|
||||
current_stream = torch.accelerator.current_stream(device)
|
||||
|
||||
# Allocate on compute stream so the allocator tracks ownership correctly
|
||||
result = torch.empty_like(tensor, device=device)
|
||||
completion_event = _register_wait(result, device)
|
||||
|
||||
transfer_stream.wait_stream(current_stream)
|
||||
|
||||
torch.accelerator.set_stream(transfer_stream)
|
||||
result.copy_(tensor, non_blocking=True)
|
||||
transfer_stream.record_event(completion_event)
|
||||
torch.accelerator.set_stream(current_stream)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@reload.register_fake
|
||||
def _(
|
||||
tensor: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty_like(tensor, device=device)
|
||||
|
||||
|
||||
# ao::wait_tensor is defined via torch.library with an aliasing schema so the
|
||||
# output can alias the input (custom_op forbids this).
|
||||
#
|
||||
# Uses CompositeExplicitAutograd (single impl for all devices) because the
|
||||
# offload case has mixed-device args: ``tensor`` is CPU (the offload result)
|
||||
# while ``keepalive`` is CUDA (the source GPU tensor). A single impl avoids
|
||||
# relying on device-priority dispatch ordering.
|
||||
#
|
||||
# Synchronization details (completion event, device) are looked up from
|
||||
# ``_wait_registry`` keyed on ``tensor.data_ptr()``.
|
||||
#
|
||||
# ``keepalive`` is not read by the op — its sole purpose is to create a graph
|
||||
# dependency that extends the tensor's lifetime in the FX graph. For offload,
|
||||
# this keeps the source GPU tensor alive until the compute stream has waited
|
||||
# on the D2H completion event, preventing the allocator from reclaiming it
|
||||
# while the async copy is still in flight.
|
||||
_lib = torch.library.Library("ao", "DEF")
|
||||
_lib.define("wait_tensor(Tensor(a) tensor, Tensor? keepalive=None) -> Tensor(a)")
|
||||
|
||||
|
||||
@torch.library.impl("ao::wait_tensor", "CompositeExplicitAutograd")
|
||||
def _ao_wait_tensor(
|
||||
tensor: torch.Tensor,
|
||||
keepalive: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
completion_event, device = _pop_wait(tensor)
|
||||
current_stream = torch.accelerator.current_stream(device)
|
||||
current_stream.wait_event(completion_event)
|
||||
return tensor
|
||||
|
||||
|
||||
@torch.library.register_fake("ao::wait_tensor")
|
||||
def _ao_wait_tensor_fake(
|
||||
tensor: torch.Tensor,
|
||||
keepalive: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return tensor
|
||||
|
||||
|
||||
has_side_effect(torch.ops.ao.wait_tensor.default)
|
||||
|
||||
|
||||
def wait_tensor(
|
||||
tensor: torch.Tensor,
|
||||
keepalive: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Callable wrapper so ``wait_tensor`` can be imported by name for op registration."""
|
||||
return torch.ops.ao.wait_tensor.default(tensor, keepalive)
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
+737
@@ -0,0 +1,737 @@
|
||||
"""
|
||||
This module provides result classes for AOT Autograd compilation.
|
||||
|
||||
Similar to how torch._inductor.output_code provides OutputCode classes for inductor
|
||||
compilation results, this module provides AOTAutogradResult classes that represent
|
||||
the compiled artifacts produced by AOT Autograd.
|
||||
|
||||
These results are:
|
||||
- Serializable: can be saved/loaded from disk without recompilation
|
||||
- Addressable: can be stored in caches with keys for later retrieval
|
||||
- Reusable: can be used for both caching and ahead-of-time compilation (precompile)
|
||||
|
||||
The main result types are:
|
||||
- GenericAOTAutogradResult: Abstract base for all AOT Autograd results
|
||||
- AOTAutogradResult: Regular result that references FxGraphCache entries
|
||||
- BundledAOTAutogradResult: Result that bundles the entire compiled code directly
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable, Sequence
|
||||
from copy import copy
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic, TYPE_CHECKING, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._dynamo.precompile_context import BackendCacheArtifact
|
||||
from torch._inductor.codecache import FxGraphCache
|
||||
from torch._inductor.output_code import (
|
||||
CompiledFxGraph,
|
||||
CompiledFxGraphConstants,
|
||||
OutputCode,
|
||||
)
|
||||
from torch._inductor.utils import should_use_remote_fx_graph_cache
|
||||
from torch._logging import getArtifactLogger
|
||||
|
||||
from .runtime_wrappers import (
|
||||
AOTDispatchAutograd,
|
||||
AOTDispatchAutogradCompileSpec,
|
||||
AOTDispatchSubclassWrapper,
|
||||
CachedAutogradLazyBackwardCompileInfo,
|
||||
CompilerWrapper,
|
||||
FunctionalizedRngRuntimeWrapper,
|
||||
post_compile,
|
||||
RuntimeWrapper,
|
||||
SerializableCompiledFunction,
|
||||
SubclassMeta,
|
||||
)
|
||||
from .schemas import AOTAutogradCacheInfo # noqa: F401
|
||||
from .utils import simple_wraps
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._inductor.compile_fx import _CompileFxKwargs
|
||||
|
||||
from .schemas import AOTConfig, ViewAndMutationMeta
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
aot_graphs_log = getArtifactLogger(__name__, "aot_graphs")
|
||||
|
||||
|
||||
TOut = TypeVar("TOut", bound=OutputCode)
|
||||
|
||||
|
||||
class InductorOutput(ABC, Generic[TOut]):
|
||||
"""
|
||||
Class representing a single inductor output
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def pre_save(self) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
def load(self, example_inputs: Sequence[Any]) -> TOut: ...
|
||||
|
||||
@abstractmethod
|
||||
def post_compile(self, result: TOut, fx_config: _CompileFxKwargs) -> TOut: ...
|
||||
|
||||
|
||||
TOutputCode = TypeVar("TOutputCode", bound=OutputCode)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BundledOutputCodeLoadable(InductorOutput[TOutputCode], Generic[TOutputCode]):
|
||||
"""
|
||||
A generic wrapper for OutputCode objects that are bundled directly in the cache
|
||||
(rather than looked up via FxGraphCache).
|
||||
|
||||
This works for any OutputCode subclass (CompiledFxGraph, RegionalOutputCode, etc.)
|
||||
"""
|
||||
|
||||
result: TOutputCode
|
||||
|
||||
def pre_save(self) -> None:
|
||||
disk_result = copy(self.result)
|
||||
disk_result.prepare_for_serialization()
|
||||
self.result = disk_result
|
||||
return
|
||||
|
||||
def load(self, example_inputs: Sequence[Any]) -> TOutputCode:
|
||||
self.example_inputs = example_inputs
|
||||
return self.result
|
||||
|
||||
def post_compile(
|
||||
self, result: TOutputCode, fx_config: _CompileFxKwargs
|
||||
) -> TOutputCode:
|
||||
constants = CompiledFxGraphConstants()
|
||||
|
||||
# Special handling for CompiledFxGraph - needs FxGraphCache.cache_hit_post_compile
|
||||
if isinstance(result, CompiledFxGraph):
|
||||
graph, cache_info = FxGraphCache.cache_hit_post_compile(
|
||||
result, {}, constants
|
||||
)
|
||||
if graph is None:
|
||||
raise RuntimeError("Failed to reload cache entry from disk")
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "fx_graph_bundled_cache_hit", # always a hit
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: json.dumps(cache_info),
|
||||
)
|
||||
result = graph # type: ignore[assignment]
|
||||
result.compile_region_name = ( # pyrefly: ignore[missing-attribute]
|
||||
fx_config.get("compile_region_name")
|
||||
)
|
||||
|
||||
# Run normal post compile
|
||||
result.post_compile(self.example_inputs, constants, fx_config)
|
||||
|
||||
# Let the CUDAGraph policy do outer-level wrapping (e.g. wrapping
|
||||
# an entire RegionalOutputCode as a single CUDA graph instead of
|
||||
# per-inner-region).
|
||||
import torch._inductor.config as _inductor_config
|
||||
|
||||
policy = _inductor_config.cudagraph_policy
|
||||
if policy is not None:
|
||||
result = policy.wrap_output(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# Backwards compatibility alias
|
||||
CompiledFxGraphLoadable: type[BundledOutputCodeLoadable[CompiledFxGraph]] = (
|
||||
BundledOutputCodeLoadable[CompiledFxGraph]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FxGraphCacheLoadable(InductorOutput[CompiledFxGraph]):
|
||||
fx_graph_cache_info: tuple[str, list[str]]
|
||||
fx_graph_guard_expr: str | None
|
||||
|
||||
def pre_save(self) -> None:
|
||||
return
|
||||
|
||||
def _is_backward(self) -> bool:
|
||||
return False
|
||||
|
||||
def load(self, example_inputs: Sequence[Any]) -> CompiledFxGraph:
|
||||
from .autograd_cache import FXGraphCacheMiss
|
||||
|
||||
# [Note: AOTAutogradCache and FXGraphCache Guard interactions]
|
||||
# As mentioned, AOTAutograd takes in the symint inputs from dynamo's list of arguments.
|
||||
# FXGraphCache serializes guards that are needed in the shape_env based on these symint inputs to the graph.
|
||||
# The invariant that AOTAutograd uses here is that the sources for symints given to it by dynamo are exactly
|
||||
# the same as the ones it passes to inductor, for both the forward and backward passes.
|
||||
# (This does not mean that the tensor values passed in are the same: only that their symints are).
|
||||
# That is, AOTAutograd and Inductor never create new guards based on symints with different sources
|
||||
# than those passed to it by inductor.
|
||||
# We pass the post compile function, which sets various fx_config boxed values,
|
||||
# so we can call it only after we're sure both forward and backward have
|
||||
# Clear CompiledTritonKernels before loading from FXGraphCache
|
||||
torch._inductor.async_compile.CompiledTritonKernels.cache_clear()
|
||||
remote_cache = None
|
||||
constants = CompiledFxGraphConstants()
|
||||
if should_use_remote_fx_graph_cache():
|
||||
remote_cache = FxGraphCache.get_remote_cache()
|
||||
(cache_key, debug_lines) = self.fx_graph_cache_info
|
||||
|
||||
def check_exact_guard_match(guard_expr: str, _hints: Any) -> bool:
|
||||
"""
|
||||
AOTAutogradCache tracks its own guards, so we just need to treat these guard expressions as a second
|
||||
cache key of sorts: we just check for equality, i.e. the FXGraphCache entry with
|
||||
the exact same guards as we originally saved into the cache.
|
||||
"""
|
||||
return guard_expr == self.fx_graph_guard_expr
|
||||
|
||||
result, cache_info = FxGraphCache.load_with_key(
|
||||
cache_key,
|
||||
debug_lines,
|
||||
example_inputs,
|
||||
local=True,
|
||||
remote_cache=remote_cache,
|
||||
is_backward=self._is_backward(),
|
||||
constants=constants,
|
||||
evaluate_guards=check_exact_guard_match,
|
||||
)
|
||||
if result is None:
|
||||
log.info("FXGraphCache cache miss for key %s", self.fx_graph_cache_info)
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "fx_graph_cache_miss", # always a hit
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: json.dumps(cache_info),
|
||||
)
|
||||
|
||||
raise FXGraphCacheMiss
|
||||
|
||||
# No need to log chromium event because AOTAutograd will log that immediately for us
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "fx_graph_cache_hit", # always a hit
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: json.dumps(cache_info),
|
||||
)
|
||||
self.example_inputs = example_inputs
|
||||
self.constants = constants
|
||||
return result
|
||||
|
||||
def post_compile(
|
||||
self, result: CompiledFxGraph, fx_config: _CompileFxKwargs
|
||||
) -> CompiledFxGraph:
|
||||
"""
|
||||
Called after FXGraphCacheLoadable.load, mutates fx_config
|
||||
"""
|
||||
result.compile_region_name = fx_config.get( # pyrefly: ignore[bad-assignment]
|
||||
"compile_region_name"
|
||||
)
|
||||
result.post_compile(self.example_inputs, self.constants, fx_config)
|
||||
|
||||
import torch._inductor.config as _inductor_config
|
||||
|
||||
policy = _inductor_config.cudagraph_policy
|
||||
if policy is not None:
|
||||
result = policy.wrap_output(result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompiledForward(FxGraphCacheLoadable):
|
||||
"""
|
||||
Cacheable entry for a forward function
|
||||
"""
|
||||
|
||||
def _is_backward(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenericCompiledBackward(InductorOutput[TOut]):
|
||||
# Used by AOTDispatchAutograd.post_compile
|
||||
backward_state_indices: list[int]
|
||||
num_symints_saved_for_bw_: int
|
||||
|
||||
def post_compile(self, result: TOut, fx_config: _CompileFxKwargs) -> TOut:
|
||||
# The concrete post_compile comes from the loadable mixin in each subclass MRO.
|
||||
compiled_bw = super().post_compile( # pyrefly: ignore[missing-attribute]
|
||||
result, fx_config
|
||||
)
|
||||
# See note [Wrapping bw_compiler in disable]
|
||||
# This is done by _wrapped_bw_compiler in torch/_dynamo/backends/common.py
|
||||
# But since on cache hit we do not call the bw_compiler, we need to reapply the disable
|
||||
return torch._dynamo.disable( # type: ignore[return-value]
|
||||
compiled_bw, reason="do not trace generated backwards pass"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompiledBackward(GenericCompiledBackward[CompiledFxGraph], FxGraphCacheLoadable):
|
||||
"""
|
||||
Cacheable entry for a backward function
|
||||
"""
|
||||
|
||||
def _is_backward(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
# Generic bundled forward/backward classes that work with any OutputCode type
|
||||
@dataclass
|
||||
class BundledCompiledForward(
|
||||
BundledOutputCodeLoadable[TOutputCode], Generic[TOutputCode]
|
||||
):
|
||||
"""
|
||||
Generic forward function for bundled compilation.
|
||||
Works with any OutputCode type (CompiledFxGraph, RegionalOutputCode, etc.)
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BundledCompiledBackward(
|
||||
GenericCompiledBackward[TOutputCode],
|
||||
BundledOutputCodeLoadable[TOutputCode],
|
||||
Generic[TOutputCode],
|
||||
):
|
||||
"""
|
||||
Generic backward function for bundled compilation.
|
||||
Works with any OutputCode type (CompiledFxGraph, RegionalOutputCode, etc.)
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class SerializedGraphModule:
|
||||
fn: Callable[[dict[Any, Any], str], torch.nn.Module]
|
||||
args: tuple[Any, ...]
|
||||
|
||||
def __init__(self, gm: torch.fx.GraphModule) -> None:
|
||||
self.fn, self.args = gm.__reduce__()
|
||||
|
||||
def deserialize(self) -> torch.fx.GraphModule:
|
||||
gm = self.fn(*self.args)
|
||||
if not isinstance(gm, torch.fx.GraphModule):
|
||||
raise AssertionError(f"expected fx.GraphModule, got {type(gm)}")
|
||||
return gm
|
||||
|
||||
|
||||
def serialize_graph_module(gm: torch.fx.GraphModule) -> SerializedGraphModule:
|
||||
# NOTE: mutates the graph module
|
||||
gm.meta = {}
|
||||
for node in gm.graph.nodes:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
node.meta = {}
|
||||
return SerializedGraphModule(gm)
|
||||
|
||||
|
||||
TForward = TypeVar("TForward", bound="InductorOutput[Any]")
|
||||
TBackward = TypeVar("TBackward", bound="GenericCompiledBackward[Any]")
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenericAOTAutogradResult(Generic[TForward, TBackward]):
|
||||
"""A single result from AOT Autograd compilation, genericized by Forward and Backward types.
|
||||
|
||||
A TForward is always an InductorOutput of some sort, which represents the
|
||||
forward graph of the compile.
|
||||
A TBackward is an InductorOutput + metadata about the backward, useful for specific
|
||||
backward-only wrappers. This type is encapsulated by GenericCompiledBackward.
|
||||
|
||||
Each AOTAutogradResult is essentially parameterized by 1. the method of loading
|
||||
from the cache (either Bundled or UnBundled), and 2. The type of the output. For now,
|
||||
the only type of output we support is Python Wrapper output, i.e. OutputCode.CompiledFxGraph,
|
||||
but the same technique works for C++ wrapper code; we'd just add an extra InductorOutput type.
|
||||
"""
|
||||
|
||||
# Forward and Backward info
|
||||
compiled_fw: TForward
|
||||
compiled_bw: TBackward | None
|
||||
|
||||
# Code of the joint graph using print_readable()
|
||||
# Used for logging purposes
|
||||
aot_joint_graph_str: str | None
|
||||
aot_forward_graph_str: str | None
|
||||
aot_backward_graph_str: str | None
|
||||
|
||||
# Runtime_metadata saved right before compilation
|
||||
runtime_metadata: ViewAndMutationMeta
|
||||
|
||||
# Wrappers that run after each aot_dispatch_* function
|
||||
dispatch_wrappers: list[CompilerWrapper]
|
||||
|
||||
# Used by AOTSubclassWrapper
|
||||
maybe_subclass_meta: SubclassMeta | None
|
||||
num_fw_outs_saved_for_bw: int | None
|
||||
|
||||
# Used by RuntimeWrapper
|
||||
indices_of_inps_to_detach: list[int]
|
||||
|
||||
# Time taken to trace/compile the forward
|
||||
# forward_time_taken includes AOTAutograd tracing time + inductor compilation time
|
||||
# backward_time_taken is essentially just the time inductor took to compile
|
||||
forward_time_taken_ns: int
|
||||
backward_time_taken_ns: int
|
||||
|
||||
# Used by standalone_compile
|
||||
sanitized_aot_config: AOTConfig
|
||||
|
||||
guards_expr: str | None
|
||||
|
||||
# Used by Compiled Autograd
|
||||
serialized_bw_module: SerializedGraphModule | None
|
||||
|
||||
def pre_save(self) -> None:
|
||||
"""
|
||||
Perform any preparations to make the result ready for serialization.
|
||||
"""
|
||||
self.compiled_fw.pre_save()
|
||||
if self.compiled_bw is not None:
|
||||
self.compiled_bw.pre_save()
|
||||
|
||||
def _log_cached_graphs(self, aot_config: AOTConfig) -> None:
|
||||
if not aot_config.enable_log:
|
||||
return
|
||||
|
||||
if self.aot_joint_graph_str is not None:
|
||||
torch._logging.trace_structured(
|
||||
"aot_joint_graph", payload_fn=lambda: self.aot_joint_graph_str
|
||||
)
|
||||
aot_graphs_log.info(
|
||||
"Joint graph (from cache)\n\n%s", self.aot_joint_graph_str
|
||||
)
|
||||
|
||||
if self.aot_forward_graph_str is not None:
|
||||
from torchgen.utils import dataclass_repr
|
||||
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "aot_forward_graph_fw_metadata",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: dataclass_repr(self.runtime_metadata),
|
||||
)
|
||||
if self.maybe_subclass_meta is not None:
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "aot_forward_graph_fw_subclass_metadata",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: dataclass_repr(self.maybe_subclass_meta),
|
||||
)
|
||||
|
||||
# It's called an inference graph if not running with autograd
|
||||
has_backward = self.aot_backward_graph_str is not None
|
||||
torch._logging.trace_structured(
|
||||
"aot_forward_graph" if has_backward else "aot_inference_graph",
|
||||
payload_fn=lambda: self.aot_forward_graph_str,
|
||||
)
|
||||
aot_graphs_log.info(
|
||||
"Forward graph (from cache)\n\n%s",
|
||||
self.aot_forward_graph_str,
|
||||
)
|
||||
|
||||
if self.aot_backward_graph_str is not None:
|
||||
torch._logging.trace_structured(
|
||||
"aot_backward_graph", payload_fn=lambda: self.aot_backward_graph_str
|
||||
)
|
||||
aot_graphs_log.info(
|
||||
"Backward graph (from cache)\n\n%s",
|
||||
self.aot_backward_graph_str,
|
||||
)
|
||||
|
||||
def _load_and_post_compile(
|
||||
self,
|
||||
args: list[torch.Tensor],
|
||||
fx_config: _CompileFxKwargs,
|
||||
) -> tuple[Callable[..., Any], Callable[..., Any] | None, bool]:
|
||||
from torch._dynamo.utils import CompileEventLogger
|
||||
|
||||
compiled_fw_func = self.compiled_fw.load(args)
|
||||
if self.compiled_bw is not None:
|
||||
compiled_bw_func = self.compiled_bw.load(args)
|
||||
needs_autograd = True
|
||||
CompileEventLogger.try_add_pt2_compile(
|
||||
"backend_compile", dispatch_mode="autograd"
|
||||
)
|
||||
# Now that we've loaded forward and backward, call post compile on both
|
||||
# This avoids setting things like BoxedBools in fx_config until
|
||||
# after both forward and backward cache hit
|
||||
fw_fx_config: _CompileFxKwargs = {
|
||||
**fx_config,
|
||||
"is_backward": False,
|
||||
}
|
||||
bw_fx_config: _CompileFxKwargs = {
|
||||
**fx_config,
|
||||
"is_backward": True,
|
||||
}
|
||||
compiled_fw_func = self.compiled_fw.post_compile(
|
||||
compiled_fw_func, fw_fx_config
|
||||
)
|
||||
compiled_bw_func = self.compiled_bw.post_compile(
|
||||
compiled_bw_func, bw_fx_config
|
||||
)
|
||||
return compiled_fw_func, compiled_bw_func, needs_autograd
|
||||
|
||||
inference_fx_config: _CompileFxKwargs = {
|
||||
**fx_config,
|
||||
"is_backward": False,
|
||||
}
|
||||
|
||||
needs_autograd = False
|
||||
CompileEventLogger.try_add_pt2_compile(
|
||||
"backend_compile", dispatch_mode="inference"
|
||||
)
|
||||
compiled_fw_func = self.compiled_fw.post_compile(
|
||||
compiled_fw_func, inference_fx_config
|
||||
)
|
||||
return compiled_fw_func, None, needs_autograd
|
||||
|
||||
def _apply_runtime_wrappers(
|
||||
self,
|
||||
compiled_fw_func: Callable[..., Any],
|
||||
compiled_bw_func: Callable[..., Any] | None,
|
||||
needs_autograd: bool,
|
||||
aot_config: AOTConfig,
|
||||
) -> Callable[..., Any]:
|
||||
from torch._dynamo.utils import CompileEventLogger
|
||||
|
||||
compiled_fw_func = AOTDispatchSubclassWrapper(
|
||||
trace_joint=needs_autograd,
|
||||
fw_only=None,
|
||||
maybe_subclass_meta=self.maybe_subclass_meta,
|
||||
num_fw_outs_saved_for_bw=self.num_fw_outs_saved_for_bw,
|
||||
).post_compile(
|
||||
compiled_fw_func, aot_config, runtime_metadata=self.runtime_metadata
|
||||
)
|
||||
|
||||
req_subclass_dispatch = self.maybe_subclass_meta is not None
|
||||
CompileEventLogger.try_add_pt2_compile(
|
||||
"backend_compile", requires_subclass_dispatch=req_subclass_dispatch
|
||||
)
|
||||
|
||||
# In autograd case, functionalizedRngWrapper should not modify outs
|
||||
return_new_outs = not needs_autograd
|
||||
compiled_fw_func = FunctionalizedRngRuntimeWrapper(
|
||||
return_new_outs=return_new_outs
|
||||
).post_compile(
|
||||
compiled_fw_func, aot_config, runtime_metadata=self.runtime_metadata
|
||||
)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
compiled_fw_func._boxed_call = True
|
||||
disable_amp = torch._C._is_any_autocast_enabled()
|
||||
|
||||
if needs_autograd:
|
||||
if self.compiled_bw is None:
|
||||
raise AssertionError("compiled_bw must not be None when needs_autograd")
|
||||
|
||||
cached_lazy_backward = None
|
||||
if self.serialized_bw_module is not None:
|
||||
cached_lazy_backward = CachedAutogradLazyBackwardCompileInfo(
|
||||
self.serialized_bw_module.deserialize
|
||||
)
|
||||
# This function is run on both cache miss and cache hit, either here
|
||||
# or in aot_dispatch_autograd. On a cache hit,
|
||||
# 1. the bw is already compiled
|
||||
# 2. we don't need to save to the cache again
|
||||
# so those corresponding arguments are set to None.
|
||||
compile_spec = AOTDispatchAutogradCompileSpec(
|
||||
compiled_fw_func=compiled_fw_func,
|
||||
compiled_bw_func=compiled_bw_func,
|
||||
maybe_subclass_meta=self.maybe_subclass_meta,
|
||||
num_symints_saved_for_bw=self.compiled_bw.num_symints_saved_for_bw_,
|
||||
backward_state_indices=self.compiled_bw.backward_state_indices,
|
||||
disable_amp=disable_amp,
|
||||
indices_of_inps_to_detach=self.indices_of_inps_to_detach,
|
||||
lazy_backward_info=cached_lazy_backward,
|
||||
aot_config=aot_config,
|
||||
fw_metadata=self.runtime_metadata,
|
||||
try_save_cache_entry=None,
|
||||
)
|
||||
compiled_function = AOTDispatchAutograd.post_compile(compile_spec)
|
||||
|
||||
else:
|
||||
compiled_function = RuntimeWrapper(
|
||||
indices_of_inps_to_detach=self.indices_of_inps_to_detach,
|
||||
trace_joint=False,
|
||||
disable_amp=disable_amp,
|
||||
).post_compile(
|
||||
compiled_fw_func, aot_config, runtime_metadata=self.runtime_metadata
|
||||
)
|
||||
|
||||
# Add serialization function back onto object
|
||||
compiled_function, _ = post_compile(
|
||||
self.dispatch_wrappers,
|
||||
compiled_function,
|
||||
aot_config,
|
||||
runtime_metadata=self.runtime_metadata,
|
||||
)
|
||||
return compiled_function
|
||||
|
||||
def _check_guards(self, args: list[torch.Tensor]) -> None:
|
||||
if self.guards_expr:
|
||||
from .autograd_cache import AOTAutogradCache
|
||||
|
||||
symints = AOTAutogradCache._filter_backed_symints(args)
|
||||
check = bool(AOTAutogradCache.evaluate_guards(self.guards_expr, symints))
|
||||
if check is not True:
|
||||
raise AssertionError(f"guards check failed: {check}")
|
||||
|
||||
# Turn result into the original callable
|
||||
def wrap_post_compile(
|
||||
self,
|
||||
args: list[torch.Tensor],
|
||||
aot_config: AOTConfig,
|
||||
fx_config: _CompileFxKwargs,
|
||||
) -> Callable[..., Any]:
|
||||
"""
|
||||
This function takes a result and carefully reconstructs the original callable
|
||||
that AOTAutograd returned the first time it was run. It does this by running the various
|
||||
post compile steps that AOTAutograd runs on its compiled artifact after running the fw/bw compilers.
|
||||
|
||||
In the inference path, this consists of the Subclass, FunctionalzedRngRuntime, and RuntimeWrappers.
|
||||
In the autograd path, this consists of AOTAutogradDispatch.post_compile.
|
||||
|
||||
The steps here should match exactly the steps that are run in aot_dispatch_base and aot_dispatch_autograd.
|
||||
|
||||
Notably absent from the cached path are:
|
||||
- DebugAssertWrapper
|
||||
- FakifiedOutWrapper
|
||||
|
||||
Which we'll handle separately later on, if necessary.
|
||||
"""
|
||||
from torch._dynamo.utils import dynamo_timed
|
||||
|
||||
self._log_cached_graphs(aot_config)
|
||||
with dynamo_timed("AOTAutogradCache.inductor_load"):
|
||||
compiled_fw_func, compiled_bw_func, needs_autograd = (
|
||||
self._load_and_post_compile(args, fx_config)
|
||||
)
|
||||
|
||||
compiled_function = self._apply_runtime_wrappers(
|
||||
compiled_fw_func, compiled_bw_func, needs_autograd, aot_config
|
||||
)
|
||||
# Now that we're pretty sure it's a successful load, add guards
|
||||
# to the existing shape environment from the cache.
|
||||
self._check_guards(args)
|
||||
return compiled_function
|
||||
|
||||
|
||||
class AOTAutogradResult(GenericAOTAutogradResult[CompiledForward, CompiledBackward]):
|
||||
"""
|
||||
Regular AOTAutogradResult: saves the forward/backward FxGraphCache keys
|
||||
and looks them up in FxGraphCache on load
|
||||
"""
|
||||
|
||||
|
||||
class BundledAOTAutogradResult(
|
||||
GenericAOTAutogradResult[
|
||||
BundledCompiledForward[TOutputCode], BundledCompiledBackward[TOutputCode]
|
||||
],
|
||||
Generic[TOutputCode],
|
||||
):
|
||||
"""
|
||||
Generic AOTAutogradResult where we bundle the entire OutputCode directly
|
||||
(rather than looking it up via FxGraphCache).
|
||||
|
||||
This works with any OutputCode type:
|
||||
- CompiledFxGraph: Traditional inductor compilation
|
||||
- RegionalOutputCode: Regional inductor compilation with GraphPickler serialization
|
||||
- Any future OutputCode subclasses
|
||||
|
||||
Type parameter:
|
||||
TOutputCode: The OutputCode subclass (e.g., CompiledFxGraph, RegionalOutputCode)
|
||||
|
||||
Usage with CompiledFxGraph:
|
||||
entry = BundledAOTAutogradResult[CompiledFxGraph](
|
||||
compiled_fw=BundledCompiledForward(result=CompiledFxGraph(...)),
|
||||
compiled_bw=BundledCompiledBackward(
|
||||
result=CompiledFxGraph(...),
|
||||
backward_state_indices=[...],
|
||||
num_symints_saved_for_bw_=...,
|
||||
),
|
||||
...
|
||||
)
|
||||
|
||||
Usage with RegionalOutputCode:
|
||||
entry = BundledAOTAutogradResult[RegionalOutputCode](
|
||||
compiled_fw=BundledCompiledForward(result=RegionalOutputCode(gm)),
|
||||
compiled_bw=BundledCompiledBackward(
|
||||
result=RegionalOutputCode(gm),
|
||||
backward_state_indices=[...],
|
||||
num_symints_saved_for_bw_=...,
|
||||
),
|
||||
...
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def deserialize_bundled_cache_entry(
|
||||
entry: BundledAOTAutogradResult[Any],
|
||||
) -> Callable[..., Any]:
|
||||
from copy import deepcopy
|
||||
|
||||
from torch._inductor.cudagraph_utils import BoxedDeviceIndex
|
||||
from torch._inductor.utils import BoxedBool
|
||||
|
||||
# In the precompile use case, guards are already serialized
|
||||
# by dynamo, so we don't need to add them to the environment
|
||||
entry.guards_expr = None
|
||||
# TODO: this isn't exactly right, because cudagraphs needs to be a shared config
|
||||
# which is set by compile_fx. But in precompile, we never actually call compile_fx
|
||||
# so we don't have a place to track cudagraphs here.
|
||||
cudagraphs = BoxedBool(torch._inductor.config.triton.cudagraphs)
|
||||
boxed_forward_device_index = BoxedDeviceIndex(None)
|
||||
# We need to make a clean copy of the cache entry
|
||||
# in case it needs to be serialized again
|
||||
serializable_copy = deepcopy(entry)
|
||||
|
||||
from torch._subclasses import FakeTensorMode
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
|
||||
context = torch._guards.TracingContext.try_get()
|
||||
if context is None:
|
||||
# Create a clean environment when running fx graph post compile
|
||||
# if one is not available
|
||||
context = torch._guards.TracingContext(FakeTensorMode(shape_env=ShapeEnv()))
|
||||
with torch._guards.tracing(context):
|
||||
compiled_fn = entry.wrap_post_compile(
|
||||
[],
|
||||
entry.sanitized_aot_config,
|
||||
{
|
||||
"cudagraphs": cudagraphs,
|
||||
"boxed_forward_device_index": boxed_forward_device_index,
|
||||
},
|
||||
)
|
||||
# Ensure the deserialized cache entry is still serializable
|
||||
|
||||
compiled_fn = SerializableCompiledFunction(compiled_fn, lambda: serializable_copy)
|
||||
|
||||
# TODO: this ignores flat_params, which can exist
|
||||
# if inline_builtin_nn_modules=False
|
||||
@simple_wraps(compiled_fn)
|
||||
def forward(*runtime_args: Any) -> Any:
|
||||
return compiled_fn(list(runtime_args))
|
||||
|
||||
if not hasattr(compiled_fn, "serialize"):
|
||||
raise AssertionError("compiled_fn must have serialize attribute")
|
||||
forward.serialize = compiled_fn.serialize # type: ignore[attr-defined]
|
||||
|
||||
return forward
|
||||
|
||||
|
||||
@dataclass
|
||||
# pyrefly: ignore [implicit-any]
|
||||
class BundledAOTAutogradCacheArtifact(BackendCacheArtifact[Callable]):
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def after_deserialization(self) -> Callable:
|
||||
return deserialize_bundled_cache_entry(self.content)
|
||||
+1395
File diff suppressed because it is too large
Load Diff
+883
@@ -0,0 +1,883 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
"""
|
||||
This module is one of the analysis modules - it takes as input a function or graph
|
||||
and some preexisting properties, and returns some data that is useful for deciding
|
||||
how to further proceed with compilation or construct runtime wrappers.
|
||||
|
||||
In particular, the analysis here constructs view and mutation metadata from running
|
||||
a functionalized version of the graph under compilation.
|
||||
"""
|
||||
|
||||
import collections
|
||||
import contextlib
|
||||
import logging
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch import Tensor
|
||||
from torch._guards import detect_fake_mode
|
||||
from torch._library.opaque_object import is_opaque_type
|
||||
from torch._logging import getArtifactLogger
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
|
||||
from torch._subclasses.meta_utils import safe_is_leaf
|
||||
from torch.fx.experimental.proxy_tensor import disable_autocast_cache
|
||||
from torch.fx.experimental.symbolic_shapes import is_concrete_int
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils._python_dispatch import (
|
||||
is_traceable_wrapper_subclass,
|
||||
transform_subclass,
|
||||
)
|
||||
|
||||
from .descriptors import (
|
||||
AOTInput,
|
||||
AOTOutput,
|
||||
InputMutationAOTOutput,
|
||||
IntermediateBaseAOTOutput,
|
||||
PlainAOTOutput,
|
||||
TangentAOTInput,
|
||||
)
|
||||
from .functional_utils import (
|
||||
are_all_mutations_hidden_from_autograd,
|
||||
are_all_mutations_under_no_grad_or_inference_mode,
|
||||
from_fun,
|
||||
has_data_mutation,
|
||||
has_metadata_mutation,
|
||||
MetadataKey,
|
||||
to_fun,
|
||||
ViewMetaSequence,
|
||||
was_inductor_storage_resized,
|
||||
)
|
||||
from .schemas import (
|
||||
InputAliasInfo,
|
||||
MemoryFormatMeta,
|
||||
MutationType,
|
||||
OutputAliasInfo,
|
||||
OutputType,
|
||||
ViewAndMutationMeta,
|
||||
)
|
||||
from .subclass_utils import create_subclass_meta
|
||||
from .utils import _get_autocast_states, KNOWN_TYPES, simple_wraps, strict_zip
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
zip = strict_zip
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
static_input_logger = getArtifactLogger("torch._dynamo", "cudagraph_static_inputs")
|
||||
|
||||
|
||||
# Note [Tangents memory format]
|
||||
# We assume tangents memory format to be similar to corresponding output's memory_format.
|
||||
# The idea is that we are technically making a guess about the strides of our tangents,
|
||||
# while we trace out the joint.
|
||||
# If runtime specified tangents will not have the same memory format as predicted traced tangents,
|
||||
# we coerce them at runtime to traced tangents memory format.
|
||||
|
||||
|
||||
# Coercing and collecting traced tangents memory format in one recursive traversal
|
||||
def coerce_tangent_and_suggest_memory_format(
|
||||
x: Tensor,
|
||||
) -> tuple[Any, MemoryFormatMeta | list[Any] | None, bool]:
|
||||
updated = False
|
||||
if not isinstance(x, Tensor):
|
||||
return x, None, updated
|
||||
|
||||
out = x.detach()
|
||||
|
||||
is_subclass = is_traceable_wrapper_subclass(out)
|
||||
|
||||
memory_format = MemoryFormatMeta.from_tensor(out)
|
||||
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
if memory_format.memory_format is not None:
|
||||
was = out
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
out = out.contiguous(memory_format=memory_format.memory_format)
|
||||
updated = was is not out
|
||||
|
||||
# For subclass we keep memory format of outer strides at the beginning of the list
|
||||
out_memory_format = [memory_format] if is_subclass else memory_format
|
||||
|
||||
# Note [Tangents memory format, Part 2]
|
||||
# In the same way that "what strides do we assigns to our tangents" is a question
|
||||
# that we can not answer (and therefore have to guess) as we trace the backward ahead-of-time,
|
||||
# The same applies to any tensor subclass metadata, when we have tangents that are subclasses.
|
||||
# To handle this situation, we have two new methods that a tensor subclass can implement:
|
||||
# (1) __coerce_tangent_metadata__(self)
|
||||
# Given a subclass with "non-standard" metadata, turn it into a new subclass with "normal" metadata.
|
||||
# The main example here is a DTensor with the "_Partial" placement.
|
||||
# If we have a forward output with a _Partial placement, and corresponding tangent
|
||||
# with a Replicate/Shard placement, we have no way to convert the tangent "back" to a _Partial placement.
|
||||
# This method lets us avoid the problem entirely by allowing subclasses to ensure that we can never
|
||||
# have a tangent with "problematic" metadata, that we cannot convert to.
|
||||
# (1) __coerce_same_metadata_as_tangent__(self, metadata)
|
||||
# Given a subclass, and a target differing metadata,
|
||||
# convert self to have the same metadata as the target.
|
||||
# With DTensor being the main example, we can use this to convert a DTensor with a Replicate()
|
||||
# placement into one with a Shard() placement, in the case that we "guessed wrong",
|
||||
# and traced tangents with a Shard() placement at compile time.
|
||||
#
|
||||
if is_subclass and hasattr(out, "__coerce_tangent_metadata__"):
|
||||
out = out.__coerce_tangent_metadata__() # type: ignore[attr-defined]
|
||||
|
||||
if is_subclass:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
attrs = out.__tensor_flatten__()[0]
|
||||
|
||||
for attr in attrs:
|
||||
elem = getattr(out, attr)
|
||||
(
|
||||
new_elem,
|
||||
new_elem_memory_format,
|
||||
elem_updated,
|
||||
) = coerce_tangent_and_suggest_memory_format(elem)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
out_memory_format.append(new_elem_memory_format)
|
||||
if elem_updated:
|
||||
setattr(out, attr, new_elem)
|
||||
|
||||
return out, out_memory_format, updated
|
||||
|
||||
|
||||
# This is a version of functionalization that is specifically designed
|
||||
# for the AOTAutograd use case.
|
||||
#
|
||||
# Unlike functorch's variant, this doesn't use the functorch level system,
|
||||
# instead it directly uses PyTorch's conventional dispatcher to hit the
|
||||
# functionalization key. In particular, this means that FunctionalTensorWrapper
|
||||
# can have autograd data stored directly on it.
|
||||
#
|
||||
# In typical AOTAutograd usage, the dispatch key order will look like:
|
||||
#
|
||||
# Autograd - Functionalization ~~~~> Proxy Mode - Fake Tensor
|
||||
# outer tensor inner tensor
|
||||
#
|
||||
# Returns:
|
||||
# - ViewAndMutationMeta, telling us metadata about the inputs and outputs, and
|
||||
# The list of outputs from the forward, but **only** the outputs that we need
|
||||
# to pass in as tangents into the backward.
|
||||
# Specifically, aliased outputs from the forward get regenerated, and don't participate
|
||||
# in the compiled backward function.
|
||||
def run_functionalized_fw_and_collect_metadata(
|
||||
f: Callable[..., Any],
|
||||
*,
|
||||
flat_args_descs: list[AOTInput],
|
||||
keep_input_mutations: bool,
|
||||
# Note: this is guaranteed to be set when running under dynamo
|
||||
static_input_indices: list[int] | None = None,
|
||||
pre_dispatch: bool = False,
|
||||
) -> Callable[..., ViewAndMutationMeta]:
|
||||
memo: dict[Tensor, Tensor] = {}
|
||||
|
||||
# TODO: see if we can rewrite this to be more accurate using
|
||||
# overload
|
||||
def _to_fun(t: object) -> object:
|
||||
if isinstance(t, Tensor):
|
||||
if t in memo:
|
||||
return memo[t]
|
||||
r = to_fun(t)
|
||||
memo[t] = r
|
||||
return r
|
||||
else:
|
||||
return t
|
||||
|
||||
@simple_wraps(f)
|
||||
def inner(*flat_args: Any) -> ViewAndMutationMeta:
|
||||
# This function is meant to be run with the forward, which expects a flat list of tensor/symint/other args.
|
||||
if not all(
|
||||
isinstance(a, tuple(KNOWN_TYPES)) or is_opaque_type(type(a))
|
||||
for a in flat_args
|
||||
):
|
||||
raise AssertionError("all flat_args must be KNOWN_TYPES or opaque types")
|
||||
|
||||
input_info: list[InputAliasInfo] = []
|
||||
output_info: list[OutputAliasInfo] = []
|
||||
|
||||
prior_grad_enabled = torch.is_grad_enabled()
|
||||
prior_autocast_states = _get_autocast_states()
|
||||
|
||||
# See Note [Disabling Functionalize TLS Above Python Functionalization]
|
||||
disable_above = torch._C._ExcludeDispatchKeyGuard(
|
||||
torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
||||
)
|
||||
|
||||
# It doesn't matter if we run this under predispatch or not because it is
|
||||
# only for figuring out metadata
|
||||
mode = FunctionalTensorMode(_allow_token_discovery=True)
|
||||
suppress_pending = contextlib.nullcontext()
|
||||
fake_mode = detect_fake_mode()
|
||||
if fake_mode and (shape_env := fake_mode.shape_env):
|
||||
suppress_pending = shape_env.ignore_fresh_unbacked_symbols()
|
||||
with disable_above, mode, suppress_pending, disable_autocast_cache():
|
||||
# precondition: The passed in function already handles unflattening inputs + flattening outputs
|
||||
flat_f_args = pytree.tree_map(_to_fun, flat_args)
|
||||
flat_f_args_descs = flat_args_descs
|
||||
flat_f_outs = f(*flat_f_args)
|
||||
|
||||
# Assert that f does NOT have an AOTOutputs in it, easy mistake to
|
||||
# make! You need to drop the second output before calling this
|
||||
# function
|
||||
if pytree.tree_any(lambda x: isinstance(x, AOTOutput), flat_f_outs):
|
||||
raise AssertionError(
|
||||
f"{f} returned AOTOutput when it shouldn't. Did you remember to wrap the "
|
||||
"function with without_output_descs before passing it here?"
|
||||
)
|
||||
|
||||
# NB: this is just to setup the input descriptors, we will
|
||||
# recreate these descriptors (with the same convention!) when we
|
||||
# actually do the trace
|
||||
flat_f_outs_descs = [PlainAOTOutput(i) for i in range(len(flat_f_outs))]
|
||||
|
||||
# We didn't do any tracing, so we don't need to process the
|
||||
# unbacked symbols, they will just disappear into the ether.
|
||||
# Also, prevent memoization from applying.
|
||||
if fake_mode:
|
||||
fake_mode.epoch += 1
|
||||
fake_mode.reset_nt_tensor_id_counter()
|
||||
|
||||
if prior_autocast_states != _get_autocast_states():
|
||||
raise RuntimeError(
|
||||
"AOTAutograd does not support tracing graphs that mutate the autocast state. "
|
||||
"Dynamo will only insert autocast context managers (e.g. with torch.autocast(..)) into the graph, "
|
||||
"which will unwind all of their mutations to autocast state before the graph exits. "
|
||||
"If you encounter this error while using torch.compile, please file a bug."
|
||||
)
|
||||
|
||||
# Inspect the state of the input tensor functional wrapper to detect input mutation info
|
||||
# If inp[i] has a metadata-only mutation, then maybe_inputs_with_mutated_metadata[i] contains the updated version
|
||||
for arg, f_arg in zip(flat_args, flat_f_args):
|
||||
# NB: Mutation of non-contiguous tensor subclass input can result in a mismatch in
|
||||
# strides between the functionalized arg inner tensors and non-functionalized arg inner
|
||||
# tensors. This is a problem as the inner tensor stride change may not be reflected
|
||||
# correctly in the outer tensor, so disallow this for now.
|
||||
mutates_data = has_data_mutation(f_arg)
|
||||
mutates_metadata = has_metadata_mutation(
|
||||
f_arg, arg, check_only_storage_mutation=False
|
||||
)
|
||||
if mutates_metadata and is_traceable_wrapper_subclass(arg):
|
||||
raise RuntimeError(
|
||||
"Metadata mutations are currently not allowed on tensor subclasses"
|
||||
)
|
||||
mutates_storage_metadata = has_metadata_mutation(
|
||||
f_arg, arg, check_only_storage_mutation=True
|
||||
)
|
||||
mutations_hidden_from_autograd = are_all_mutations_hidden_from_autograd(
|
||||
f_arg
|
||||
)
|
||||
mutations_under_no_grad_or_inference_mode = (
|
||||
mutates_data
|
||||
and are_all_mutations_under_no_grad_or_inference_mode(f_arg)
|
||||
)
|
||||
mutation_inductor_storage_resize = was_inductor_storage_resized(f_arg)
|
||||
|
||||
if mutates_storage_metadata:
|
||||
mutates_data = False
|
||||
|
||||
requires_grad = isinstance(f_arg, torch.Tensor) and f_arg.requires_grad
|
||||
|
||||
input_info.append(
|
||||
InputAliasInfo(
|
||||
is_leaf=isinstance(arg, Tensor) and safe_is_leaf(arg),
|
||||
mutates_data=mutates_data,
|
||||
mutates_metadata=mutates_metadata,
|
||||
mutations_hidden_from_autograd=mutations_hidden_from_autograd,
|
||||
mutates_storage_metadata=mutates_storage_metadata,
|
||||
mutations_under_no_grad_or_inference_mode=mutations_under_no_grad_or_inference_mode,
|
||||
mutation_inductor_storage_resize=mutation_inductor_storage_resize,
|
||||
requires_grad=requires_grad,
|
||||
keep_input_mutations=keep_input_mutations,
|
||||
)
|
||||
)
|
||||
|
||||
# If a function involves creating a tensor, and returning a view of it, such that its _base is the intermediate,
|
||||
# We need to make sure our graph returns the _base as a graph output, and we manually recreate the view
|
||||
# to return to the user. Why? The backend compiler is free to (incorrectly) not set requires_grad
|
||||
# on the base tensor, but we are obligated to properly set requires-gradness on the real output.
|
||||
|
||||
inp_storage_refs = {
|
||||
StorageWeakRef(inpt.untyped_storage()): idx
|
||||
for idx, inpt in enumerate(flat_f_args)
|
||||
if isinstance(inpt, Tensor)
|
||||
}
|
||||
|
||||
# We need inp tensor id's to be able to tell if an outputs **are** inputs.
|
||||
inp_tensor_ids = {id(inpt) for inpt in flat_f_args if isinstance(inpt, Tensor)}
|
||||
# We need output tensor id's to tell if any output._base` attributes **are** other outputs.
|
||||
# (This is also a dict because we need to know that output's index, so we can regenerate
|
||||
# the alias from it).
|
||||
out_tensor_ids = {id(o): i for i, o in enumerate(flat_f_outs)}
|
||||
|
||||
# Keep track of which outputs alias other outputs
|
||||
out_tensor_alias_counts: collections.defaultdict[StorageWeakRef | None, int] = (
|
||||
collections.defaultdict(int)
|
||||
)
|
||||
# This tells us, for a given group of outputs that alias each other,
|
||||
# whether they e.g. all came from an unbind call
|
||||
num_aliased_tensors_that_are_multi_output_views: collections.defaultdict[
|
||||
StorageWeakRef | None, int
|
||||
] = collections.defaultdict(int)
|
||||
|
||||
out_storage_to_metadata_key_to_tensors: collections.defaultdict[
|
||||
StorageWeakRef | None,
|
||||
collections.defaultdict[MetadataKey, set[torch.Tensor]],
|
||||
] = collections.defaultdict(lambda: collections.defaultdict(set))
|
||||
|
||||
curr_storage = None
|
||||
for o in flat_f_outs:
|
||||
if isinstance(o, torch.Tensor):
|
||||
curr_storage = StorageWeakRef(o.untyped_storage())
|
||||
out_tensor_alias_counts[curr_storage] += 1
|
||||
# Note: [AOTAutograd: differentiable outputs that alias each other from a multi-output view call]
|
||||
# This is an optimization on top of the "alias of intermediates" logic,
|
||||
# which you can read more about under Note [AOT Autograd: outputs aliasing inputs or intermediates!]
|
||||
#
|
||||
# Before describing the optimization: this is important for AOTAutograd to have good
|
||||
# perf around, multi-output views. HOWEVER:
|
||||
# - There is a more generic change to AOTAutograd that we'd like to make, that subsumes this case,
|
||||
# around using pre-dispatch tracing to partition out a graph so we can faithfully replay all
|
||||
# views without having to regenerate them at runtime.
|
||||
# - It's loosely described in this doc (more details will be added soon):
|
||||
# https://docs.google.com/document/d/1DlfFq8TKbuAn2zyJxLfoW-X1qkkm5PLdHFtySo03QAk/edit
|
||||
# - Once that change lands, we should just rip out this "optimization", since:
|
||||
# (1) It will be fully unnecessary
|
||||
# (2) Although it is only a few lines of code, it is a bit difficult to reason about
|
||||
# its correctness with the autograd engine in all cases.
|
||||
#
|
||||
#
|
||||
# What is this optimization? Consider the below case:
|
||||
# def f(x):
|
||||
# intermediate = x.mul(2)
|
||||
# # x and intermediate here require grad
|
||||
# o1, o2, ... o10 = intermediate.unbind(-1)
|
||||
# return intermediate, o1, o2, ... o10
|
||||
# Now, the "intermediate base" handling in AOTAutograd implies that we must do the following:
|
||||
# (1) return "intermediate as an extra output of the compiled graph
|
||||
# (2) regenerate each aliased output off of "intermediate", **outside** of the autograd.Function.
|
||||
# The reason AOTAutograd ordinarily does this is for safety: the autograd engine needs to know
|
||||
# that o1 through o10 are all aliased, and if we blindly return o1 through o10 from the autograd.Function,
|
||||
# this information will be hidden.
|
||||
# In particular, mutating one alias might require autograd to update autograd metadata on the other aliases
|
||||
# (like their grad_fn, for example, when the autograd engine needs to do view-replay).
|
||||
#
|
||||
# However, intermediate_base logic can be bad for backward performance (we sometimes generate
|
||||
# as_strided calls during the intermediate base logic, which can have a slow backward formula).
|
||||
# Is it possible to find a set of conditions where it is **safe** to hide the output aliasing from autograd?
|
||||
#
|
||||
# For a set of outputs of the graph that alias each other, o_1...o_k, consider:
|
||||
# (1) They came from the same multi-output view op, e.g. o_1, ..., o_k = intermediate.unbind(0)
|
||||
# (2) If there are any other aliases of o_1 through o_k (in the example above, intermediate),
|
||||
# **at most** 1 can escape from the graph (e.g. there is not some other graph input/output
|
||||
# o_other, that aliases these outputs)
|
||||
# (3) o_1...o_k all require_grad, they all share the same ._base, and their ._base requires grad.
|
||||
# This condition is important because it's what causes slowness in the intermediate_base
|
||||
# codepath of aot_autograd. Ordinarily, o_1...o_k would all get a grad_fn, and
|
||||
# aot_autograd's view-replay might give each output an AsStridedBackward as its grad_fn.
|
||||
# "K" AsStridedBackward calls will be *much* slower than a single UnbindBackward.
|
||||
# In this setup, is it possible to mutate one of the outputs o_i in a way that would affect the autograd meta
|
||||
# of the other aliases?
|
||||
#
|
||||
# Claim: No! Consider a few example (which I'm pretty sure cover all cases of mutation w.r.t. autograd):
|
||||
# (a) What happens if we mutate any of o_1 through o_k directly?
|
||||
# Autograd raises an error:
|
||||
# "RuntimeError: Output 0 of UnbindBackward0 is a view and is being modified inplace. This view is
|
||||
# the output of a function that returns multiple views. Such functions do not allow the output
|
||||
# views to be modified inplace. You should replace the inplace operation by an out-of-place one."
|
||||
# (b) What if we take a view of o_k and mutate it, o_k.view(o_k.shape).mul_(2)?
|
||||
# Autograd raises the same error- the "multi-output-view"ness of an alias propagates to future views.
|
||||
# (c) What if we mutate o_k under no_grad?
|
||||
# Autograd raises the same error
|
||||
# (d) What if we detach and mutate, e.g. o_k.detach().mul_(2)?
|
||||
# Autograd allows this, *but* autograd updates all alias's grad_fn's to be error functions when accessed.
|
||||
# Autograd raises the same error
|
||||
# (e) What if we try to mutate another alias of o_1...o_k, that was **not** created from a multi-output view?
|
||||
# We promised that there is at most **one** such alias, e.g. intermediate in the example above.
|
||||
# You can mutate intermediate, but in eager mode this will change the grad_fn of o_1...o_k
|
||||
# to be error fn's.
|
||||
# Since intermediate was the *only* non-multi-output-alias, there are no other aliases
|
||||
# of `intermediate` around that were produced by the compiled fn and have a valid grad_fn.
|
||||
#
|
||||
# Coming back to this optimization:
|
||||
# Given that it is not possible for mutating one of these aliases to affect the autograd metadata of another alias
|
||||
# without causing an error in eager mode, we will simple hide the aliasing from autograd during torch.compile
|
||||
# if all of the above conditions are met.
|
||||
# This has the slight downside that it's possible to write some "bad" code that autograd will raise an error on
|
||||
# in eager but fail to during torch.compile, but it has the benefit that this code has much better performance.
|
||||
# NOTE: if and when we eventually update AOTAutograd to do the "view graph slicing" defined here:
|
||||
# https://docs.google.com/document/d/1DlfFq8TKbuAn2zyJxLfoW-X1qkkm5PLdHFtySo03QAk/edit,
|
||||
# then this optimization will probably matter less and might be ok to remove.
|
||||
is_cur_tensor_multi_out_view = isinstance(
|
||||
o, FunctionalTensor
|
||||
) and torch._functionalize_is_multi_output_view( # type: ignore[attr-defined]
|
||||
o.elem
|
||||
)
|
||||
if is_cur_tensor_multi_out_view:
|
||||
num_aliased_tensors_that_are_multi_output_views[curr_storage] += 1
|
||||
if o.requires_grad:
|
||||
out_storage_to_metadata_key_to_tensors[curr_storage][
|
||||
MetadataKey.make(o)
|
||||
].add(o)
|
||||
|
||||
# maps the id of an intermediate base to its index in the output of the compiled forward
|
||||
intermediate_base_tensor_id_to_output_idx: dict[int, int] = {}
|
||||
intermediate_bases: list[torch.Tensor] = []
|
||||
intermediate_bases_descs: list[AOTInput] = []
|
||||
# Why Do We Care If Storage Changed?
|
||||
# It's important to understand the implications of storage changes in complex scenarios. Take this example:
|
||||
#
|
||||
# def f(x):
|
||||
# x_storage = x.untyped_storage()
|
||||
# non_leaf_tensor = torch.ones(4, requires_grad=True).clone()
|
||||
#
|
||||
# # Using no_grad() and _unsafe_preserve_version_counter to simulate the .data = operation
|
||||
# with torch.no_grad(), torch.autograd._unsafe_preserve_version_counter(x):
|
||||
# x.set_(non_leaf_tensor.untyped_storage())
|
||||
#
|
||||
# out = x.view(-1)
|
||||
#
|
||||
# # Restoring x to its original storage, again simulating .data = operation
|
||||
# with torch.no_grad(), torch.autograd._unsafe_preserve_version_counter(x):
|
||||
# x.set_(x_storage)
|
||||
#
|
||||
# return out
|
||||
#
|
||||
# In this scenario, 'x' and 'out' have different shapes and are stored at different memory addresses, aka no aliasing.
|
||||
# However, due to how set_() and more specificlaly, set is functionalized, is defined to preserve eager semantics,
|
||||
# the autograd engine mistakenly assumes that 'x' and 'out' are aliased, treating 'x' as 'out._base'.
|
||||
# This misinterpretation leads to an 'alias_of_input' flag, causing an unnecessary as_strided() call to be generated,
|
||||
# which could lead to issues later in the code.
|
||||
for o, desc in zip(flat_f_outs, flat_f_outs_descs):
|
||||
functional_tensor_storage_changed = isinstance(
|
||||
o, FunctionalTensor
|
||||
) and torch._functionalize_was_storage_changed( # type: ignore[attr-defined]
|
||||
o.elem
|
||||
)
|
||||
curr_storage = (
|
||||
None
|
||||
if not isinstance(o, torch.Tensor)
|
||||
else StorageWeakRef(o.untyped_storage())
|
||||
)
|
||||
outs_with_identical_metadata_that_require_grad: list[torch.Tensor] = (
|
||||
[]
|
||||
if not isinstance(o, Tensor)
|
||||
else [
|
||||
curr
|
||||
for curr in out_storage_to_metadata_key_to_tensors[curr_storage][
|
||||
MetadataKey.make(o)
|
||||
]
|
||||
if o is not curr
|
||||
]
|
||||
)
|
||||
|
||||
# See Note [Accessing .grad_fn on FunctionalTensor]
|
||||
# In-place operations on views will trigger a lazy rebase of the autograd graph;
|
||||
# this runs during access to the .grad_fn. The rebase logic will invoke view ops
|
||||
# on FunctionalTensors, so we must enable a FunctionalTensorMode here to ensure
|
||||
# these op calls succeed.
|
||||
grad_fn = None
|
||||
if isinstance(o, Tensor):
|
||||
with FunctionalTensorMode():
|
||||
grad_fn = o.grad_fn
|
||||
|
||||
is_result_of_custom_autograd_fn = False
|
||||
# Need to check for both custom cpp (CppFunction) and python (BackwardCFunction)
|
||||
# autograd fns
|
||||
if type(grad_fn).__name__ == "CppFunction":
|
||||
is_result_of_custom_autograd_fn = True
|
||||
if isinstance(grad_fn, torch.autograd.function.BackwardCFunction):
|
||||
is_result_of_custom_autograd_fn = True
|
||||
|
||||
if not isinstance(o, Tensor):
|
||||
output_type = OutputType.non_alias
|
||||
base_idx = None
|
||||
elif (
|
||||
curr_storage in inp_storage_refs
|
||||
and grad_fn is not None
|
||||
and is_result_of_custom_autograd_fn
|
||||
):
|
||||
output_type = OutputType.custom_function_view
|
||||
base_idx = None
|
||||
elif (
|
||||
curr_storage in inp_storage_refs
|
||||
and not functional_tensor_storage_changed
|
||||
):
|
||||
# pyrefly: ignore [bad-index, index-error]
|
||||
base_idx = inp_storage_refs[curr_storage]
|
||||
is_input_tensor = id(o) in inp_tensor_ids
|
||||
num_aliased_outs = out_tensor_alias_counts[curr_storage]
|
||||
num_multi_output_view_outs = (
|
||||
num_aliased_tensors_that_are_multi_output_views[curr_storage]
|
||||
)
|
||||
num_aliased_outs_that_are_not_multi_output_views = (
|
||||
num_aliased_outs - num_multi_output_view_outs
|
||||
)
|
||||
if (
|
||||
grad_fn is not None
|
||||
and num_aliased_outs_that_are_not_multi_output_views == 0
|
||||
):
|
||||
# See Note: [AOTAutograd: differentiable outputs that alias each other from a multi-output view call]
|
||||
# In particular, given:
|
||||
# def f(x):
|
||||
# return list(x.unbind(0))
|
||||
# The main reason we ordinarily try to regenerate these output aliases outside of the
|
||||
# compiled autograd.Function is because if any of the outputs are later mutated,
|
||||
# autograd needs to perform view-replay to regenerate them.
|
||||
# However, autograd does not allow users to mutate multi-output views
|
||||
# in any way that can change the autograd metadata of other aliases.
|
||||
# So we hide this aliasing from autograd here.
|
||||
log.debug(
|
||||
"Encountered AOTAutograd case: differentiable outputs that \
|
||||
alias each other from a multi-output view call"
|
||||
)
|
||||
output_type = OutputType.non_alias
|
||||
elif is_input_tensor:
|
||||
output_type = OutputType.is_input
|
||||
else:
|
||||
output_type = OutputType.alias_of_input
|
||||
elif functional_tensor_storage_changed and id(o) in inp_tensor_ids:
|
||||
# When there is a set_() on an input, we cannot rely on checking storages
|
||||
# to detect if we are returning an input (since the inputs storage is different)
|
||||
if curr_storage is None:
|
||||
raise AssertionError("curr_storage must not be None")
|
||||
base_idx = inp_storage_refs[curr_storage]
|
||||
output_type = OutputType.is_input
|
||||
|
||||
# We only need to handle the intermediate base case when both
|
||||
# the intermediate base and the output require gradients.
|
||||
# See Note [AOT Autograd: outputs aliasing inputs or intermediates!]
|
||||
elif o._base is not None and o.requires_grad and o._base.requires_grad:
|
||||
num_aliased_outs = out_tensor_alias_counts[curr_storage]
|
||||
num_multi_output_view_outs = (
|
||||
num_aliased_tensors_that_are_multi_output_views[curr_storage]
|
||||
)
|
||||
num_aliased_outs_that_are_not_multi_output_views = (
|
||||
num_aliased_outs - num_multi_output_view_outs
|
||||
)
|
||||
# Note: [AOTAutograd: differentiable outputs that alias each other from a multi-output view call]
|
||||
if (
|
||||
out_tensor_alias_counts[curr_storage] == 1
|
||||
or num_aliased_outs_that_are_not_multi_output_views <= 1
|
||||
):
|
||||
# Note [Intermediate Bases Optimization]
|
||||
# Normally if we have an output that aliases an intermediate,
|
||||
# we need to add the extra "intermediate base" logic further down
|
||||
# to prevent autograd from yelling at us if the user later tries to
|
||||
# mutate that output.
|
||||
# However, the common case here is if we have an output that aliases an intermediate,
|
||||
# but doesn't alias any other outputs.
|
||||
# In that case, autograd shouldn't have to worry about the aliasing at all
|
||||
# (if that output is mutated, there are no other live aliases for autograd to worry about).
|
||||
# The "intermediate bases" can hurt inductor perf by forcing more variables to become outputs.
|
||||
# So as an optimization, we won't do intermediate base handling in this case.
|
||||
# Instead, we'll hide the aliasing from autograd using aten._unsafe_view().
|
||||
if (
|
||||
out_tensor_alias_counts[curr_storage] != 1
|
||||
and num_aliased_outs_that_are_not_multi_output_views <= 1
|
||||
):
|
||||
log.debug(
|
||||
"Encountered AOTAutograd case: differentiable outputs that alias each other \
|
||||
from a multi-output view call"
|
||||
)
|
||||
output_type = OutputType.unsafe_view_alias
|
||||
base_idx = None
|
||||
else:
|
||||
# First, check if o's ._base is an existing output
|
||||
maybe_existing_out_idx = out_tensor_ids.get(id(o._base))
|
||||
if maybe_existing_out_idx is not None:
|
||||
# Special case where the output is an alias of a graph intermediate, but that intermediate
|
||||
# is itself also a user output.
|
||||
output_type = (
|
||||
OutputType.alias_of_intermediate_base_is_user_output
|
||||
)
|
||||
base_idx = maybe_existing_out_idx
|
||||
else:
|
||||
# Next, check if o's ._base is an intermediate base that we already returned
|
||||
maybe_existing_base_output_idx = (
|
||||
intermediate_base_tensor_id_to_output_idx.get(id(o._base))
|
||||
)
|
||||
if maybe_existing_base_output_idx is not None:
|
||||
output_type = OutputType.alias_of_intermediate
|
||||
base_idx = maybe_existing_base_output_idx
|
||||
else:
|
||||
# Otherwise, take o._base and explicitly return it as an output in the compiled graph
|
||||
new_out_idx = len(intermediate_bases)
|
||||
base_idx = new_out_idx
|
||||
# Indicate to the logic later on (when we trace the joint)
|
||||
# that this particular output should get it's ._base appended to the forward graph outputs
|
||||
output_type = (
|
||||
OutputType.alias_of_intermediate_save_as_output
|
||||
)
|
||||
intermediate_base_tensor_id_to_output_idx[id(o._base)] = (
|
||||
new_out_idx
|
||||
)
|
||||
intermediate_bases.append(o._base)
|
||||
# NB: The desc we picked here is guaranteed to be
|
||||
# synchronized with the one in
|
||||
# graph_capture_wrappers.py because we
|
||||
# SPECIFICALLY notated this output as
|
||||
# alias_of_intermediate_save_as_output
|
||||
intermediate_bases_descs.append(
|
||||
TangentAOTInput(IntermediateBaseAOTOutput(desc))
|
||||
)
|
||||
elif (
|
||||
# See https://github.com/pytorch/pytorch/issues/100348 for this case.
|
||||
# This protects against the specific case where a user fn returns (output, output.detach())
|
||||
out_tensor_alias_counts[curr_storage] > 1
|
||||
and len(outs_with_identical_metadata_that_require_grad) > 0
|
||||
and not o.requires_grad
|
||||
):
|
||||
# In theory we could use any of these tensors to regenerate the aliased outputs from,
|
||||
# since they all alias each other and have identical metadata
|
||||
out_alias = outs_with_identical_metadata_that_require_grad[0]
|
||||
existing_out_idx = out_tensor_ids[id(out_alias)]
|
||||
output_type = OutputType.alias_of_intermediate_base_is_user_output
|
||||
base_idx = existing_out_idx
|
||||
else:
|
||||
output_type = OutputType.non_alias
|
||||
base_idx = None
|
||||
|
||||
if isinstance(o, torch.Tensor):
|
||||
dynamic_dims = {
|
||||
i for i, s in enumerate(o.shape) if not is_concrete_int(s)
|
||||
}
|
||||
else:
|
||||
dynamic_dims = None
|
||||
|
||||
# Save the current FunctionalTensor output.
|
||||
#
|
||||
# This will be used at runtime for reconstructing output views from
|
||||
# their respective base tensors.
|
||||
#
|
||||
# The FunctionalTensor will be saved if one of the 2 conditions below
|
||||
# is true:
|
||||
view_meta_sequence = None
|
||||
if (
|
||||
# 1. If the output_type is either of:
|
||||
# (i) alias_of_intermediate;
|
||||
# (ii) alias_of_intermediate_save_as_output; or
|
||||
# (iii) alias_of_intermediate_base_is_user_output.
|
||||
#
|
||||
# No need to worry about in-place view operations here, since
|
||||
# this functionalization step elimitates mutations.
|
||||
#
|
||||
# i.e. we have access to the actual base tensor, before the
|
||||
# in-place operation was applied.
|
||||
output_type
|
||||
in (
|
||||
OutputType.alias_of_intermediate,
|
||||
OutputType.alias_of_intermediate_save_as_output,
|
||||
OutputType.alias_of_intermediate_base_is_user_output,
|
||||
)
|
||||
) or (
|
||||
# 2. If the output_type is alias_of_input, and no in-place view
|
||||
# operationthe was run on the input (base tensor).
|
||||
#
|
||||
# In this case, we need to check for metadata mutation because
|
||||
# the runtime explicitly reconstructs the inputs, before actually
|
||||
# reconstructing the outputs. Due to in-place view operations, the
|
||||
# fully reconstructed input may not be this output base tensor
|
||||
# anymore.
|
||||
output_type == OutputType.alias_of_input
|
||||
and base_idx is not None
|
||||
and not input_info[base_idx].mutates_metadata
|
||||
):
|
||||
if isinstance(o, FunctionalTensor):
|
||||
view_meta_sequence = ViewMetaSequence(o)
|
||||
|
||||
requires_grad = isinstance(o, torch.Tensor) and o.requires_grad
|
||||
out_info = OutputAliasInfo(
|
||||
output_type=output_type,
|
||||
raw_type=type(o),
|
||||
base_idx=base_idx,
|
||||
dynamic_dims=dynamic_dims,
|
||||
requires_grad=requires_grad,
|
||||
# A view created under no_grad() inherits requires_grad from
|
||||
# its base but has no grad_fn and does not participate in
|
||||
# differentiation.
|
||||
requires_grad_for_backward=requires_grad
|
||||
and (o._base is None or grad_fn is not None),
|
||||
view_meta_sequence=view_meta_sequence,
|
||||
)
|
||||
output_info.append(out_info)
|
||||
|
||||
# See Note [AOT Autograd: Views to avoid tangents aliasing inputs]
|
||||
def view_avoid_dupes_with_primals(t: object) -> object:
|
||||
if isinstance(t, Tensor) and is_traceable_wrapper_subclass(t):
|
||||
return transform_subclass(
|
||||
t, lambda _, inner_t: view_avoid_dupes_with_primals(inner_t)
|
||||
)
|
||||
if isinstance(t, Tensor):
|
||||
return t.view(t.shape)
|
||||
return t
|
||||
|
||||
# This analysis function returns *only* the outputs that are meant to be tangents to the backwards.
|
||||
# Anything that aliases (inputs returned in the fw due to metadata mutations, or outputs that alias inputs/intermediates)
|
||||
# are *regenerated* later, and not used directly in the autograd graph
|
||||
def _plain_fake_tensor_like_subclass(x: Any) -> torch.Tensor:
|
||||
# pyrefly: ignore [bad-context-manager]
|
||||
with detect_fake_mode():
|
||||
return torch.empty(
|
||||
x.shape, dtype=x.dtype, device=x.device, layout=x.layout
|
||||
)
|
||||
|
||||
def _is_subclass_mutated_input_tangent_always_subclass(inp: object) -> bool:
|
||||
return (
|
||||
isinstance(inp, torch.nested._internal.nested_tensor.NestedTensor)
|
||||
or torch._functorch.config.disable_guess_zero_tangent_for_mutated_input_subclass
|
||||
)
|
||||
|
||||
f_input_tangents_pairs = [
|
||||
# Note: [AOTAutograd Tangent Subclassness for mutated inputs]
|
||||
# Generally when creating tangents to trace with, we assume that tangents will have
|
||||
# the same subclass-ness as their forward outs
|
||||
# however: for tangents that correspond to input mutations, in practice it is more likely
|
||||
# that these tangents will be plain tensors of zeros at runtime, so we tweak our guess
|
||||
# to assume that these tangents should always be plaint tensors.
|
||||
# Example:
|
||||
# def f(x):
|
||||
# x.mul_(2)
|
||||
# return x + 1
|
||||
# out = f(x)
|
||||
# out.sum().backward()
|
||||
# In the above code, we will have a tangent "x_updated_tangent",
|
||||
# which will be a plain tensor of zeros, *unless* x is used in some compute after executing f
|
||||
#
|
||||
# However, there are exceptions to this logic. If a view is created from mutated input and is used in backward,
|
||||
# The tangent for this subclass input will be a subclass tensor.
|
||||
# Example:
|
||||
# def f(a, b):
|
||||
# a.mul_(2)
|
||||
# b.mul_(3)
|
||||
# return b.view(b.shape), a + b
|
||||
# a_out, b_out = f(..., Subclass)
|
||||
# (a * b).sum().backward()
|
||||
#
|
||||
# We can not deduce it easily now, so introducing a debug config to be able to turn off this for specific cases.
|
||||
# NJT guarantees to have its tangent as NJT, because it has dedicated integration in Autograd
|
||||
# See torch/csrc/autograd/python_function.cpp, use_zeros_like.
|
||||
(
|
||||
(
|
||||
_plain_fake_tensor_like_subclass(inp)
|
||||
if is_traceable_wrapper_subclass(inp)
|
||||
and not _is_subclass_mutated_input_tangent_always_subclass(inp)
|
||||
else inp
|
||||
),
|
||||
TangentAOTInput(InputMutationAOTOutput(inp_desc)),
|
||||
)
|
||||
for inp, inp_desc, info in zip(flat_f_args, flat_f_args_descs, input_info)
|
||||
if info.mutation_type == MutationType.MUTATED_OUT_GRAPH
|
||||
and info.mutates_data
|
||||
and info.requires_grad
|
||||
]
|
||||
f_input_tangents, f_input_tangents_descs = (
|
||||
[x[0] for x in f_input_tangents_pairs],
|
||||
[x[1] for x in f_input_tangents_pairs],
|
||||
)
|
||||
|
||||
f_output_tangents_pairs = [
|
||||
(o, TangentAOTInput(desc))
|
||||
for o, info, desc in zip(flat_f_outs, output_info, flat_f_outs_descs)
|
||||
if info.output_type
|
||||
in [
|
||||
OutputType.non_alias,
|
||||
OutputType.unsafe_view_alias,
|
||||
OutputType.custom_function_view,
|
||||
]
|
||||
and issubclass(info.raw_type, torch.Tensor)
|
||||
and info.requires_grad_for_backward
|
||||
]
|
||||
f_output_tangents, f_output_tangents_descs = (
|
||||
[x[0] for x in f_output_tangents_pairs],
|
||||
[x[1] for x in f_output_tangents_pairs],
|
||||
)
|
||||
|
||||
# intermediate bases are also included in the backward graph
|
||||
f_tangents = f_input_tangents + f_output_tangents + intermediate_bases
|
||||
f_tangents_descs = (
|
||||
f_input_tangents_descs + f_output_tangents_descs + intermediate_bases_descs
|
||||
)
|
||||
|
||||
# TODO: I'm pretty sure you don't need a tree_map here
|
||||
traced_tangents = pytree.tree_map(from_fun, f_tangents)
|
||||
traced_tangents = pytree.tree_map(
|
||||
view_avoid_dupes_with_primals, traced_tangents
|
||||
)
|
||||
traced_tangents = [
|
||||
coerce_tangent_and_suggest_memory_format(tt)[0]
|
||||
for i, tt in enumerate(traced_tangents)
|
||||
]
|
||||
# NB: update this if the maps above ever change structure.
|
||||
# Also, it might be helpful to add coercion information to the tangent desc!
|
||||
traced_tangents_descs = f_tangents_descs
|
||||
|
||||
nonlocal static_input_indices
|
||||
static_input_indices = static_input_indices or []
|
||||
if torch._dynamo.compiled_autograd.in_compiled_autograd_region:
|
||||
passed_indices = set(static_input_indices)
|
||||
static_input_indices = [
|
||||
i
|
||||
for i, arg in enumerate(flat_args)
|
||||
if (isinstance(arg, torch.nn.Parameter) or i in passed_indices)
|
||||
]
|
||||
|
||||
static_input_logger.debug(
|
||||
"static input indices metadata analysis: %s", static_input_indices
|
||||
)
|
||||
|
||||
f_mutated_inputs = [
|
||||
inp
|
||||
for inp, info in zip(flat_f_args, input_info)
|
||||
if info.mutation_type == MutationType.MUTATED_OUT_GRAPH
|
||||
]
|
||||
# Build the full list of forward graph outputs so the subclass wrapping
|
||||
# code knows exactly which graph outputs to wrap back into subclasses.
|
||||
# Including intermediate_bases unconditionally is safe: they are only
|
||||
# populated when outputs require grad (line ~539), so they are naturally
|
||||
# empty during pure inference. In the "downgrade from training to
|
||||
# inference" path, num_intermediate_bases > 0 is already gated behind
|
||||
# `assert not req_subclass_dispatch` (aot_autograd.py), so the subclass
|
||||
# wrapping code that consumes subclass_fw_graph_out_meta never sees them.
|
||||
f_fw_graph_outs = [*f_mutated_inputs, *flat_f_outs, *intermediate_bases]
|
||||
fw_graph_outs = pytree.tree_map(from_fun, f_fw_graph_outs)
|
||||
|
||||
grad_enabled_mutation = None
|
||||
if torch.is_grad_enabled() != prior_grad_enabled:
|
||||
grad_enabled_mutation = torch.is_grad_enabled()
|
||||
torch.set_grad_enabled(
|
||||
prior_grad_enabled
|
||||
) # Restore the prior state after tracing it
|
||||
log.debug(
|
||||
(
|
||||
"grad_mode mutation encountered in graph. "
|
||||
"Will emit mutation epilogue, to set grad_mode=%s"
|
||||
),
|
||||
grad_enabled_mutation,
|
||||
)
|
||||
|
||||
subclass_inp_meta = create_subclass_meta(flat_args)
|
||||
subclass_fw_graph_out_meta = create_subclass_meta(fw_graph_outs)
|
||||
subclass_tangent_meta = create_subclass_meta(
|
||||
traced_tangents, count_symints=False, with_memory_format=True
|
||||
)
|
||||
|
||||
metadata = ViewAndMutationMeta(
|
||||
input_info=input_info,
|
||||
output_info=output_info,
|
||||
num_intermediate_bases=len(intermediate_bases),
|
||||
keep_input_mutations=keep_input_mutations,
|
||||
traced_tangents=traced_tangents,
|
||||
traced_tangents_descs=traced_tangents_descs,
|
||||
subclass_inp_meta=subclass_inp_meta,
|
||||
subclass_fw_graph_out_meta=subclass_fw_graph_out_meta,
|
||||
subclass_tangent_meta=subclass_tangent_meta,
|
||||
grad_enabled_mutation=grad_enabled_mutation,
|
||||
static_input_indices=static_input_indices,
|
||||
tokens=mode._tokens,
|
||||
)
|
||||
return metadata
|
||||
|
||||
return inner
|
||||
+782
@@ -0,0 +1,782 @@
|
||||
"""
|
||||
AOTAutograd descriptors are a path-like data structure (similar to pytree
|
||||
paths and sources) that describe the semantic meaning of an input/output to FX
|
||||
graphs. Although you may know the input/output meaning at the top level of
|
||||
the original function you traced, because we have many graph capture wrappers
|
||||
that change the calling convention, it can be difficult to tell how these
|
||||
correspond to the actual FX graph you get back, to say nothing about the extra
|
||||
arguments/outputs for tangents, gradients, etc. Descriptors describe the meaning
|
||||
of arguments.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Before we talk about the precise semantics, it's helpful to look at some
|
||||
examples to get some intuition for the meaning of descriptors. Here are some
|
||||
input descriptors you might find on the joint FX graph:
|
||||
|
||||
* PlainAOTInput(idx=0) - the first input from the original callable, as is
|
||||
|
||||
* ParamAOTInput(target="mod.weight") - the parameter with FQN mod.weight
|
||||
|
||||
* TangentAOTInput(output=PlainAOTOutput(idx=1)) - the input tangent
|
||||
corresponding to the gradients for the second output in the forward graph
|
||||
|
||||
* ViewBaseAOTInput(base_of=PlainAOTInput(idx=0)) - it turned out the first
|
||||
input was actually a (differentiable) view of a tensor which aliased with
|
||||
another input tensor. We replaced this input with a single input for the
|
||||
base of all of these inputs, replacing the original inputs (one of which is
|
||||
mentioned in base_of). We would generate a GradAOTOutput for *this* input
|
||||
(and not the original PlainAOTInputs!) If you have a joint graph where a
|
||||
view base like this is undesirable, you can eliminate this by cloning
|
||||
the views outside of the compiled region (assuming you aren't mutating this
|
||||
tensor).
|
||||
|
||||
* SubclassGetAttrAOTInput(base=AOTInput(idx=0), attr="inner") - this tensor
|
||||
corresponds to the "inner" tensor from the tensor subclass that is at the
|
||||
first index. In general, joint graphs from AOTAutograd never take tensor
|
||||
subclasses as inputs; they are always unpacked into their constituent plain
|
||||
tensor pieces; use the descriptors to identify the parts of the tensor that
|
||||
are related. Note that this can be nested (if you have nested tensor
|
||||
subclasses!)
|
||||
|
||||
Here are some output descriptors you might find on the Joint FX graph:
|
||||
|
||||
* PlainAOTOutput(idx=0) - the first output from the original forward function,
|
||||
as is
|
||||
|
||||
* GradAOTOutput(grad_of=PlainAOTInput(idx=1)) - the computed gradient for the
|
||||
second input to the graph, an output of the backward graph
|
||||
|
||||
* InputMutationAOTOutput(mutated_input=PlainAOTInput(idx=0)) - when the first
|
||||
input is mutated, the new value to be copied into the first input of the
|
||||
graph. Sometimes, these outputs can be elided and the ``copy_`` is done directly
|
||||
in the graph (controlled by keep_input_mutations), but if the input
|
||||
mutation must be differentiated through we always generate an output like this
|
||||
|
||||
* IntermediateBaseAOTOutput(base_of=PlainAOTOutput(idx=0)) - if we return
|
||||
multiple outputs which alias each other, we instead replace them with a single
|
||||
output tensor representing the base of all the aliases. This output indicates
|
||||
it is the base for /one/ of those original outputs. If this is undesirable in
|
||||
the joint graph, clone all outputs before returning from the graph.
|
||||
|
||||
* SubclassGetAttrAOTOutput(base=PlainAOTOutput(idx=0), idx="inner") - this
|
||||
tensor correspondings to the inner tensor of the first original output which
|
||||
is a tensor subclass. This and other subclass components of that output will
|
||||
get repacked into a tensor subclass.
|
||||
|
||||
High level semantics
|
||||
--------------------
|
||||
|
||||
OK, let's formally define a descriptor. Intuitively, suppose we have::
|
||||
|
||||
def wrapped_graph(*args):
|
||||
ret = graph(*in_transform(args))
|
||||
return out_transform(ret)
|
||||
|
||||
Then the descriptor for input[i] to graph describes a function fin_i such that::
|
||||
|
||||
fin_i(args) == in_transform(args)[i]
|
||||
|
||||
and the descriptor for output[j] from graph describes a function fout_j such that::
|
||||
|
||||
fout_j(out_transform(ret)) == ret[j]
|
||||
|
||||
AKA input descriptors tell you how to get from outer inputs to inner inputs,
|
||||
while output descriptors tell you how to get from outer outputs to inner
|
||||
outputs (inverse data flow!)
|
||||
|
||||
We haven't said anything about what these transformations actually do. There
|
||||
are three major transformations AOTAutograd does (performed in this order):
|
||||
|
||||
* View/mutation handling
|
||||
* Autograd
|
||||
* Subclasses
|
||||
|
||||
So intuitively, descriptors are built like this:
|
||||
|
||||
1. **PlainAOTInput, PlainAOTOutput.**
|
||||
|
||||
We start off descriptors describing the exact inputs/outputs of the
|
||||
original flattened user function. This user function is assumed to already
|
||||
be flattened; you would chain on pytree KeyPaths to further describe where
|
||||
in the pytree each input/output lived if you needed to deal with
|
||||
unflattened functions: this can be done from userland on top of
|
||||
descriptors, so the main descriptors mechanism doesn't handle it.
|
||||
|
||||
2. **SyntheticBaseAOTInput, ViewBaseAOTInput, MetadataMutationAOTOutput,
|
||||
InputMutationAOTOutput, IntermediateBaseAOTOutput**
|
||||
|
||||
We deal with mutations and aliasing by removing duplicate PlainAOTInputs
|
||||
and introduce some new artificial inputs/outputs. These inputs do not
|
||||
have a straightforward correspondence to the original user inputs, but if
|
||||
you are implementing a pass that doesn't care about the exact semantics of
|
||||
inputs, you should handle all of these uniformly in the same way as regular
|
||||
inputs.
|
||||
|
||||
3. **TangentAOTInput, GradAOTOutput**
|
||||
|
||||
We deal with autograd by introducing a tangent input for every
|
||||
differentiable AOTOutput (including the new ones introduced above), and a
|
||||
gradient output for every differentiable AOTInput (also including new ones
|
||||
introduced above.) The arguments to these AOTInput/AOTOutput can ONLY be
|
||||
the ones we already have above (from steps 1-2). As AOTAutograd does not
|
||||
currently support double backwards, you never have tangents of grads or
|
||||
vice versa (but in the future we could!)
|
||||
|
||||
4. **SubclassGetAttrAOTInput, SubclassGetAttrAOTOutput, et al.**
|
||||
|
||||
We deal with subclasses by introducing flattened inputs/outputs (including
|
||||
potentially symbolic sizes/strides) for every AOTInput/AOTOutput that was a
|
||||
subclass. As above, the arguments to these AOTInput/AOTOutput can ONLY be
|
||||
the ones we have above (from steps 1-3). Recursive subclasses are
|
||||
supported, so these descriptors can nest with each other (so descriptors
|
||||
from step 4 are fair game as well.)
|
||||
|
||||
5. **ForwardTokenAOTInput, ForwardTokenAOTOutput, BackwardTokenAOTInput, BackwardTokenAOTOutput.**
|
||||
|
||||
Some extra token inputs/outputs get added, these are synthetic and are just here to
|
||||
prevent DCE/reordering.
|
||||
|
||||
The important thing about the pipeline is that descriptors can ONLY be
|
||||
created from top-to-bottom. So for example, you can have::
|
||||
|
||||
SubclassGetAttrAOTInput(TangentAOTInput(PlainAOTOutput(...))) # OK
|
||||
|
||||
As you can see that PlainAOTOutput -> TangentAOTInput ->
|
||||
SubclassGetAttrAOTInput is consistent with the pipeline ordering), but you can
|
||||
NEVER have::
|
||||
|
||||
TangentAOTInput(SubclassGetAttrAOTOutput(PlainAOTOutput(...)) # BAD
|
||||
|
||||
This is inconsistent; we always do autograd BEFORE we process subclasses!
|
||||
|
||||
Similarly, for example, this is illegal::
|
||||
|
||||
GradAOTOutput(SubclassGetAttrAOTInput(PlainAOTInput(...))) # BAD
|
||||
|
||||
It is illegal because subclasses are handled *after* create joint during
|
||||
wrapper construction. Instead, you would have::
|
||||
|
||||
SubclassGetAttrAOTOutput(GradAOTOutput(PlainAOTInput(...))) # OK
|
||||
|
||||
This intuitively captures the fact that we always to autograd directly on the
|
||||
subclass, rather than after desugaring the subclass into its inner tensors.
|
||||
|
||||
Descriptor index
|
||||
----------------
|
||||
|
||||
Here is a list of all AOTInput/AOTOutput, organized by how likely you need to
|
||||
handle them:
|
||||
|
||||
* AOTInput
|
||||
|
||||
* Important:
|
||||
|
||||
* PlainAOTInput (the primals!)
|
||||
* ParamAOTInput
|
||||
* TangentAOTInput
|
||||
* SubclassGetAttrAOTInput et al. (if you use subclasses)
|
||||
|
||||
* View related (can be eliminated by cloning inputs to graph; if you don't
|
||||
eliminate them, make sure to handle pairing them with GradAOTOutput):
|
||||
|
||||
* ViewBaseAOTInput
|
||||
* SyntheticBaseAOTInput
|
||||
|
||||
* Non-tensor, mostly just ignore them:
|
||||
|
||||
* DummyAOTInput
|
||||
* PhiloxForwardSeedAOTInput
|
||||
* PhiloxForwardBaseOffsetAOTInput
|
||||
* PhiloxBackwardSeedAOTInput
|
||||
* PhiloxBackwardBaseOffsetAOTInput
|
||||
* ForwardTokenAOTInput
|
||||
* BackwardTokenAOTInput
|
||||
|
||||
* AOTOutput
|
||||
|
||||
* Important:
|
||||
|
||||
* PlainAOTOutput
|
||||
* GradAOTOutput
|
||||
* SubclassGetAttrAOTOutput et al. (if you use subclasses)
|
||||
|
||||
* More obscure (if not eliminated, make sure you handle pairing them with
|
||||
TangentAOTInput):
|
||||
|
||||
* InputMutationAOTOutput (can be eliminated if mutations are non-differentiable)
|
||||
* IntermediateBaseAOTOutput (can be eliminated by cloning outputs of graph)
|
||||
* MetadataMutationAOTOutput (uhh, just don't mutate metadata?)
|
||||
|
||||
* Non-tensor, mostly just ignore them:
|
||||
|
||||
* PhiloxUpdatedForwardOffsetAOTOutput
|
||||
* PhiloxUpdatedBackwardOffsetAOTOutput
|
||||
* ForwardTokenAOTOutput
|
||||
* BackwardTokenAOTOutput
|
||||
* DummyAOTOutput
|
||||
|
||||
For convenience, we also have DifferentiableAOTInput and
|
||||
DifferentiableAOTOutput to help you classify which inputs/outputs can be
|
||||
wrapped by GradAOTOutput/TangentAOTInput (respectively), which are essentially
|
||||
all tensor AOTInput/AOTOutput excluding the subclass descriptors.
|
||||
|
||||
Implementation details
|
||||
----------------------
|
||||
|
||||
The stylized view above is good for understanding how to interpret
|
||||
descriptors, but the way that descriptors are generated in code is a bit more
|
||||
complicated. Specifically, AOTAutograd is structured as a series of wrappers
|
||||
on the original user function, which are composed together to form the final
|
||||
function to trace. As a result of this, AOTAutograd ends up first building
|
||||
the full AOTInputs for a function to be traced (as it builds the wrappers and
|
||||
modifies the flat arguments to be compatible with the new input signature of
|
||||
the wrapper), and then in reverse builds up the AOTOutput as it is tracing.
|
||||
|
||||
There is one major exception to this general idea of "build AOTInput first",
|
||||
and then "build AOTOutput second": when we create TangentAOTInput, we need to
|
||||
reference AOTOutputs (which output we are the tangents of) which we generally
|
||||
haven't created yet. There's two ways we deal with this:
|
||||
|
||||
- After the precompile steps (dedup and synthetic base handling), we do an
|
||||
initial pass to collect forward metadata that produces the initial set of
|
||||
PlainAOTOutputs which we use to create the tangent inputs.
|
||||
|
||||
- We also sometimes just violate causality and predict that an AOTOutput will
|
||||
be created in a particular way at some later point in time when we build an
|
||||
AOTInput.
|
||||
|
||||
As of July 2025, here is an exhaustive description of how inputs/outputs
|
||||
traverse the wrappers from AOTAutograd, and what descriptors can be introduced
|
||||
at these phases.
|
||||
|
||||
::
|
||||
|
||||
Build wrappers (FLOWS DOWN) Run trace (FLOWS UP)
|
||||
-------------------------------------------------------------------------------------------------
|
||||
Begin PlainAOTInput (n/a)
|
||||
ParamAOTInput
|
||||
|
||||
Precompile dedupe (remove dupes) (nothing)
|
||||
|
||||
Precompile synthetic base SyntheticBaseAOTInput MetadataMutationAOTOutput
|
||||
ViewBaseAOTInput
|
||||
|
||||
Forward metadata trace PlainAOTOutput (n/a)
|
||||
MetadataMutationAOTOutput
|
||||
|
||||
Prepare for autograd (nothing) InputMutationAOTOutput
|
||||
IntermediateBaseAOTOutput
|
||||
|
||||
Create joint TangentAOTInput GradAOTOutput
|
||||
w/ InputMutationAOTOutput
|
||||
w/ IntermediateBaseAOTOutput
|
||||
|
||||
Precompile subclass SubclassGetAttrAOTInput et al. SubclassGetAttrAOTOutput et al.
|
||||
|
||||
Effect tokens ForwardTokenAOTInput ForwardTokenAOTOutput
|
||||
BackwardTokenAOTInput BackwardTokenAOTOutput
|
||||
|
||||
End (n/a) PlainAOTOutput
|
||||
|
||||
It can be helpful to separately write down the input flow and the output flow
|
||||
for ease of understanding the data flow:
|
||||
|
||||
* Input desc propagation (happens as we build wrappers)
|
||||
|
||||
* [IN] Begin with original calling convention (PlainAOTInput, ParamAOTInput)
|
||||
* [IN] Precompile dedupe: (removes duplicate AOTInputs)
|
||||
* [IN] Precompile synthetic base: SyntheticBaseAOTInput, ViewBaseAOTInput
|
||||
* Forward metadata trace (mini output desc propagation)
|
||||
|
||||
* [OUT] Original output convention: PlainAOTOutput
|
||||
* [OUT] Precompile synthetic base: MetadataMutationAOTOutput
|
||||
|
||||
* [IN] Prepare for autograd: (nothing)
|
||||
* [IN] Create joint: TangentAOTInput (potentially w/
|
||||
IntermediateBaseAOTOutput, InputMutationAOTOutput)
|
||||
* [IN] Precompile subclass: SubclassGetAttrAOTInput et al.
|
||||
* [IN] Effect tokens: ForwardTokenAOTInput, BackwardTokenAOTInput
|
||||
(Note: BackwardTokenAOTInput is technically generated not by a wrapper but
|
||||
actually done by token_discovery which implicitly adds extra arguments
|
||||
to the FX trace on-the-fly.)
|
||||
|
||||
* Trigger a trace with the modified inputs on the wrapper
|
||||
* Output desc propagation (happens as we unwind from the user function call in trace)
|
||||
|
||||
* [OUT] Begin with original calling convention: PlainAOTOutput
|
||||
* [OUT] Effect tokens: ForwardTokenAOTOutput, BackwardTokenAOTOutput
|
||||
* [OUT] Precompile subclass: SubclassGetAttrAOTOutput et al.
|
||||
* [OUT] Create joint: GradAOTOutput
|
||||
* [OUT] Prepare for autograd: InputMutationAOTOutput, IntermediateBaseAOTOutput
|
||||
* [OUT] Precompile synthetic base: MetadataMutationAOTOutput
|
||||
* [OUT] Precompile dedupe: (nothing)
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
|
||||
|
||||
# TODO: the is_* predicates are a little suspicious because (1) they're not
|
||||
# used by anything and (2) they always report False even when a parameter got
|
||||
# swizzled into a view base or deduped with a non-parameter. It is pretty
|
||||
# difficult to exercise these cases but it's not clear if you will write code
|
||||
# that works correctly in those cases.
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class AOTInput:
|
||||
"""Describes where an input from an AOTAutograd produced FX graph comes from"""
|
||||
|
||||
def expr(self) -> str:
|
||||
raise NotImplementedError("Subclasses must implement expr()")
|
||||
|
||||
def is_param(self) -> bool:
|
||||
"""True if this input is a parameter or derived from a parameter (e.g., subclass attr)"""
|
||||
return False
|
||||
|
||||
def is_buffer(self) -> bool:
|
||||
"""True if this input is a buffer or derived from a buffer (e.g., subclass attr)"""
|
||||
return False
|
||||
|
||||
def is_tangent(self) -> bool:
|
||||
"""True if this input is a tangent or derived from a tangent (e.g., subclass attr)"""
|
||||
return False
|
||||
|
||||
|
||||
# Note: Currently, our typing discipline for differentiable versus not is not
|
||||
# very good, so feel free to rely on runtime tests instead.
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class DifferentiableAOTInput(AOTInput):
|
||||
"""A subclass that classifies AOTInput that can be wrapped by GradAOTOutput"""
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class AOTOutput:
|
||||
"""Describes where an output from an AOTAutograd produced FX graph will
|
||||
eventually be bundled into the final output"""
|
||||
|
||||
def expr(self) -> str:
|
||||
raise NotImplementedError("Subclasses must implement expr()")
|
||||
|
||||
def is_grad(self) -> bool:
|
||||
"""True if this output is a grad or derived from a grad (e.g., subclass attr)"""
|
||||
return False
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class DifferentiableAOTOutput(AOTOutput):
|
||||
"""A subclass that classifies AOTOutput that can be wrapped by TangentAOTInput"""
|
||||
|
||||
|
||||
# ------------
|
||||
|
||||
# AOTInput
|
||||
|
||||
# ------------
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ParamAOTInput(DifferentiableAOTInput):
|
||||
"""The input is a parameter, whose FQN is target"""
|
||||
|
||||
target: str
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"self.get_parameter({self.target!r})"
|
||||
|
||||
def is_param(self) -> bool:
|
||||
return True
|
||||
|
||||
def is_buffer(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class BufferAOTInput(DifferentiableAOTInput):
|
||||
"""The input is a buffer, whose FQN is target"""
|
||||
|
||||
target: str
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"self.get_buffer({self.target!r})"
|
||||
|
||||
def is_param(self) -> bool:
|
||||
return False
|
||||
|
||||
def is_buffer(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class DummyAOTInput(AOTInput):
|
||||
"""In some circumstances, we want to call into a function that expects AOTInput, but
|
||||
we don't actually care about that logic (most typically, because some code is being used
|
||||
for both compile-time and run-time; AOTInput processing is not needed in this situation.
|
||||
Pass a dummy in this situation; but it is better to just have a version of the function
|
||||
that doesn't have this at all."""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__dummy{self.idx}"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PlainAOTInput(DifferentiableAOTInput):
|
||||
"""The input is a plain input, corresponding to a particular positional index.
|
||||
|
||||
Note that AOTInput is always relative to a function with a *flat* calling convention,
|
||||
e.g., as accepted by `aot_module_simplified`. There are some AOTAutograd APIs that
|
||||
flatten pytrees, and we don't record PyTree key paths from the flattening (but we
|
||||
could and should!)
|
||||
"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"args[{self.idx}]"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassGetAttrAOTInput(AOTInput):
|
||||
"""Subclass inputs get unpacked into their constituent pieces before going into an FX
|
||||
graph. This tells you which particular attribute of the subclass this particular
|
||||
input corresponds to (of the 'base' originally subclass argument.)
|
||||
"""
|
||||
|
||||
base: AOTInput
|
||||
attr: str
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.{self.attr}"
|
||||
|
||||
def is_param(self) -> bool:
|
||||
return self.base.is_param()
|
||||
|
||||
def is_buffer(self) -> bool:
|
||||
return self.base.is_buffer()
|
||||
|
||||
def is_tangent(self) -> bool:
|
||||
return self.base.is_tangent()
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassSizeAOTInput(AOTInput):
|
||||
"""Which subclass this particular outer size SymInt input (at dim idx) came from."""
|
||||
|
||||
base: AOTInput
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.size({self.idx})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassStrideAOTInput(AOTInput):
|
||||
"""Which subclass this particular outer stride SymInt input (at dim idx) came from."""
|
||||
|
||||
base: AOTInput
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.stride({self.idx})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ViewBaseAOTInput(DifferentiableAOTInput):
|
||||
"""
|
||||
When multiple differentiable inputs are views of the same input, AOTAutograd will replace all of these
|
||||
views with a single input representing the base. If this is undesirable, you can clone the views
|
||||
example inputs before passing them into AOTAutograd.
|
||||
|
||||
TODO: In principle we could report ALL of the inputs who this is a base of.
|
||||
"""
|
||||
|
||||
base_of: AOTInput
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base_of.expr()}._base"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SyntheticBaseAOTInput(DifferentiableAOTInput):
|
||||
"""This is similar to ViewBaseAOTInput, but this happens when none of the views were differentiable, so
|
||||
we weren't able to get our hands on the true original view and constructed a synthetic one instead
|
||||
for the sake of autograd.
|
||||
"""
|
||||
|
||||
base_of: AOTInput
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__make_synthetic_base({self.base_of.expr()})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxForwardSeedAOTInput(AOTInput):
|
||||
"""The seed for functionalized Philox RNG calls, specifically for forward graph."""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_forward_seed"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxForwardBaseOffsetAOTInput(AOTInput):
|
||||
"""The offset for functionalized Philox RNG calls, specifically for forward graph."""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_forward_base_offset"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxBackwardSeedAOTInput(AOTInput):
|
||||
"""The seed for functionalized Philox RNG calls, specifically for backward graph."""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_backward_seed"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxBackwardBaseOffsetAOTInput(AOTInput):
|
||||
"""The offset for functionalized Philox RNG calls, specifically for backward graph."""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_backward_base_offset"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ForwardTokenAOTInput(AOTInput):
|
||||
"""The world token which is threaded through side-effectful operations"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__forward_token{self.idx}"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class BackwardTokenAOTInput(AOTInput):
|
||||
"""The world token which is threaded through side-effectful operations, for backwards"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__backward_token{self.idx}"
|
||||
|
||||
|
||||
# Technically the "output" here is redundant, tangents always correspond to
|
||||
# outputs
|
||||
# NB: this is marked differentiable as it /would/ be differentiable if we
|
||||
# support double backwards, but we never generate this today because we
|
||||
# don't support double backwards.
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class TangentAOTInput(DifferentiableAOTInput):
|
||||
"""An input to the joint graph representing the tangent of an output."""
|
||||
|
||||
output: DifferentiableAOTOutput
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.output, DifferentiableAOTOutput):
|
||||
raise AssertionError(
|
||||
f"expected output to be DifferentiableAOTOutput, got {type(self.output)}"
|
||||
)
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__output_tangent({self.output.expr()})"
|
||||
|
||||
def is_tangent(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
# ------------
|
||||
|
||||
# AOTOutput
|
||||
|
||||
# ------------
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PlainAOTOutput(DifferentiableAOTOutput):
|
||||
"""A plain tensor output at position idx of the output tuple"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"output[{self.idx}]"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class InputMutationAOTOutput(DifferentiableAOTOutput):
|
||||
"""The mutated value of an input tensor, returned so we can appropriately propagate autograd."""
|
||||
|
||||
mutated_input: AOTInput
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__input_mutation({self.mutated_input.expr()})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class IntermediateBaseAOTOutput(DifferentiableAOTOutput):
|
||||
"""An intermediate base of multiple outputs which alias each other. We only report ONE of
|
||||
the outputs that contributed to this base"""
|
||||
|
||||
base_of: "AOTOutput"
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__intermediate_base({self.base_of.expr()})"
|
||||
|
||||
|
||||
# TODO: it's a little dodgy this is differentiable lol, but we do generate
|
||||
# these BEFORE autograd is handled
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class MetadataMutationAOTOutput(DifferentiableAOTOutput):
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__aliased_arg_with_metadata_mutation{self.idx}"
|
||||
|
||||
|
||||
# NB: this is marked differentiable as it /would/ be differentiable if we
|
||||
# support double backwards, but we never generate this today because we
|
||||
# don't support double backwards.
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class GradAOTOutput(DifferentiableAOTOutput):
|
||||
"""An output representing the computed gradient for a differentiable input, in the joint graph"""
|
||||
|
||||
grad_of: DifferentiableAOTInput
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.grad_of, DifferentiableAOTInput):
|
||||
raise AssertionError(
|
||||
f"expected grad_of to be DifferentiableAOTInput, got {type(self.grad_of)}"
|
||||
)
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__grad({self.grad_of.expr()})"
|
||||
|
||||
def is_grad(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxUpdatedForwardOffsetAOTOutput(AOTOutput):
|
||||
"""The final offset from the functionalized RNG calls, forward only"""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_updated_forward_offset"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class PhiloxUpdatedBackwardOffsetAOTOutput(AOTOutput):
|
||||
"""The final offset from the functionalized RNG calls, backward only"""
|
||||
|
||||
def expr(self) -> str:
|
||||
return "__philox_updated_backward_offset"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ForwardTokenAOTOutput(AOTOutput):
|
||||
"""The world token output for side-effectful calls, returned so we cannot DCE it, forward only"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__forward_token{self.idx}"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class BackwardTokenAOTOutput(AOTOutput):
|
||||
"""The world token output for side-effectful calls, returned so we cannot DCE it, backward only"""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__backward_token{self.idx}"
|
||||
|
||||
|
||||
# These are seemingly symmetric with their AOTInput counterparts. The way to
|
||||
# think about it is that a subclass could be an input or an output, and they
|
||||
# get exploded into plain tensors on the way in and out. So we need
|
||||
# descriptors for both.
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassGetAttrAOTOutput(AOTOutput):
|
||||
"""This output will be bundled into a subclass at this location"""
|
||||
|
||||
base: AOTOutput
|
||||
attr: str
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.{self.attr}"
|
||||
|
||||
def is_grad(self) -> bool:
|
||||
return self.base.is_grad()
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassSizeAOTOutput(AOTOutput):
|
||||
"""This output size will be bundled into a subclass at this location"""
|
||||
|
||||
base: AOTOutput
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.size({self.idx})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SubclassStrideAOTOutput(AOTOutput):
|
||||
"""This output stride will be bundled into a subclass at this location"""
|
||||
|
||||
base: AOTOutput
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"{self.base.expr()}.stride({self.idx})"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class DummyAOTOutput(AOTOutput):
|
||||
"""For cases when you don't actually care about descriptor propagation, do not use under normal
|
||||
circumstances."""
|
||||
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__dummy{self.idx}"
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SavedForBackwardsAOTOutput(AOTOutput):
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__saved_for_backwards_{self.idx}"
|
||||
|
||||
|
||||
# Note [Activations with no version counter checks in eager]
|
||||
# In eager, when a tensor is saved for backward, the autograd engine
|
||||
# generally performs version counter checks when grabbing the activation
|
||||
# at backward time.
|
||||
# One exception is in a custom autograd.Function: if the user stashes a tensor
|
||||
# for use in the backward using `ctx.foo = foo`, rather than
|
||||
# `ctx.save_for_backward(foo)`, then the autograd engine will not know
|
||||
# to perform any checks.
|
||||
# In torch.compile, we handle autograd by tracing through it ahead of time,
|
||||
# and wrapping the fw + bw graphs into a custom autograd.Function.
|
||||
# In order to provide parity with eager around VC checks, though, we need
|
||||
# to know which activations had gotten this "no VC checks" treatment in eager,
|
||||
# so we can plumb this info into the AOTAutograd runtime,
|
||||
# and avoid calling ctx.save_for_backward on these tensors.
|
||||
#
|
||||
# This dataclass tells us that a given AOTOutput corresponds to:
|
||||
# - An activation output from the forward graph
|
||||
# - that should *not* have its version counter checked at runtime in the backward.
|
||||
# This is done by stashing this activation on the ctx object directly
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class SavedForBackwardsNoVcCheckAOTOutput(AOTOutput):
|
||||
idx: int
|
||||
|
||||
def expr(self) -> str:
|
||||
return f"__saved_for_backwards_no_vc_check_{self.idx}"
|
||||
+421
@@ -0,0 +1,421 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, cast, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._guards import detect_fake_mode
|
||||
from torch._library.opaque_object import is_opaque_type
|
||||
from torch._opaque_base import OpaqueBase
|
||||
from torch._subclasses import FakeTensor, FakeTensorMode
|
||||
from torch.fx.experimental.proxy_tensor import _pytree_subclasses_that_lose_info
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
|
||||
|
||||
from .. import config
|
||||
from .descriptors import BufferAOTInput, DifferentiableAOTInput, ParamAOTInput
|
||||
from .schemas import AOTConfig, FakifiedFlatArgs
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator, KeysView
|
||||
|
||||
|
||||
static_inputs_log = torch._logging.getArtifactLogger(
|
||||
__name__, "cudagraph_static_inputs"
|
||||
)
|
||||
|
||||
|
||||
def process_inputs(
|
||||
flat_args: list[Any],
|
||||
aot_config: AOTConfig,
|
||||
fake_mode: FakeTensorMode,
|
||||
shape_env: ShapeEnv | None,
|
||||
ignore_shape_env: bool = False,
|
||||
) -> tuple[FakifiedFlatArgs, list[int]]:
|
||||
"""Convert real tensor inputs into fake tensors for AOT autograd tracing.
|
||||
|
||||
Called at compile time (not runtime) to produce the fake inputs that AOT
|
||||
autograd traces through. Each real tensor is converted to a FakeTensor
|
||||
via ``fake_mode.from_tensor``, preserving shape, dtype, device, and
|
||||
symbolic shape information from the ShapeEnv. Non-tensor inputs (ints,
|
||||
SymInts, ScriptObjects) are converted or passed through as appropriate.
|
||||
|
||||
Tensor subclass inputs (DTensor, etc.) are fakified recursively by
|
||||
walking their ``__tensor_flatten__`` attrs. AsyncCollectiveTensors are
|
||||
resolved via ``trigger_wait()`` before fakification so they don't appear
|
||||
in the traced metadata (see below).
|
||||
|
||||
Called from ``aot_function``, ``aot_module_simplified``, and
|
||||
``aot_export_module`` — anywhere AOT autograd needs fake inputs before
|
||||
graph capture.
|
||||
|
||||
Returns:
|
||||
A tuple of (fakified_args, act_input_indices) where act_input_indices
|
||||
records which positions held AsyncCollectiveTensors. These indices are
|
||||
stored on ViewAndMutationMeta so that the runtime wrapper can emit
|
||||
direct trigger_wait() calls on those positions.
|
||||
"""
|
||||
# Resolve AsyncCollectiveTensors before tracing. ACTs are transient
|
||||
# eager-mode wrappers for async collective overlap; if they leak into the
|
||||
# traced graph as input types, AOT autograd records them in
|
||||
# SubclassCreationMeta for output tangent metadata. At runtime, autograd
|
||||
# produces plain tensor tangents, causing a type mismatch. Unwrapping
|
||||
# here prevents ACT from appearing in the traced metadata.
|
||||
try:
|
||||
from torch.distributed._functional_collectives import AsyncCollectiveTensor
|
||||
except ImportError:
|
||||
AsyncCollectiveTensor = None
|
||||
|
||||
act_input_indices: list[int] = []
|
||||
if AsyncCollectiveTensor is not None:
|
||||
for i, a in enumerate(flat_args):
|
||||
if isinstance(a, AsyncCollectiveTensor):
|
||||
act_input_indices.append(i)
|
||||
flat_args[i] = a.trigger_wait()
|
||||
|
||||
with fake_mode:
|
||||
|
||||
def convert(idx: int, x: Any) -> Any:
|
||||
nonlocal ignore_shape_env
|
||||
if shape_env is not None and not ignore_shape_env:
|
||||
from torch._dynamo.source import ConstantSource
|
||||
|
||||
if isinstance(x, int):
|
||||
# We always specialize on scalar values in export.
|
||||
if aot_config.is_export:
|
||||
return x
|
||||
source = ConstantSource(f"sym_{idx}")
|
||||
return shape_env.create_symintnode(
|
||||
shape_env.create_symbol(x, source, positive=x >= 0),
|
||||
hint=x,
|
||||
source=source,
|
||||
)
|
||||
if isinstance(x, torch.ScriptObject) or is_opaque_type(type(x)):
|
||||
return torch._library.fake_class_registry.maybe_to_fake_obj(
|
||||
fake_mode, x
|
||||
)
|
||||
if not isinstance(x, torch.Tensor):
|
||||
return x
|
||||
if isinstance(x, FakeTensor):
|
||||
# In the case of cross compilation we will have example inputs
|
||||
# with a different fake mode than our tracing fake mode.
|
||||
# In these cases we want to clone the fake tensor into our
|
||||
# inner fake mode.
|
||||
if x.fake_mode is not fake_mode:
|
||||
return fake_mode.from_tensor(x)
|
||||
return x
|
||||
if is_traceable_wrapper_subclass(x):
|
||||
attrs, _ = x.__tensor_flatten__()
|
||||
# See if all inner tensors are FakeTensors from this mode
|
||||
all_this_fake = True
|
||||
for a in attrs:
|
||||
match getattr(x, a):
|
||||
case FakeTensor() as v:
|
||||
if v.fake_mode is not fake_mode:
|
||||
# FakeTensor subclass from a different mode.
|
||||
# Fall through to refakify.
|
||||
all_this_fake = False
|
||||
break
|
||||
case torch.Tensor():
|
||||
all_this_fake = False
|
||||
break
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
|
||||
if all_this_fake:
|
||||
return x
|
||||
|
||||
# see note [Tensor Fakification and Symbol Caching]
|
||||
symbolic_context = None
|
||||
source = None
|
||||
trace = True
|
||||
if tracing_context := torch._guards.TracingContext.try_get():
|
||||
if x in tracing_context.tensor_to_context:
|
||||
symbolic_context = tracing_context.tensor_to_context[x]
|
||||
source = symbolic_context.tensor_source
|
||||
# We already fakeified this tensor in Dynamo, don't
|
||||
# dump the trace for it again
|
||||
trace = False
|
||||
if (
|
||||
idx < aot_config.num_params_buffers
|
||||
and config.static_weight_shapes
|
||||
and not symbolic_context
|
||||
):
|
||||
# TODO: Ensure that this codepath is never exercised from
|
||||
# Dynamo
|
||||
return fake_mode.from_tensor(x, static_shapes=True)
|
||||
|
||||
result = fake_mode.from_tensor(
|
||||
x,
|
||||
static_shapes=ignore_shape_env,
|
||||
symbolic_context=symbolic_context,
|
||||
source=source,
|
||||
trace=trace,
|
||||
)
|
||||
return result
|
||||
|
||||
return FakifiedFlatArgs(
|
||||
[convert(idx, x) for idx, x in enumerate(flat_args)]
|
||||
), act_input_indices
|
||||
|
||||
|
||||
def construct_fake_mode(
|
||||
flat_args: list[Any], aot_config: AOTConfig
|
||||
) -> tuple[FakeTensorMode, ShapeEnv | None]:
|
||||
fake_mode = detect_fake_mode(flat_args)
|
||||
if fake_mode is None:
|
||||
shape_env = ShapeEnv() if aot_config.dynamic_shapes else None
|
||||
fake_mode = FakeTensorMode(shape_env=shape_env)
|
||||
else:
|
||||
shape_env = fake_mode.shape_env
|
||||
return (fake_mode, shape_env)
|
||||
|
||||
|
||||
def _try_get_metadata_from_dynamo(
|
||||
mod: torch.nn.Module,
|
||||
param_keys: KeysView[str],
|
||||
full_args_num: int,
|
||||
full_args_descs: list[DifferentiableAOTInput],
|
||||
) -> tuple[list[torch._guards.Source | None] | None, list[int]]:
|
||||
"""
|
||||
Metadata is forwarded from Dynamo to AOTDispatch via special fields on GraphModule.
|
||||
We first verify that `mod` does come from Dynamo, then we handle cases where
|
||||
metadata might be missing.
|
||||
|
||||
Returns:
|
||||
aot_autograd_arg_pos_to_source: used to dedup params and their guards
|
||||
static_input_indices: used to identify static inputs for cudagraphs
|
||||
"""
|
||||
# Note [Assumption on Dynamo Metadata]
|
||||
# This function assumes a graph module from dynamo provides `dynamo_compiled_id`,
|
||||
# _param_name_to_source, and every placeholder node has `_dynamo_source` attributes.
|
||||
# When gm is modified (e.g., DDPOptimizer via split_module), metadata needs to
|
||||
# be propagated in order to be recognized as a dynamo graph
|
||||
|
||||
if not (isinstance(mod, torch.fx.GraphModule) and "dynamo_compile_id" in mod.meta):
|
||||
# graph was not captured by dynamo
|
||||
return None, []
|
||||
|
||||
if not hasattr(mod, "_param_name_to_source"):
|
||||
# is from export
|
||||
static_input_indices = [
|
||||
i
|
||||
for i, node in enumerate(full_args_descs)
|
||||
if isinstance(node, (ParamAOTInput, BufferAOTInput))
|
||||
]
|
||||
return None, static_input_indices
|
||||
|
||||
# We now know this came from dynamo, and (1) we care about guards,
|
||||
# so setting up aot_autograd_arg_pos_to_source for downstream dedup guards
|
||||
# can now be done safely. (2) Dynamo logic protects the 1:1 sizing below.
|
||||
# Additionally, we mark static indices for cudagraphs.
|
||||
param_name_to_source = cast(
|
||||
dict[str, torch._guards.Source], mod._param_name_to_source
|
||||
)
|
||||
seen_sources = set()
|
||||
|
||||
aot_autograd_arg_pos_to_source: list[torch._guards.Source | None] = []
|
||||
static_input_indices = []
|
||||
# Collect the new inputs lifted by aotdispatch
|
||||
for i, name in enumerate(param_keys):
|
||||
if name not in param_name_to_source:
|
||||
raise AssertionError(f"{name} not found in param_name_to_source")
|
||||
source = param_name_to_source[name]
|
||||
if source in seen_sources:
|
||||
raise AssertionError(f"source {source} already in seen_sources")
|
||||
if source is None:
|
||||
raise AssertionError(f"source must not be None for {name}")
|
||||
seen_sources.add(source)
|
||||
aot_autograd_arg_pos_to_source.append(source)
|
||||
|
||||
static_input_indices.append(i)
|
||||
|
||||
# Collect the dynamo graph inputs
|
||||
# TODO(mlazos): Revisit if this is still needed. With Dynamo install ID
|
||||
# matched tensors back into the Fx graph, this might not be necessary.
|
||||
for pos, node in enumerate(mod.graph.find_nodes(op="placeholder")):
|
||||
if not hasattr(node, "_dynamo_source"):
|
||||
raise AssertionError(f"node {node} must have _dynamo_source attribute")
|
||||
source = node._dynamo_source
|
||||
# `source`` specifies the source from user code. ddp optimizer may have
|
||||
# intermediate values becoming submodule placeholders which does not
|
||||
# have a source
|
||||
if source is not None and source in seen_sources:
|
||||
raise AssertionError(f"source {source} already in seen_sources")
|
||||
seen_sources.add(source)
|
||||
aot_autograd_arg_pos_to_source.append(source)
|
||||
source_name = source.name if source else str(source)
|
||||
|
||||
# input[i] in dynamo is now:
|
||||
# input[i + len(extra_params)] in AOT,
|
||||
# where extra_params are the params/buffers that dynamo baked into the
|
||||
# OutputGraph
|
||||
actual_pos = pos + len(param_keys)
|
||||
|
||||
if "tensor_dict" in node.meta and node.meta["tensor_dict"].get(
|
||||
"_dynamo_static_input_type", None
|
||||
):
|
||||
static_inputs_log.debug(
|
||||
"Adding static input pos %s for source %s", actual_pos, source_name
|
||||
)
|
||||
static_input_indices.append(actual_pos)
|
||||
else:
|
||||
static_inputs_log.debug(
|
||||
"Non-static input pos %s for source %s", actual_pos, source_name
|
||||
)
|
||||
|
||||
if full_args_num != len(aot_autograd_arg_pos_to_source):
|
||||
raise AssertionError(
|
||||
f"full_args_num={full_args_num} != len(aot_autograd_arg_pos_to_source)={len(aot_autograd_arg_pos_to_source)}"
|
||||
)
|
||||
return aot_autograd_arg_pos_to_source, static_input_indices
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _detect_attribute_assignment(mod: torch.nn.Module) -> Generator[None, None, None]:
|
||||
# Do not allow assignment of tensor attributes during export unless
|
||||
# the attribute is registered as a buffer.
|
||||
|
||||
NN_MODULE_STD_ATTRS = [
|
||||
"_backward_hooks",
|
||||
"_backward_pre_hooks",
|
||||
"_buffers",
|
||||
"_forward_hooks",
|
||||
"_forward_hooks_always_called",
|
||||
"_forward_hooks_with_kwargs",
|
||||
"_forward_pre_hooks",
|
||||
"_forward_pre_hooks_with_kwargs",
|
||||
"_is_full_backward_hook",
|
||||
"_load_state_dict_post_hooks",
|
||||
"_load_state_dict_pre_hooks",
|
||||
"_modules",
|
||||
"_non_persistent_buffers_set",
|
||||
"_parameters",
|
||||
"_state_dict_hooks",
|
||||
"_state_dict_pre_hooks",
|
||||
"training",
|
||||
]
|
||||
NN_MODULE_LAZY_STD_ATTRS = [
|
||||
"_initialize_hook",
|
||||
"_load_hook",
|
||||
]
|
||||
STD_ATTRS = {
|
||||
*NN_MODULE_STD_ATTRS,
|
||||
*NN_MODULE_LAZY_STD_ATTRS,
|
||||
}
|
||||
|
||||
def _get_attributes(mod: torch.nn.Module) -> dict[str, Any]:
|
||||
# return any attributes of a module that are not standard attributes
|
||||
return {k: v for k, v in mod.__dict__.items() if k not in STD_ATTRS}
|
||||
|
||||
def _get_all_module_attributes(mod: torch.nn.Module) -> dict[str, dict[str, Any]]:
|
||||
# return attributes from all modules and submodules
|
||||
result = {}
|
||||
for name, submodule in mod.named_modules():
|
||||
result[name] = _get_attributes(submodule)
|
||||
return result
|
||||
|
||||
def _restore_all_module_attributes(
|
||||
mod: torch.nn.Module, snapshot: dict[str, dict[str, Any]]
|
||||
) -> None:
|
||||
# restore attributes to all modules and submodules
|
||||
for name, submodule in mod.named_modules():
|
||||
if name in snapshot:
|
||||
submodule.__dict__.update(snapshot[name])
|
||||
|
||||
# save state of attributes before enter
|
||||
snapshot = pytree.tree_map(
|
||||
lambda x: x,
|
||||
_get_all_module_attributes(mod),
|
||||
is_leaf=lambda x: type(x) in _pytree_subclasses_that_lose_info,
|
||||
)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# after exit, compare state of attributes with snapshot
|
||||
# to detect which tensor attributes were assigned
|
||||
|
||||
def _collect_assigned_tensor_attributes(
|
||||
snapshot: dict[str, dict[str, Any]], new_attrs: dict[str, dict[str, Any]]
|
||||
) -> list[str]:
|
||||
assigned_tensor_attributes = []
|
||||
|
||||
def _compare_values(path: str, old_val: Any, new_val: Any) -> None:
|
||||
"""Recursively compare values, handling containers."""
|
||||
# Same object, no change
|
||||
if old_val is new_val:
|
||||
return
|
||||
|
||||
if old_val is None or new_val is None:
|
||||
if isinstance(new_val, torch.Tensor):
|
||||
assigned_tensor_attributes.append(path)
|
||||
return
|
||||
|
||||
# Check if it's a tensor that was reassigned
|
||||
if isinstance(new_val, torch.Tensor):
|
||||
assigned_tensor_attributes.append(path)
|
||||
return
|
||||
|
||||
# Handle dict containers
|
||||
if isinstance(old_val, dict) and isinstance(new_val, dict):
|
||||
all_keys = set(old_val.keys()) | set(new_val.keys())
|
||||
for key in all_keys:
|
||||
old_item = old_val.get(key)
|
||||
new_item = new_val.get(key)
|
||||
_compare_values(f"{path}[{key!r}]", old_item, new_item)
|
||||
return
|
||||
|
||||
# Handle list/tuple containers
|
||||
if isinstance(old_val, (list, tuple)) and isinstance(
|
||||
new_val, (list, tuple)
|
||||
):
|
||||
# Different lengths = mutation happened
|
||||
max_len = max(len(old_val), len(new_val))
|
||||
for i in range(max_len):
|
||||
old_item = old_val[i] if i < len(old_val) else None
|
||||
new_item = new_val[i] if i < len(new_val) else None
|
||||
_compare_values(f"{path}[{i}]", old_item, new_item)
|
||||
return
|
||||
|
||||
# For other types, just check if they're different objects
|
||||
# (we don't care about non-tensor mutations)
|
||||
|
||||
for module_name in snapshot.keys() | new_attrs.keys():
|
||||
old_module_attrs = snapshot.get(module_name, {})
|
||||
new_module_attrs = new_attrs.get(module_name, {})
|
||||
|
||||
for attr_name in old_module_attrs.keys() | new_module_attrs.keys():
|
||||
module_prefix = f"self.{module_name}." if module_name else "self."
|
||||
full_path = f"{module_prefix}{attr_name}"
|
||||
|
||||
old_val = old_module_attrs.get(attr_name)
|
||||
new_val = new_module_attrs.get(attr_name)
|
||||
_compare_values(full_path, old_val, new_val)
|
||||
|
||||
return assigned_tensor_attributes
|
||||
|
||||
new_attrs = _get_all_module_attributes(mod)
|
||||
assigned_tensor_attributes = _collect_assigned_tensor_attributes(
|
||||
snapshot, new_attrs
|
||||
)
|
||||
# restore state of all attributes (including, e.g., of primitive types)
|
||||
_restore_all_module_attributes(mod, snapshot)
|
||||
|
||||
if assigned_tensor_attributes:
|
||||
if len(assigned_tensor_attributes) > 1:
|
||||
noun, verb = "attributes", "were"
|
||||
else:
|
||||
noun, verb = "attribute", "was"
|
||||
warnings.warn(
|
||||
f"The tensor {noun} {', '.join(assigned_tensor_attributes)} {verb} assigned during export. "
|
||||
"Such attributes must be registered as buffers using the `register_buffer` API "
|
||||
"(https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer).",
|
||||
stacklevel=2,
|
||||
)
|
||||
+658
@@ -0,0 +1,658 @@
|
||||
"""
|
||||
This file contains utilities related to functionalization in AOTAutograd:
|
||||
1. converting to/from functional tensors
|
||||
2. detecting Tensor mutations - both metadata and Tensor value
|
||||
3. regenerating/replaying views from their base
|
||||
4. checking if a graph is functional i.e. whether it contains any mutation ops
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TypeGuard
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._C import _functionalization
|
||||
from torch._logging import getArtifactLogger
|
||||
from torch._opaque_base import OpaqueBase
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
from torch._subclasses.meta_utils import is_sparse_any
|
||||
from torch.fx.experimental.symbolic_shapes import guard_or_false, sym_eq, SymIntEqByExpr
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils._python_dispatch import (
|
||||
is_traceable_wrapper_subclass,
|
||||
transform_subclass,
|
||||
)
|
||||
|
||||
|
||||
aot_joint_log = getArtifactLogger(__name__, "aot_joint_graph")
|
||||
|
||||
|
||||
def to_fun(t: object) -> Any:
|
||||
if isinstance(t, Tensor):
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
# See Note [Functionalization always runs last]
|
||||
# This means that if we want to "functionalize" a subclass, we need to ensure that the functional wrapper
|
||||
# goes at the bottom.
|
||||
# recurse here, so we can support nested wrapper subclasses
|
||||
out = transform_subclass(t, lambda _, inner_t: to_fun(inner_t))
|
||||
torch._mirror_autograd_meta_to(t, out) # type: ignore[attr-defined]
|
||||
return out
|
||||
else:
|
||||
return FunctionalTensor.to_functional(t)
|
||||
else:
|
||||
return t
|
||||
|
||||
|
||||
def sync_functional_tensor(t: torch.Tensor) -> None:
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
attrs, _ctx = t.__tensor_flatten__() # type: ignore[attr-defined]
|
||||
for attr in attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as inner:
|
||||
sync_functional_tensor(inner)
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
else:
|
||||
torch._sync(t)
|
||||
|
||||
|
||||
# When subclasses are involved, t here will usually look something like:
|
||||
# SubclassA(SubclassB(FunctionalTensor(_to_fun_tensor(FakeTensor))))
|
||||
def from_fun(t: object) -> object:
|
||||
if isinstance(t, Tensor) and is_traceable_wrapper_subclass(t):
|
||||
# See Note [Functionalization always runs last]
|
||||
# This means that if we want to "functionalize" a subclass, we need to ensure that the functional wrapper
|
||||
# goes at the bottom.
|
||||
# recurse here, so we can support nested wrapper subclasses
|
||||
out = transform_subclass(t, lambda _, inner_t: from_fun(inner_t))
|
||||
torch._mirror_autograd_meta_to(t, out) # type: ignore[attr-defined]
|
||||
return out
|
||||
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
# quick sanity assert
|
||||
if isinstance(t, torch.Tensor):
|
||||
if torch._is_functional_tensor(t): # type: ignore[attr-defined]
|
||||
raise AssertionError("expected non-functional tensor")
|
||||
return t
|
||||
sync_functional_tensor(t)
|
||||
return torch._from_functional_tensor(t.elem)
|
||||
|
||||
|
||||
def is_fun(t: object) -> TypeGuard[FunctionalTensor | Tensor]:
|
||||
if isinstance(t, Tensor) and is_traceable_wrapper_subclass(t):
|
||||
# See Note [Functionalization always runs last]
|
||||
# This means that if we want to "functionalize" a subclass, we need to ensure that the functional wrapper
|
||||
# goes at the bottom.
|
||||
# recurse here, so we can support nested wrapper subclasses
|
||||
t_attrs, _ = t.__tensor_flatten__() # type: ignore[attr-defined]
|
||||
got_fun: bool | None = None
|
||||
for attr in t_attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as v:
|
||||
fun = is_fun(v)
|
||||
if got_fun is None:
|
||||
got_fun = fun
|
||||
elif got_fun != fun:
|
||||
raise AssertionError(
|
||||
"mixed functional/non-functional inner tensors"
|
||||
)
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return got_fun or False
|
||||
|
||||
return isinstance(t, FunctionalTensor)
|
||||
|
||||
|
||||
# t here is either
|
||||
# (1) A FunctionalTensor(_to_functional_tensor(FakeTensor))
|
||||
# (2) A traceable tensor subclass that holds a FunctionalTensor
|
||||
# (3) Not a tensor
|
||||
def has_data_mutation(t: object) -> bool:
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
attrs, _ = t.__tensor_flatten__()
|
||||
# A tensor subclass was updated if any of its inner elements were updated
|
||||
for attr in attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as v:
|
||||
if has_data_mutation(v):
|
||||
return True
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return False
|
||||
else:
|
||||
if isinstance(t, torch.Tensor):
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
|
||||
return torch._functionalize_has_data_mutation(t.elem) # type: ignore[attr-defined]
|
||||
return False
|
||||
|
||||
|
||||
def are_all_mutations_hidden_from_autograd(t: object) -> bool:
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
attrs, _ = t.__tensor_flatten__()
|
||||
# If all inner elements are mutations hidden from autograd, then it is a mutation hidden from autograd.
|
||||
for attr in attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as v:
|
||||
if not are_all_mutations_hidden_from_autograd(v):
|
||||
return False
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return True
|
||||
elif isinstance(t, torch.Tensor):
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
|
||||
return torch._functionalize_are_all_mutations_hidden_from_autograd(t.elem)
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def are_all_mutations_under_no_grad_or_inference_mode(t: torch.Tensor) -> bool:
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
attrs, _ = t.__tensor_flatten__()
|
||||
for attr in attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as v:
|
||||
if not are_all_mutations_under_no_grad_or_inference_mode(v):
|
||||
return False
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return True
|
||||
else:
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
|
||||
return torch._functionalize_are_all_mutations_under_no_grad_or_inference_mode(
|
||||
t.elem
|
||||
)
|
||||
|
||||
|
||||
def was_inductor_storage_resized(t: object) -> bool:
|
||||
if is_traceable_wrapper_subclass(t):
|
||||
attrs, _ = t.__tensor_flatten__()
|
||||
for attr in attrs:
|
||||
match getattr(t, attr):
|
||||
case Tensor() as v:
|
||||
if was_inductor_storage_resized(v):
|
||||
raise RuntimeError(
|
||||
f"storage resizing is not supported on tensor subclass: {type(t)}"
|
||||
)
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return False
|
||||
elif not isinstance(t, torch.Tensor):
|
||||
return False
|
||||
else:
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
raise AssertionError(f"expected FunctionalTensor, got {type(t)}")
|
||||
return torch._functionalize_was_inductor_storage_resized(t.elem)
|
||||
|
||||
|
||||
# f_arg here is either
|
||||
# (1) A FunctionalTensor(_to_functional_tensor(FakeTensor))
|
||||
# (2) A traceable tensor subclass that holds a FunctionalTensor
|
||||
# (3) Not a tensor
|
||||
# Assumption: arg promises to be the "original" tensor wrapped by f_arg
|
||||
# Note: "storage mutations" coming from set_() are a type of metadata mutation. So:
|
||||
# - check_only_storage_mutation=True: only return true if there was a storage mutation
|
||||
# - check_only_storage_mutation=Flse: return true if there was any metadata mutation (including a storage mutation)
|
||||
def has_metadata_mutation(
|
||||
f_arg: object, arg: object, *, check_only_storage_mutation: bool
|
||||
) -> bool:
|
||||
if is_traceable_wrapper_subclass(f_arg):
|
||||
attrs, _ = f_arg.__tensor_flatten__()
|
||||
# A tensor subclass was updated if any of its inner elements were updated
|
||||
for attr in attrs:
|
||||
match getattr(f_arg, attr):
|
||||
case Tensor():
|
||||
f_inner_t = getattr(f_arg, attr)
|
||||
inner_t = getattr(arg, attr)
|
||||
if has_metadata_mutation(
|
||||
f_inner_t,
|
||||
inner_t,
|
||||
check_only_storage_mutation=check_only_storage_mutation,
|
||||
):
|
||||
return True
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return False
|
||||
else:
|
||||
if not isinstance(f_arg, torch.Tensor):
|
||||
if isinstance(arg, torch.Tensor):
|
||||
raise AssertionError(
|
||||
f"f_arg is not a Tensor but arg is: {type(f_arg)} vs {type(arg)}"
|
||||
)
|
||||
return False
|
||||
if not isinstance(f_arg, FunctionalTensor):
|
||||
raise AssertionError(
|
||||
f"expected FunctionalTensor for f_arg, got {type(f_arg)}"
|
||||
)
|
||||
if not isinstance(arg, FakeTensor):
|
||||
raise AssertionError(f"expected FakeTensor for arg, got {type(arg)}")
|
||||
|
||||
arg_after = torch._from_functional_tensor(f_arg.elem)
|
||||
# This is true if the current tensor experienced at least one set_() call
|
||||
maybe_storage_changed = torch._functionalize_was_storage_changed(f_arg.elem) # type: ignore[attr-defined]
|
||||
# However, multiple set_() calls can cancel out. So we also check whether the
|
||||
# storage of the tensor has changed.
|
||||
# Note: if an input experienced two set_() calls that cancel out, **and**
|
||||
# it experiences an data mutation, we pessimistically think that the set_()
|
||||
# call is necessary here. We could in theory fix this, but this will
|
||||
# hopefully never happen in user code, and is not needed for fsdp.
|
||||
if is_sparse_any(arg):
|
||||
# TODO:add sparse tensors support to functionalization
|
||||
same_storages = False
|
||||
else:
|
||||
same_storages = StorageWeakRef(arg.untyped_storage()) == StorageWeakRef(
|
||||
arg_after.untyped_storage()
|
||||
)
|
||||
has_storage_metadata_mutation = maybe_storage_changed and not same_storages
|
||||
if check_only_storage_mutation:
|
||||
return has_storage_metadata_mutation
|
||||
|
||||
# storage metadata mutation is a type of metadata mutation, so return true if we saw one
|
||||
if has_storage_metadata_mutation:
|
||||
return True
|
||||
|
||||
maybe_metadata_mutated = torch._functionalize_has_metadata_mutation(f_arg.elem) # type: ignore[attr-defined]
|
||||
# This is true if the current tensor experienced at least one metadata mutation.
|
||||
# So if false, we know there was no metadata mutation
|
||||
if not maybe_metadata_mutated:
|
||||
return False
|
||||
|
||||
# However, multi metadata mutations can cancel out.
|
||||
# So we also check if the concrete sizes/strides on the tensor have changed.
|
||||
same_sizes = arg.shape == arg_after.shape
|
||||
same_strides = arg.stride() == arg_after.stride()
|
||||
same_offsets = arg.storage_offset() == arg_after.storage_offset()
|
||||
has_metadata_mutation_ = maybe_metadata_mutated and not (
|
||||
same_sizes and same_strides and same_offsets
|
||||
)
|
||||
# We consider a tensor to have been metadata mutated if its storage was mutated through a set_() call.
|
||||
return has_metadata_mutation_
|
||||
|
||||
|
||||
def gen_alias_from_base(
|
||||
aliased_base_tensor: Tensor,
|
||||
target_meta_tensor: Tensor,
|
||||
target_requires_grad: bool,
|
||||
target_view_meta_sequence: ViewMetaSequence | None = None,
|
||||
*,
|
||||
replay_views: bool,
|
||||
) -> Tensor:
|
||||
# Patch the correct requires_grad field of the output tensor, depending on whether:
|
||||
# (i) the reconstructed output (out) was came from a tensor that requires grad or not;
|
||||
# and (ii) the concrete returned output does require grad or not.
|
||||
def patch_requires_grad(out: Tensor) -> Tensor:
|
||||
if aliased_base_tensor.requires_grad and not target_requires_grad:
|
||||
out = out.detach()
|
||||
elif not aliased_base_tensor.requires_grad and target_requires_grad:
|
||||
out.requires_grad_(True)
|
||||
return out
|
||||
|
||||
# If provided, use the target functional tensor for replaying the views.
|
||||
#
|
||||
# In summary, we use the fact that FunctionalTensorWrapper saves the view
|
||||
# functions applied to itself (collected during functionalization) so as
|
||||
# to replay them (view functions) on the aliased_base_tensor.
|
||||
if (
|
||||
replay_views
|
||||
and target_view_meta_sequence is not None
|
||||
and not any(vm.has_symbolic_inputs for vm in target_view_meta_sequence.sequence)
|
||||
):
|
||||
out = _functionalization.apply_view_meta_sequence(
|
||||
aliased_base_tensor, target_view_meta_sequence.sequence
|
||||
)
|
||||
# If re-applying the ViewMeta sequence succeeded, there should be no more
|
||||
# problems going forward. We just check we got to the target shape and
|
||||
# patch requires_grad flag.
|
||||
if out.shape != target_meta_tensor.shape:
|
||||
raise AssertionError(
|
||||
"incorrect out shape after application of ViewMeta sequence: "
|
||||
f"{tuple(out.shape)} (actual) vs {tuple(target_meta_tensor.shape)} (expected)"
|
||||
)
|
||||
return patch_requires_grad(out)
|
||||
|
||||
# Try to do view-replay if possible.
|
||||
# fall back to .as_strided() if we can't.
|
||||
if target_meta_tensor._base is not None:
|
||||
# The base that we want to replay our view off of might have a different shape than the view's original base.
|
||||
b = target_meta_tensor._base
|
||||
abt = aliased_base_tensor
|
||||
# Don't unnecessarily call as_strided if nothing changed; as_strided's
|
||||
# backward is poorly implemented and slow
|
||||
if abt is not b and (
|
||||
abt.size() != b.size()
|
||||
or abt.stride() != b.stride()
|
||||
or abt.storage_offset() != b.storage_offset()
|
||||
):
|
||||
reshaped_base_tensor = aliased_base_tensor.as_strided(
|
||||
b.size(), b.stride(), b.storage_offset()
|
||||
)
|
||||
else:
|
||||
reshaped_base_tensor = aliased_base_tensor
|
||||
out = target_meta_tensor._view_func(reshaped_base_tensor) # type: ignore[attr-defined]
|
||||
# This shape mismatch can happen due to a bug in inplace/view handling in autograd.
|
||||
# Try putting a breakpoint here and running
|
||||
# `test/functorch/test_aotdispatch TestAOTAutograd.test_output_all_alias_types`
|
||||
# Also, https://github.com/pytorch/pytorch/issues/49825
|
||||
#
|
||||
# As a stopgap, we'll fall back to as_strided.
|
||||
if out is not None and out.shape == target_meta_tensor.shape:
|
||||
return patch_requires_grad(out)
|
||||
|
||||
size = target_meta_tensor.size()
|
||||
stride = target_meta_tensor.stride()
|
||||
storage_offset = target_meta_tensor.storage_offset()
|
||||
if aliased_base_tensor.is_complex() and not target_meta_tensor.is_complex():
|
||||
aliased_out = torch.view_as_real(aliased_base_tensor).as_strided(
|
||||
size, stride, storage_offset
|
||||
)
|
||||
elif not aliased_base_tensor.is_complex() and target_meta_tensor.is_complex():
|
||||
aliased_out = torch.view_as_complex(aliased_base_tensor).as_strided(
|
||||
size, stride, storage_offset
|
||||
)
|
||||
else:
|
||||
aliased_out = aliased_base_tensor.as_strided(size, stride, storage_offset)
|
||||
# For outputs aliasing inputs, we need to check if the requires-gradness has changed.
|
||||
aliased_out = patch_requires_grad(aliased_out)
|
||||
# For outputs aliasing inputs, we need to check if the dtype has changed.
|
||||
# as_strided() is the "most generic" view, but it does not cover cross-dtype views
|
||||
if aliased_out.dtype != target_meta_tensor.dtype:
|
||||
aliased_out = aliased_out.view(target_meta_tensor.dtype)
|
||||
return aliased_out
|
||||
|
||||
|
||||
def has_same_metadata(t1: Tensor, t2: Tensor) -> bool:
|
||||
return (
|
||||
guard_or_false(sym_eq(t1.size(), t2.size()))
|
||||
and guard_or_false(t1.layout == t2.layout)
|
||||
and (
|
||||
is_sparse_any(t1)
|
||||
or (
|
||||
guard_or_false(sym_eq(t1.stride(), t2.stride()))
|
||||
and guard_or_false(t1.storage_offset() == t2.storage_offset())
|
||||
)
|
||||
)
|
||||
and t1.is_conj() == t2.is_conj()
|
||||
and t1.is_neg() == t2.is_neg()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetadataKey:
|
||||
"""
|
||||
This should be equal whenever has_same_metadata would return True
|
||||
"""
|
||||
|
||||
size: tuple[SymIntEqByExpr, ...]
|
||||
layout: torch.layout
|
||||
is_sparse: bool
|
||||
# these are empty when is_sparse
|
||||
stride: tuple[SymIntEqByExpr, ...] | None
|
||||
storage_offset: SymIntEqByExpr | None
|
||||
is_conj: bool
|
||||
is_neg: bool
|
||||
|
||||
@staticmethod
|
||||
def make(t: Tensor) -> MetadataKey:
|
||||
is_sparse = is_sparse_any(t)
|
||||
return MetadataKey(
|
||||
size=tuple(SymIntEqByExpr(s) for s in t.size()),
|
||||
layout=t.layout,
|
||||
is_sparse=is_sparse,
|
||||
stride=None if is_sparse else tuple(SymIntEqByExpr(s) for s in t.stride()),
|
||||
storage_offset=None if is_sparse else SymIntEqByExpr(t.storage_offset()),
|
||||
is_conj=t.is_conj(),
|
||||
is_neg=t.is_neg(),
|
||||
)
|
||||
|
||||
|
||||
# ViewMeta sequence wrapper for equality comparisons.
|
||||
#
|
||||
# Even though we can compare each ViewMeta instance, we compare the resulting
|
||||
# tensor metadata, instead. That's because the creation of synthetic bases + the
|
||||
# re-generation of input views might end-up creating a different sequence of
|
||||
# ViewMeta that is semantically equivalent. i.e. gets to a tensor with the same
|
||||
# metadata.
|
||||
#
|
||||
# Therefore, we store what the end result should look like as serializable
|
||||
# metadata.
|
||||
#
|
||||
# When logging, this class should look like:
|
||||
#
|
||||
# ViewMetaSequence(view, select_int, slice_Tensor)
|
||||
#
|
||||
# i.e. a parenthesized list of view operations within that ViewMeta sequence.
|
||||
class ViewMetaSequence:
|
||||
def __init__(self, tensor: FunctionalTensor) -> None:
|
||||
if not torch._is_functional_tensor(tensor.elem):
|
||||
raise AssertionError("expected tensor.elem to be a functional tensor")
|
||||
self.sequence = _functionalization.get_view_meta_sequence(tensor.elem)
|
||||
self.metadata = MetadataKey.make(tensor)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
suffix = len("_ViewMeta")
|
||||
types = ", ".join(type(vm).__name__[:-suffix] for vm in self.sequence)
|
||||
return f"ViewMetaSequence({types})"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
# If other is None, then it probably means that we weren't able to recreate
|
||||
# the ViewMeta sequence. One example is when we update the view metadata by
|
||||
# calling: create_synthetic_base_metadata.
|
||||
if other is None:
|
||||
return True
|
||||
|
||||
# Comparison against any other type is not implemented.
|
||||
if not isinstance(other, ViewMetaSequence):
|
||||
return NotImplemented
|
||||
|
||||
return self.metadata == other.metadata
|
||||
|
||||
|
||||
# new_arg and arg here are either:
|
||||
# (1) both a FakeTensor
|
||||
# (2) both a traceable tensor subclass that holds a FakeTensor
|
||||
# Pre-condition: the two args are the "old" and "new" inputs from running functionalization.
|
||||
# When we run functionalization and wrap our inputs into FunctionalTensors,
|
||||
# we can detect whether or not an input was mutated by checking to see if the inner tensor has changed
|
||||
#
|
||||
# Normally it would be enough just to check if arg is new_arg, which is normally enough for functionalization
|
||||
# to confirm that inputs were not mutated when running the user's model with functionalization on.
|
||||
# But when we have subclass inputs, we can't rely on that:
|
||||
# `from_fun(to_fun(x)) is x` will return False, because the call to `from_fun` constructs
|
||||
# a brand new subclass instance: we are calling __tensor_unflatten__, and going
|
||||
# from Subclass(FakeTensor) to Subclass(FunctionalTensor(FakeTensor))
|
||||
def was_tensor_updated(arg: torch.Tensor, new_arg: torch.Tensor) -> bool:
|
||||
if is_traceable_wrapper_subclass(arg):
|
||||
if not is_traceable_wrapper_subclass(new_arg):
|
||||
raise AssertionError(
|
||||
f"expected new_arg to be traceable wrapper subclass, got {type(new_arg)}"
|
||||
)
|
||||
attrs, _ = arg.__tensor_flatten__()
|
||||
new_attrs, _ = new_arg.__tensor_flatten__()
|
||||
if attrs != new_attrs:
|
||||
raise AssertionError(f"attrs mismatch: {attrs} != {new_attrs}")
|
||||
# A tensor subclass was updated if any of its inner elements were updated
|
||||
for attr in attrs:
|
||||
match getattr(arg, attr):
|
||||
case Tensor() as v:
|
||||
if was_tensor_updated(v, getattr(new_arg, attr)):
|
||||
return True
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return arg is not new_arg
|
||||
|
||||
|
||||
# new_arg and arg here are either:
|
||||
# (1) both a FakeTensor
|
||||
# (2) both a traceable tensor subclass that holds a FakeTensor
|
||||
# Pre-condition: the two args are the "old" and "new" inputs from running functionalization.
|
||||
# When we run functionalization and wrap our inputs into FunctionalTensors,
|
||||
# we can detect whether or not an input was mutated by checking to see if the inner tensor has changed,
|
||||
# but shares storage with the old input
|
||||
def was_tensor_metadata_updated(arg: Any, new_arg: Any) -> bool:
|
||||
if is_traceable_wrapper_subclass(arg):
|
||||
if not is_traceable_wrapper_subclass(new_arg):
|
||||
raise AssertionError(
|
||||
f"expected new_arg to be traceable wrapper subclass, got {type(new_arg)}"
|
||||
)
|
||||
attrs, _ = arg.__tensor_flatten__()
|
||||
new_attrs, _ = new_arg.__tensor_flatten__()
|
||||
if attrs != new_attrs:
|
||||
raise AssertionError(f"attrs mismatch: {attrs} != {new_attrs}")
|
||||
# A tensor subclass was updated if any of its inner elements were updated
|
||||
for attr in attrs:
|
||||
match getattr(arg, attr):
|
||||
case Tensor() as v:
|
||||
if was_tensor_metadata_updated(v, getattr(new_arg, attr)):
|
||||
return True
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return arg is not new_arg and StorageWeakRef(
|
||||
arg.untyped_storage()
|
||||
) == StorageWeakRef(new_arg.untyped_storage())
|
||||
|
||||
|
||||
# Returns the number of detected copy_
|
||||
def _is_functional_graph(fx_g: torch.fx.Graph) -> tuple[str | None, int]:
|
||||
allowed_mutation_ops = [
|
||||
torch.ops.aten.copy_.default,
|
||||
torch.ops.aten.set_.source_Tensor,
|
||||
]
|
||||
if hasattr(torch.ops.fsdp, "copy_"):
|
||||
allowed_mutation_ops.append(torch.ops.fsdp.copy_.default)
|
||||
|
||||
placeholders = set()
|
||||
mutation_count = 0
|
||||
# NB: It would also be nice to verify that the mutations all happen at the
|
||||
# end, but we also do some administrative views after mutations so this
|
||||
# isn't actually true. (TODO: Could this cause problems for Inductor?)
|
||||
error = None
|
||||
for n in fx_g.nodes:
|
||||
if n.op == "placeholder":
|
||||
placeholders.add(n)
|
||||
if isinstance(n.target, torch._ops.OpOverload):
|
||||
if n.target in allowed_mutation_ops:
|
||||
# Can only copy_/set_ into an input
|
||||
# this is mostly a hack to avoid failing XLA tests.
|
||||
# See https://github.com/pytorch/pytorch/pull/122434#issuecomment-2101012113
|
||||
if "set_buffer_donor_" not in str(n.args[0]):
|
||||
if n.args[0] not in placeholders:
|
||||
error = f"n={str(n)}, n.args[0]={str(n.args[0])}, placeholders={str(placeholders)}, graph={str(fx_g)}"
|
||||
mutation_count += 1
|
||||
else:
|
||||
if n.target._schema.is_mutable:
|
||||
error = f"aot_autograd expected to have an entirely functional graph, but found {n.format_node()}"
|
||||
return error, mutation_count
|
||||
|
||||
|
||||
def assert_functional_graph(fx_g: torch.fx.Graph) -> int:
|
||||
error, mutation_count = _is_functional_graph(fx_g)
|
||||
if error is not None:
|
||||
raise AssertionError(error)
|
||||
return mutation_count
|
||||
|
||||
|
||||
def propagate_input_mutation_stacktraces(fx_g: torch.fx.Graph) -> None:
|
||||
placeholders = set()
|
||||
for n in fx_g.nodes:
|
||||
if n.op == "placeholder":
|
||||
placeholders.add(n)
|
||||
if isinstance(n.target, torch._ops.OpOverload):
|
||||
if n.target is torch.ops.aten.copy_.default:
|
||||
# Can only copy_ into an input, and can only do so once
|
||||
if "set_buffer_donor_" not in str(n.args[0]):
|
||||
if n.args[0] not in placeholders:
|
||||
raise AssertionError(
|
||||
f"n={str(n)}, n.args[0]={str(n.args[0])}, placeholders={str(placeholders)}, graph={str(fx_g)}"
|
||||
)
|
||||
placeholders.remove(n.args[0])
|
||||
copy_from_node = n.args[1]
|
||||
# Pre-condition: every node has a "stack_trace" field in its meta,
|
||||
# but copy_() nodes do not (since we manually added them during functionalization).
|
||||
# Instead, we manually propagate here.
|
||||
if "stack_trace" in copy_from_node.meta:
|
||||
n.meta["stack_trace"] = copy_from_node.meta["stack_trace"]
|
||||
|
||||
|
||||
def _check_if_mutation_can_be_in_graph(
|
||||
keep_input_mutations: bool,
|
||||
mutates_data: bool,
|
||||
mutates_metadata: bool,
|
||||
mutations_hidden_from_autograd: bool,
|
||||
mutations_under_no_grad_or_inference_mode: bool,
|
||||
mutates_storage_metadata: bool,
|
||||
mutation_inductor_storage_resize: bool,
|
||||
requires_grad: bool,
|
||||
) -> bool:
|
||||
if keep_input_mutations:
|
||||
in_graph = (
|
||||
mutates_data or mutates_storage_metadata or mutation_inductor_storage_resize
|
||||
) and (
|
||||
(not mutates_metadata and not requires_grad)
|
||||
or mutations_hidden_from_autograd
|
||||
or mutations_under_no_grad_or_inference_mode
|
||||
)
|
||||
else:
|
||||
in_graph = False
|
||||
# See Note [set_() Input Mutations in AOTAutograd]
|
||||
# If there was a `set_()`, we require that all mutations were under no_grad,
|
||||
# so we can (safely) emit the set_() in the graph at runtime
|
||||
# resize_() gets the same treatment
|
||||
if mutation_inductor_storage_resize or mutates_storage_metadata:
|
||||
op_name = "resize_" if mutation_inductor_storage_resize else "set_"
|
||||
if not in_graph:
|
||||
raise AssertionError(f"""\
|
||||
Encountered a {op_name} on a graph input, but the input has other mutations that we cannot
|
||||
keep in the graph. This is not supported today. Current state:
|
||||
keep_input_mutations={keep_input_mutations}
|
||||
mutates_data={mutates_data}
|
||||
mutates_metadata={mutates_metadata}
|
||||
mutations_hidden_from_autograd={mutations_hidden_from_autograd}
|
||||
mutations_under_no_grad_or_inference_mode={mutations_under_no_grad_or_inference_mode}
|
||||
mutation_inductor_storage_resize={mutation_inductor_storage_resize}
|
||||
requires_grad={requires_grad}""")
|
||||
return in_graph
|
||||
@@ -0,0 +1,324 @@
|
||||
"""
|
||||
This module contains utility functions for working with joint FX graphs with descriptors
|
||||
that are produced by AOTAutograd. They will NOT work on generic FX graphs. See also
|
||||
:func:`torch._functorch.aot_autograd.aot_export_joint_with_descriptors`. We also
|
||||
recommend reading :mod:torch._functorch._aot_autograd.descriptors`.
|
||||
"""
|
||||
|
||||
from typing import NoReturn
|
||||
|
||||
import torch.fx as fx
|
||||
|
||||
from .descriptors import (
|
||||
AOTInput,
|
||||
AOTOutput,
|
||||
BufferAOTInput,
|
||||
DifferentiableAOTInput,
|
||||
DifferentiableAOTOutput,
|
||||
GradAOTOutput,
|
||||
ParamAOTInput,
|
||||
PlainAOTInput,
|
||||
PlainAOTOutput,
|
||||
SubclassGetAttrAOTInput,
|
||||
SubclassGetAttrAOTOutput,
|
||||
TangentAOTInput,
|
||||
)
|
||||
|
||||
|
||||
def _raise_autograd_subclass_not_implemented(
|
||||
n: fx.Node, desc: AOTInput | AOTOutput
|
||||
) -> NoReturn:
|
||||
raise RuntimeError(
|
||||
"Subclasses are currently not supported by this function, but a desugared subclass input "
|
||||
f"was found at {n} ({desc}). The problem is "
|
||||
"that there may not necessarily be a 1-1 correspondence between primals/tangents/outputs/grads "
|
||||
"when subclasses are involved: for example, the primal might be a plain tensor "
|
||||
"but the tangent a tensor subclass that desugared into multiple plain tensors. "
|
||||
"It is not clear what exactly you would like this function to do in this case "
|
||||
"(Collect all nodes for the subclass together? Match up the inner nodes if "
|
||||
"subclasses match exactly?) If you have a concrete use case, please file an "
|
||||
"issue so we can understand it and design an API that works for your case."
|
||||
)
|
||||
|
||||
|
||||
def get_all_input_and_grad_nodes(
|
||||
g: fx.Graph,
|
||||
) -> dict[DifferentiableAOTInput, tuple[fx.Node, fx.Node | None]]:
|
||||
"""
|
||||
Given a joint graph with descriptors (meta['desc'] on placeholders and
|
||||
output), returns the node for every input and its corresponding grad
|
||||
output node if it exists. These tuples are in a dict that is indexed by
|
||||
the AOTInput descriptor that describes the input.
|
||||
|
||||
NB: *all* forward tensor inputs are returned, including non-differentiable
|
||||
inputs (which simply have a None grad), so it is safe to use this function
|
||||
to perform operations on all inputs. (Non-tensor inputs like symbolic
|
||||
integers, tokens or RNG state are NOT traversed by this function.)
|
||||
|
||||
Args:
|
||||
g: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each DifferentiableAOTInput descriptor to a tuple
|
||||
containing:
|
||||
- The input node itself
|
||||
- The grad (output) node if it exists, None otherwise
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the joint graph has subclass tensor inputs/outputs; this
|
||||
is not supported by API as there is not necessarily a 1-1 correspondence
|
||||
between inputs and grads when subclasses are involved.
|
||||
"""
|
||||
input_index: dict[DifferentiableAOTInput, tuple[fx.Node, fx.Node | None]] = {}
|
||||
for n in g.nodes:
|
||||
if n.op == "placeholder":
|
||||
desc = n.meta["desc"]
|
||||
# Skip inputs that cannot possibly be differentiable
|
||||
if not isinstance(desc, DifferentiableAOTInput):
|
||||
continue
|
||||
if isinstance(desc, SubclassGetAttrAOTInput):
|
||||
_raise_autograd_subclass_not_implemented(n, desc)
|
||||
|
||||
input_index[desc] = (n, None)
|
||||
elif n.op == "output":
|
||||
if "desc" not in n.meta:
|
||||
raise AssertionError(f"'desc' not in n.meta for {n}: {n.meta}")
|
||||
desc = n.meta["desc"]
|
||||
for sub_n, sub_desc in zip(n.args[0], desc):
|
||||
if isinstance(sub_desc, SubclassGetAttrAOTOutput):
|
||||
_raise_autograd_subclass_not_implemented(sub_n, sub_desc)
|
||||
if isinstance(sub_desc, GradAOTOutput):
|
||||
inp, grad = input_index[sub_desc.grad_of]
|
||||
if grad is not None:
|
||||
raise AssertionError(
|
||||
f"grad already set for {sub_n}, {sub_desc}, {input_index}"
|
||||
)
|
||||
input_index[sub_desc.grad_of] = (inp, sub_n)
|
||||
return input_index
|
||||
|
||||
|
||||
def get_all_output_and_tangent_nodes(
|
||||
g: fx.Graph,
|
||||
) -> dict[DifferentiableAOTOutput, tuple[fx.Node, fx.Node | None]]:
|
||||
"""Get all output nodes and their corresponding tangent nodes from a joint graph.
|
||||
|
||||
Similar to get_all_input_and_grad_nodes, but returns output nodes paired with
|
||||
their tangent nodes (if they exist). This function traverses the graph to find
|
||||
all differentiable outputs and matches them with their corresponding tangent
|
||||
inputs used in forward-mode autodiff.
|
||||
|
||||
NB: *all* forward tensor output sare turned, including non-differentiable outputs,
|
||||
so you can use this function to perform operations on all outputs.
|
||||
|
||||
Args:
|
||||
g: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each DifferentiableAOTOutput descriptor to a tuple
|
||||
containing:
|
||||
- The output node itself
|
||||
- The tangent (input) node if it exists, None otherwise
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the joint graph has subclass tensor inputs/outputs; this
|
||||
is not supported by API as there is not necessarily a 1-1 correspondence
|
||||
between outputs and tangents when subclasses are involved.
|
||||
"""
|
||||
output_index: dict[DifferentiableAOTOutput, tuple[fx.Node, fx.Node | None]] = {}
|
||||
for n in g.nodes:
|
||||
if n.op == "output":
|
||||
desc = n.meta["desc"]
|
||||
for sub_n, sub_d in zip(n.args[0], desc):
|
||||
# Skip outputs that cannot possibly be differentiable
|
||||
if not isinstance(sub_d, DifferentiableAOTOutput):
|
||||
continue
|
||||
if isinstance(sub_d, SubclassGetAttrAOTOutput):
|
||||
_raise_autograd_subclass_not_implemented(sub_n, sub_d)
|
||||
|
||||
output_index[sub_d] = (sub_n, None)
|
||||
for n in g.nodes:
|
||||
if n.op == "placeholder":
|
||||
desc = n.meta["desc"]
|
||||
if isinstance(desc, SubclassGetAttrAOTInput):
|
||||
_raise_autograd_subclass_not_implemented(n, desc)
|
||||
if isinstance(desc, TangentAOTInput):
|
||||
out, tangent = output_index[desc.output]
|
||||
if tangent is not None:
|
||||
raise AssertionError(
|
||||
f"tangent already set for {n}, {desc}, {output_index}"
|
||||
)
|
||||
output_index[desc.output] = (out, n)
|
||||
return output_index
|
||||
|
||||
|
||||
def get_param_and_grad_nodes(
|
||||
graph: fx.Graph,
|
||||
) -> dict[ParamAOTInput, tuple[fx.Node, fx.Node | None]]:
|
||||
"""Get parameter nodes and their corresponding gradient nodes from a joint graph.
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each ParamAOTInput descriptor to a tuple containing:
|
||||
- The parameter input node
|
||||
- The gradient (output) node if it exists, None otherwise
|
||||
"""
|
||||
return {
|
||||
desc: (n, g)
|
||||
for desc, (n, g) in get_all_input_and_grad_nodes(graph).items()
|
||||
if isinstance(desc, ParamAOTInput)
|
||||
}
|
||||
|
||||
|
||||
def get_plain_input_and_grad_nodes(
|
||||
graph: fx.Graph,
|
||||
) -> dict[PlainAOTInput, tuple[fx.Node, fx.Node | None]]:
|
||||
"""Get plain input nodes and their corresponding gradient nodes from a joint graph.
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each PlainAOTInput descriptor to a tuple containing:
|
||||
- The plain input node
|
||||
- The gradient (output) node if it exists, None otherwise
|
||||
"""
|
||||
return {
|
||||
desc: (n, g)
|
||||
for desc, (n, g) in get_all_input_and_grad_nodes(graph).items()
|
||||
if isinstance(desc, PlainAOTInput)
|
||||
}
|
||||
|
||||
|
||||
def get_plain_output_and_tangent_nodes(
|
||||
graph: fx.Graph,
|
||||
) -> dict[PlainAOTOutput, tuple[fx.Node, fx.Node | None]]:
|
||||
"""Get plain output nodes and their corresponding tangent nodes from a joint graph.
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping each PlainAOTOutput descriptor to a tuple containing:
|
||||
- The plain output node
|
||||
- The tangent (input) node if it exists, None otherwise
|
||||
"""
|
||||
return {
|
||||
desc: (n, g)
|
||||
for desc, (n, g) in get_all_output_and_tangent_nodes(graph).items()
|
||||
if isinstance(desc, PlainAOTOutput)
|
||||
}
|
||||
|
||||
|
||||
def _raise_fqn_subclass_not_implemented(
|
||||
n: fx.Node, desc: AOTInput | AOTOutput
|
||||
) -> NoReturn:
|
||||
raise RuntimeError(
|
||||
"Subclasses are currently not supported by this function, but a desugared subclass input "
|
||||
f"was found at {n} ({desc}). The problem is "
|
||||
"that there may not necessarily be a 1-1 correspondence between a FQN and a plain tensor "
|
||||
"when subclasses are involved: for example, a parameter that is a subclass "
|
||||
"would desugar into multiple plain tensors, which we can't uniquely assign the "
|
||||
"FQN to. It's not clear what you want the API to do in this case: do you want to "
|
||||
"instead return a struct of nodes showing how to assemble the subclass? But you "
|
||||
"don't (directly) have the metadata for the subclass? If you have a concrete use "
|
||||
"case, please file an issue so we can understand it and design an API that works for your case."
|
||||
)
|
||||
|
||||
|
||||
def get_named_param_nodes(graph: fx.Graph) -> dict[str, fx.Node]:
|
||||
"""Get parameter nodes mapped by their fully qualified names.
|
||||
|
||||
This function traverses the graph to find all parameter input nodes and
|
||||
returns them in a dictionary where keys are the parameter names (FQNs)
|
||||
and values are the corresponding FX nodes.
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping parameter names (str) to their corresponding FX nodes.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If subclass tensors are encountered (not yet supported), as
|
||||
with subclasses a FQN does not necessarily map to a single plain tensor.
|
||||
"""
|
||||
r = {}
|
||||
for n in graph.nodes:
|
||||
if n.op == "placeholder":
|
||||
desc = n.meta["desc"]
|
||||
if isinstance(desc, SubclassGetAttrAOTInput):
|
||||
_raise_fqn_subclass_not_implemented(n, desc)
|
||||
elif isinstance(desc, ParamAOTInput):
|
||||
r[desc.target] = n
|
||||
return r
|
||||
|
||||
|
||||
def get_named_buffer_nodes(graph: fx.Graph) -> dict[str, fx.Node]:
|
||||
"""Get buffer nodes mapped by their fully qualified names.
|
||||
|
||||
This function traverses the graph to find all buffer input nodes and
|
||||
returns them in a dictionary where keys are the buffer names (FQNs)
|
||||
and values are the corresponding FX nodes.
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A dictionary mapping buffer names (str) to their corresponding FX nodes.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If subclass tensors are encountered (not yet supported), as
|
||||
with subclasses a FQN does not necessarily map to a single plain tensor.
|
||||
"""
|
||||
r = {}
|
||||
for n in graph.nodes:
|
||||
if n.op == "placeholder":
|
||||
desc = n.meta["desc"]
|
||||
if isinstance(desc, SubclassGetAttrAOTInput):
|
||||
_raise_fqn_subclass_not_implemented(n, desc)
|
||||
elif isinstance(desc, BufferAOTInput):
|
||||
r[desc.target] = n
|
||||
return r
|
||||
|
||||
|
||||
def get_param_nodes(graph: fx.Graph) -> list[fx.Node]:
|
||||
"""Get all parameter nodes from a graph as a list.
|
||||
|
||||
You can rely on this providing the correct order of parameters you need
|
||||
to feed into the joint graph (at the very beginning of the argument list,
|
||||
before buffers).
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A list of FX nodes representing all parameters in the graph.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If subclass tensors are encountered (not yet supported), as
|
||||
it is not clear if you wanted each individual constituent piece of the
|
||||
subclasses, or have them grouped up in some way.
|
||||
"""
|
||||
return list(get_named_param_nodes(graph).values())
|
||||
|
||||
|
||||
def get_buffer_nodes(graph: fx.Graph) -> list[fx.Node]:
|
||||
"""Get all buffer nodes from a graph as a list.
|
||||
|
||||
You can rely on this providing the correct order of buffers you need
|
||||
to feed into the joint graph (after parameters).
|
||||
|
||||
Args:
|
||||
graph: The FX joint graph with descriptors
|
||||
|
||||
Returns:
|
||||
A list of FX nodes representing all buffers in the graph.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If subclass tensors are encountered (not yet supported), as
|
||||
it is not clear if you wanted each individual constituent piece of the
|
||||
subclasses, or have them grouped up in some way.
|
||||
"""
|
||||
return list(get_named_buffer_nodes(graph).values())
|
||||
+585
@@ -0,0 +1,585 @@
|
||||
"""
|
||||
This module dispatches the graphs to either the forward-only or joint compilation
|
||||
pathways, taking into account the AOTConfig and the collected ViewAndMutationMetadata.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import dataclasses
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
import torch.utils.dlpack
|
||||
from torch._dispatch.python import enable_python_dispatcher
|
||||
from torch._dynamo.utils import detect_fake_mode, lazy_format_graph_code
|
||||
from torch._logging import getArtifactLogger, trace_structured
|
||||
from torch._subclasses.functional_tensor import FunctionalTensorMode
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
from torchgen.utils import dataclass_repr
|
||||
|
||||
from .. import config
|
||||
from .descriptors import AOTInput, BackwardTokenAOTInput
|
||||
from .functional_utils import (
|
||||
assert_functional_graph,
|
||||
propagate_input_mutation_stacktraces,
|
||||
)
|
||||
from .graph_capture_wrappers import (
|
||||
aot_dispatch_subclass,
|
||||
create_functionalized_fn,
|
||||
create_joint,
|
||||
fn_input_mutations_to_outputs,
|
||||
fn_prepped_for_autograd,
|
||||
handle_effect_tokens_fn,
|
||||
)
|
||||
from .schemas import AOTConfig, FxValue, SubclassMeta, TraceFn, ViewAndMutationMeta
|
||||
from .streams import (
|
||||
assign_backward_streams,
|
||||
assign_epilogue_copy_streams,
|
||||
insert_backward_syncs,
|
||||
populate_fw_metadata_with_stream_indices,
|
||||
sync_deallocations,
|
||||
wrap_all_sync_nodes_with_control_deps,
|
||||
)
|
||||
from .utils import (
|
||||
call_and_expect_output_descs,
|
||||
copy_fwd_metadata_to_bw_nodes,
|
||||
fn_wrappers,
|
||||
register_buffer_assignment_hook,
|
||||
root_module_when_exporting_non_strict,
|
||||
simple_wraps,
|
||||
unlift_tokens,
|
||||
)
|
||||
|
||||
|
||||
aot_graphs_log = getArtifactLogger(__name__, "aot_graphs")
|
||||
|
||||
|
||||
def _extract_tangent_source_stack_traces(
|
||||
fx_g: torch.fx.GraphModule,
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
) -> None:
|
||||
from .descriptors import PlainAOTOutput, TangentAOTInput
|
||||
|
||||
if not fw_metadata.traced_tangents_descs:
|
||||
return
|
||||
|
||||
output_node = list(fx_g.graph.nodes)[-1]
|
||||
all_outputs = output_node.args[0]
|
||||
|
||||
stack_traces: list[str | None] = []
|
||||
got_one = False
|
||||
|
||||
for desc in fw_metadata.traced_tangents_descs:
|
||||
stack_trace = None
|
||||
|
||||
if isinstance(desc, TangentAOTInput):
|
||||
output_desc = desc.output
|
||||
if isinstance(output_desc, PlainAOTOutput) and output_desc.idx < len(
|
||||
all_outputs
|
||||
):
|
||||
output_arg = all_outputs[output_desc.idx]
|
||||
if isinstance(output_arg, torch.fx.Node):
|
||||
stack_trace = output_arg.meta.get("stack_trace", None)
|
||||
got_one = True
|
||||
|
||||
stack_traces.append(stack_trace)
|
||||
|
||||
if got_one:
|
||||
fw_metadata.tangent_source_stack_traces = stack_traces
|
||||
|
||||
|
||||
def _create_graph(
|
||||
f: Callable[..., Any],
|
||||
args: list[torch.Tensor],
|
||||
args_descs: list[AOTInput]
|
||||
| None = None, # keep compat with old clients; maybe we should split into two impls
|
||||
*,
|
||||
aot_config: AOTConfig,
|
||||
) -> torch.fx.GraphModule:
|
||||
# FunctionalTensorMode must be enabled here.
|
||||
# See Note [Accessing .grad_fn on FunctionalTensor]
|
||||
out_descs = None
|
||||
|
||||
if args_descs is None:
|
||||
inner_f = f
|
||||
else:
|
||||
|
||||
@simple_wraps(f)
|
||||
def inner_f(*args: Any) -> Any:
|
||||
nonlocal out_descs
|
||||
if out_descs is not None:
|
||||
raise AssertionError("out_descs must be None")
|
||||
out, out_descs = call_and_expect_output_descs(f, args)
|
||||
return out
|
||||
|
||||
if aot_config.disable_functionalization:
|
||||
ctx = contextlib.nullcontext()
|
||||
else:
|
||||
ctx = FunctionalTensorMode( # type: ignore[assignment]
|
||||
pre_dispatch=aot_config.pre_dispatch,
|
||||
export=aot_config.is_export,
|
||||
# Allow token discovery for joint fn tracing as tokens can be used in backward.
|
||||
_allow_token_discovery=True,
|
||||
)
|
||||
|
||||
with (
|
||||
enable_python_dispatcher(),
|
||||
ctx,
|
||||
):
|
||||
fx_g = make_fx(
|
||||
inner_f,
|
||||
decomposition_table=aot_config.decompositions,
|
||||
record_module_stack=True,
|
||||
pre_dispatch=aot_config.pre_dispatch,
|
||||
_disable_torch_fn_metadata_mode=aot_config._disable_torch_fn_metadata_mode,
|
||||
)(*args)
|
||||
|
||||
if args_descs is not None:
|
||||
flat_args_descs, _ = pytree.tree_flatten(args_descs)
|
||||
flat_out_descs, _ = pytree.tree_flatten(out_descs)
|
||||
|
||||
# Unfortunately, flat_args_descs is not guaranteed to match the
|
||||
# number of actual arguments that show up on the FX graph.
|
||||
# Specifically, allow_token_discovery=True means that we will
|
||||
# silently add extra token arguments to the backwards graph.
|
||||
#
|
||||
# Although there are a few ways to detect what these tokens are,
|
||||
# we are going to settle for something dodgy but simple to
|
||||
# implement: match tangents_token placeholders specifically,
|
||||
# as these are the only placeholders that are created by token
|
||||
# discovery (NB: there is NO other code that treats this name
|
||||
# as load bearing, so this is a bit naughty!)
|
||||
#
|
||||
# I originally wanted to detect tokens in exactly the same way
|
||||
# that they are detected at normal runtime, but to be honest
|
||||
# the normal runtime detection is pretty strange: it seems the
|
||||
# backward tokens are not reliably at the end of the argument list
|
||||
# but *precede* the RNG arguments (I don't understand why this is
|
||||
# the case). And in unlift_tokens, token arguments are detected
|
||||
# by seeing if they feed into an effects call! Dastardly. Why
|
||||
# didn't we just introduce a new type.
|
||||
|
||||
i = 0
|
||||
j = 0
|
||||
for n in fx_g.graph.nodes:
|
||||
if n.op == "placeholder":
|
||||
if n.name.startswith("tangents_token"):
|
||||
n.meta["desc"] = BackwardTokenAOTInput(j)
|
||||
j += 1
|
||||
else:
|
||||
if i >= len(flat_args_descs):
|
||||
raise AssertionError(
|
||||
f"i={i} >= len(flat_args_descs)={len(flat_args_descs)}: "
|
||||
f"fn_wrappers={fn_wrappers(inner_f)}, "
|
||||
f"placeholders={[n for n in fx_g.graph.nodes if n.op == 'placeholder']}"
|
||||
)
|
||||
n.meta["desc"] = flat_args_descs[i]
|
||||
i += 1
|
||||
elif n.op == "output":
|
||||
n.meta["desc"] = flat_out_descs
|
||||
|
||||
return fx_g
|
||||
|
||||
|
||||
# TODO: Refactor the following code so detach() persists item_memo
|
||||
def _detach_and_copy_item_memo(t: torch.Tensor) -> torch.Tensor:
|
||||
detached_t = t.detach()
|
||||
if hasattr(t, "item_memo"):
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
detached_t.item_memo = t.item_memo
|
||||
return detached_t
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _GraphCaptureTracingResult:
|
||||
fn_to_trace: Callable[..., Any]
|
||||
flat_args: Any
|
||||
flat_args_descs: Any
|
||||
maybe_subclass_meta: SubclassMeta | None
|
||||
|
||||
|
||||
def _detach_traced_inputs(flat_args: Any) -> Any:
|
||||
if detect_fake_mode():
|
||||
detach_tensor = _detach_and_copy_item_memo
|
||||
else:
|
||||
|
||||
def detach_tensor(t: torch.Tensor) -> torch.Tensor:
|
||||
return t.detach()
|
||||
|
||||
return pytree.tree_map_only(torch.Tensor, detach_tensor, flat_args)
|
||||
|
||||
|
||||
def _prepare_graph_capture_tracing(
|
||||
fn_to_trace: Callable[..., Any],
|
||||
flat_args: Any,
|
||||
flat_args_descs: Any,
|
||||
flat_fn: TraceFn,
|
||||
*,
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
aot_config: AOTConfig,
|
||||
trace_joint: bool,
|
||||
joint_fn_handle: Any | None = None,
|
||||
) -> _GraphCaptureTracingResult:
|
||||
if aot_config.disable_functionalization:
|
||||
updated_flat_args, updated_flat_args_descs = flat_args, flat_args_descs
|
||||
else:
|
||||
fn_to_trace, updated_flat_args, updated_flat_args_descs = (
|
||||
create_functionalized_fn(
|
||||
fn_to_trace,
|
||||
flat_args,
|
||||
flat_args_descs,
|
||||
meta=fw_metadata,
|
||||
aot_config=aot_config,
|
||||
trace_joint=trace_joint,
|
||||
joint_fn_handle=joint_fn_handle,
|
||||
)
|
||||
)
|
||||
|
||||
subclass_tracing_info = aot_dispatch_subclass(
|
||||
fn_to_trace,
|
||||
updated_flat_args,
|
||||
updated_flat_args_descs,
|
||||
is_joint_structure=trace_joint,
|
||||
meta=fw_metadata,
|
||||
fw_only=flat_fn,
|
||||
)
|
||||
fn_to_trace = subclass_tracing_info.plain_tensor_trace_fn
|
||||
updated_flat_args = subclass_tracing_info.plain_tensor_args
|
||||
updated_flat_args_descs = subclass_tracing_info.plain_tensor_args_descs
|
||||
|
||||
if not aot_config.disable_functionalization:
|
||||
fn_to_trace, updated_flat_args, updated_flat_args_descs = (
|
||||
handle_effect_tokens_fn(
|
||||
fn_to_trace,
|
||||
updated_flat_args,
|
||||
updated_flat_args_descs,
|
||||
meta=fw_metadata,
|
||||
trace_joint=trace_joint,
|
||||
)
|
||||
)
|
||||
|
||||
return _GraphCaptureTracingResult(
|
||||
fn_to_trace=fn_to_trace,
|
||||
flat_args=updated_flat_args,
|
||||
flat_args_descs=updated_flat_args_descs,
|
||||
maybe_subclass_meta=subclass_tracing_info.maybe_subclass_meta,
|
||||
)
|
||||
|
||||
|
||||
def _create_graph_and_save_traced_inputs(
|
||||
fn_to_trace: Callable[..., Any],
|
||||
flat_args: Any,
|
||||
flat_args_descs: Any,
|
||||
*,
|
||||
aot_config: AOTConfig,
|
||||
) -> tuple[torch.fx.GraphModule, Any]:
|
||||
saved_flat_args = _detach_traced_inputs(flat_args)
|
||||
return (
|
||||
_create_graph(fn_to_trace, flat_args, flat_args_descs, aot_config=aot_config),
|
||||
saved_flat_args,
|
||||
)
|
||||
|
||||
|
||||
def aot_dispatch_base_graph(
|
||||
flat_fn: TraceFn,
|
||||
flat_args: list[FxValue],
|
||||
flat_args_descs: list[AOTInput],
|
||||
aot_config: AOTConfig,
|
||||
*,
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
) -> tuple[torch.fx.GraphModule, list[FxValue], list[AOTInput], SubclassMeta | None]:
|
||||
# aot_dispatch_base requires functionalization, but doesn't need to handle as many cases as the autograd case.
|
||||
# The cases that aot_dispatch_base doesn't need to handle include:
|
||||
# - outputs that are aliases of graph intermediates
|
||||
# - outputs that are aliases of graph inputs
|
||||
# While cases that it does need to handle include:
|
||||
# - input mutations (including when inputs are aliases of each other)
|
||||
# - input metadata mutations
|
||||
fn_to_trace = fn_input_mutations_to_outputs(
|
||||
flat_fn,
|
||||
flat_args_descs,
|
||||
fw_metadata,
|
||||
keep_data_input_mutations=aot_config.keep_inference_input_mutations,
|
||||
)
|
||||
# TODO: replace with AOTDispatchSubclassWrapper once we refactor
|
||||
# fn_input_mutations_to_outputs and create_functionalized_fn
|
||||
# into CompilerWrappers.
|
||||
tracing_state = _prepare_graph_capture_tracing(
|
||||
fn_to_trace,
|
||||
flat_args,
|
||||
flat_args_descs,
|
||||
flat_fn,
|
||||
fw_metadata=fw_metadata,
|
||||
aot_config=aot_config,
|
||||
trace_joint=False,
|
||||
)
|
||||
fn_to_trace = tracing_state.fn_to_trace
|
||||
updated_flat_args_subclasses_desugared = tracing_state.flat_args
|
||||
updated_flat_args_subclasses_desugared_descs = tracing_state.flat_args_descs
|
||||
maybe_subclass_meta = tracing_state.maybe_subclass_meta
|
||||
|
||||
aot_graphs_log.debug(
|
||||
"aot_config id: %s, fw_metadata=%s,subclass_metadata=%s",
|
||||
aot_config.aot_id,
|
||||
fw_metadata,
|
||||
maybe_subclass_meta,
|
||||
)
|
||||
|
||||
# We track buffer assignments when exporting in non-strict mode.
|
||||
# (In contrast, strict mode errors on any attribute assignment.)
|
||||
mod_when_exporting_non_strict = root_module_when_exporting_non_strict(flat_fn)
|
||||
if aot_config.is_export and mod_when_exporting_non_strict is not None:
|
||||
# For any buffer that is assigned, we want to associate it to the final proxy node
|
||||
# that it is assigned to. This node can then be added as a buffer mutation output.
|
||||
assigned_buffers: dict[str, str] = {}
|
||||
hook = register_buffer_assignment_hook(
|
||||
mod_when_exporting_non_strict, assigned_buffers
|
||||
)
|
||||
|
||||
(
|
||||
fw_module,
|
||||
saved_updated_flat_args_subclasses_desugared,
|
||||
) = _create_graph_and_save_traced_inputs(
|
||||
fn_to_trace,
|
||||
updated_flat_args_subclasses_desugared,
|
||||
updated_flat_args_subclasses_desugared_descs,
|
||||
aot_config=aot_config,
|
||||
)
|
||||
saved_updated_flat_args_subclasses_desugared_descs = (
|
||||
updated_flat_args_subclasses_desugared_descs
|
||||
)
|
||||
|
||||
if aot_config.is_export and mod_when_exporting_non_strict is not None:
|
||||
# We update metadata to consider any assigned buffers as buffer mutations.
|
||||
i = len(dict(mod_when_exporting_non_strict.named_parameters()))
|
||||
for name, _ in mod_when_exporting_non_strict.named_buffers():
|
||||
if name in assigned_buffers and not fw_metadata.input_info[i].mutates_data: # type: ignore[possibly-undefined]
|
||||
fw_metadata.input_info[i] = dataclasses.replace(
|
||||
fw_metadata.input_info[i], mutates_data=True
|
||||
)
|
||||
fw_metadata.num_mutated_inp_runtime_indices += 1
|
||||
i += 1
|
||||
|
||||
# We add nodes corresponding to buffer assignments as output nodes in the graph.
|
||||
add_nodes = []
|
||||
output_node = list(fw_module.graph.nodes)[-1]
|
||||
for name in assigned_buffers.values(): # type: ignore[possibly-undefined]
|
||||
for node in fw_module.graph.nodes:
|
||||
if node.name == name:
|
||||
add_nodes.append(node)
|
||||
node.users[output_node] = None
|
||||
output_node.args = ((*add_nodes, *output_node.args[0]),)
|
||||
|
||||
hook.remove() # type: ignore[possibly-undefined]
|
||||
|
||||
# As long as we opted to remove input mutations, then
|
||||
# there should be *NO* mutating ops in the graph at this point.
|
||||
if not aot_config.disable_functionalization:
|
||||
copy_count = assert_functional_graph(fw_module.graph)
|
||||
assign_epilogue_copy_streams(fw_module)
|
||||
# Wrap sync nodes with control_deps to prevent reordering
|
||||
wrap_all_sync_nodes_with_control_deps(fw_module)
|
||||
# Populate fw_metadata with stream indices from the compiled graph
|
||||
populate_fw_metadata_with_stream_indices(fw_module, fw_metadata)
|
||||
fw_module.graph.eliminate_dead_code()
|
||||
fw_module.recompile()
|
||||
copy_count2 = assert_functional_graph(fw_module.graph)
|
||||
propagate_input_mutation_stacktraces(fw_module.graph)
|
||||
if copy_count != copy_count2:
|
||||
raise AssertionError(
|
||||
f"copy_count={copy_count} != copy_count2={copy_count2}"
|
||||
)
|
||||
else:
|
||||
fw_module.graph.eliminate_dead_code()
|
||||
|
||||
# See Note [Side-Effectful Tokens in AOTAutograd]
|
||||
num_tokens = len(fw_metadata.tokens)
|
||||
if num_tokens != 0 and config.unlift_effect_tokens:
|
||||
unlift_tokens(fw_module, fw_metadata, aot_config)
|
||||
saved_updated_flat_args_subclasses_desugared = (
|
||||
saved_updated_flat_args_subclasses_desugared[num_tokens:]
|
||||
)
|
||||
saved_updated_flat_args_subclasses_desugared_descs = (
|
||||
saved_updated_flat_args_subclasses_desugared_descs[num_tokens:]
|
||||
)
|
||||
|
||||
if aot_config.enable_log:
|
||||
aot_graphs_log.info(
|
||||
"%s",
|
||||
lazy_format_graph_code(
|
||||
"Forward graph",
|
||||
fw_module,
|
||||
aot_config.aot_id,
|
||||
include_stride=True,
|
||||
include_device=True,
|
||||
colored=True,
|
||||
# For more expanded output set this to True (but can't default
|
||||
# to this because it affects tests):
|
||||
expanded_def=False,
|
||||
),
|
||||
)
|
||||
|
||||
trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "aot_forward_graph_fw_metadata",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: dataclass_repr(fw_metadata),
|
||||
)
|
||||
if maybe_subclass_meta is not None:
|
||||
trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "aot_forward_graph_fw_subclass_metadata",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: dataclass_repr(maybe_subclass_meta),
|
||||
)
|
||||
|
||||
trace_structured(
|
||||
"aot_inference_graph",
|
||||
payload_fn=lambda: fw_module.print_readable(
|
||||
print_output=False,
|
||||
include_stride=True,
|
||||
include_device=True,
|
||||
expanded_def=True,
|
||||
),
|
||||
)
|
||||
|
||||
# TODO: should factor this into a separate function for export that always only returns just the graph.
|
||||
if aot_config.is_export and maybe_subclass_meta is not None:
|
||||
raise AssertionError(
|
||||
"aot_export_module does not support tensor subclass inputs for now."
|
||||
)
|
||||
return (
|
||||
fw_module,
|
||||
saved_updated_flat_args_subclasses_desugared,
|
||||
saved_updated_flat_args_subclasses_desugared_descs,
|
||||
maybe_subclass_meta,
|
||||
)
|
||||
|
||||
|
||||
# Has the precondition that there
|
||||
# are no duplicate arguments in flat_args (e.g., the same Tensor
|
||||
# object never shows up twice. However, two tensor inputs MAY alias
|
||||
# the same storage, so long as they have separate TensorImpls.)
|
||||
def aot_dispatch_autograd_graph(
|
||||
flat_fn: TraceFn,
|
||||
flat_args: list[Any],
|
||||
flat_args_descs: list[AOTInput],
|
||||
aot_config: AOTConfig,
|
||||
*,
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
) -> tuple[
|
||||
torch.fx.GraphModule,
|
||||
tuple[list[Any], list[Any]],
|
||||
tuple[list[AOTInput], list[AOTInput]],
|
||||
SubclassMeta | None,
|
||||
]:
|
||||
# NB: flat_fn here is the original user function (as far as
|
||||
# aot_module_simplified is concerned)
|
||||
|
||||
# traced_tangents corresponds to the set of outputs in the traced forward that should get grad_outputs in the traced backward.
|
||||
# It includes outputs of the original forward, *and* any updated inputs due to input mutations.
|
||||
# However, it does *not* include any outputs that are aliases of inputs or intermediates, or any metadata-only input mutations.
|
||||
joint_inputs = (flat_args, fw_metadata.traced_tangents)
|
||||
joint_inputs_descs = (flat_args_descs, fw_metadata.traced_tangents_descs)
|
||||
|
||||
fn_prepared_for_autograd = fn_prepped_for_autograd(
|
||||
flat_fn,
|
||||
flat_args_descs,
|
||||
fw_metadata,
|
||||
aot_config,
|
||||
)
|
||||
joint_fn_to_trace = create_joint(
|
||||
fn_prepared_for_autograd, flat_args_descs, aot_config=aot_config
|
||||
)
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
joint_fn_handle = joint_fn_to_trace.handle
|
||||
# TODO: replace with AOTDispatchSubclassWrapper once we refactor
|
||||
# fn_input_mutations_to_outputs and create_functionalized_fn
|
||||
# into CompilerWrappers.
|
||||
tracing_state = _prepare_graph_capture_tracing(
|
||||
joint_fn_to_trace,
|
||||
joint_inputs,
|
||||
joint_inputs_descs,
|
||||
flat_fn,
|
||||
fw_metadata=fw_metadata,
|
||||
aot_config=aot_config,
|
||||
trace_joint=True,
|
||||
joint_fn_handle=joint_fn_handle,
|
||||
)
|
||||
joint_fn_to_trace = tracing_state.fn_to_trace
|
||||
updated_joint_inputs = tracing_state.flat_args
|
||||
updated_joint_inputs_descs = tracing_state.flat_args_descs
|
||||
maybe_subclass_meta = tracing_state.maybe_subclass_meta
|
||||
|
||||
# When we call _create_graph, this may mutate the metadata of joint
|
||||
# inputs. But callers are expecting to get the original joint inputs. So
|
||||
# we make aliases of all the inputs to make sure we have a copy that
|
||||
# doesn't get modified.
|
||||
#
|
||||
# This destroys requires_grad/grad_fn information. However, backends
|
||||
# beneath AOTAutograd are indifferent to this information, so it doesn't
|
||||
# matter.
|
||||
fx_g, saved_updated_joint_inputs = _create_graph_and_save_traced_inputs(
|
||||
joint_fn_to_trace,
|
||||
updated_joint_inputs,
|
||||
updated_joint_inputs_descs,
|
||||
aot_config=aot_config,
|
||||
)
|
||||
|
||||
# Redundant with the check above, but worth having in case tracing introduced
|
||||
# a fake tensor. Unlikely.
|
||||
# See Note: [Fake Modules and AOTAutograd]
|
||||
torch._dynamo.utils.assert_no_fake_params_or_buffers(fx_g)
|
||||
|
||||
# Have to copy before eliminate_dead_code otherwise the
|
||||
# fw node match might be erased
|
||||
copy_fwd_metadata_to_bw_nodes(fx_g)
|
||||
|
||||
# After copying metadata, assign streams to gradient accumulation nodes
|
||||
assign_backward_streams(fx_g)
|
||||
|
||||
assign_epilogue_copy_streams(fx_g)
|
||||
|
||||
# Insert syncs for newly assigned backward streams
|
||||
insert_backward_syncs(fx_g)
|
||||
|
||||
# Sync deallocations for tensors where the stream w/ their last usage
|
||||
# is distinct from their allocation stream
|
||||
sync_deallocations(fx_g)
|
||||
|
||||
# Wrap sync nodes with control_deps to prevent reordering
|
||||
# (must be after sync_deallocations which inserts additional sync nodes)
|
||||
wrap_all_sync_nodes_with_control_deps(fx_g)
|
||||
|
||||
# Populate fw_metadata with stream indices from the compiled graph
|
||||
# NB: This needs to be done after the above stream assignments
|
||||
populate_fw_metadata_with_stream_indices(fx_g, fw_metadata)
|
||||
|
||||
# this helps users identify which forward output to call .detach() on.
|
||||
_extract_tangent_source_stack_traces(fx_g, fw_metadata)
|
||||
|
||||
fx_g.graph.eliminate_dead_code()
|
||||
if not aot_config.disable_functionalization:
|
||||
# There should be *NO* mutating ops in the graph at this point.
|
||||
assert_functional_graph(fx_g.graph)
|
||||
|
||||
fx_g.recompile()
|
||||
|
||||
# TODO: in AOTAutograd, we create metadata like _indices_of_inps_to_detach to detect
|
||||
# when we need to manually detach() some inputs in the forward.
|
||||
# Higher order ops might eventually need to do the same.
|
||||
if aot_config.is_export and maybe_subclass_meta is not None:
|
||||
raise AssertionError(
|
||||
"aot_export_module does not support tensor subclass inputs for now."
|
||||
)
|
||||
return (
|
||||
fx_g,
|
||||
saved_updated_joint_inputs,
|
||||
updated_joint_inputs_descs,
|
||||
maybe_subclass_meta,
|
||||
)
|
||||
+1554
File diff suppressed because it is too large
Load Diff
+2641
File diff suppressed because it is too large
Load Diff
+54
@@ -0,0 +1,54 @@
|
||||
from collections.abc import Iterator, MutableMapping
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
|
||||
K = TypeVar("K")
|
||||
V = TypeVar("V")
|
||||
|
||||
|
||||
# Used for fast next key access (using the fact that the dict is ordered)
|
||||
# Note: doesn't support deletion but we don't need it!
|
||||
class IndexedDict(MutableMapping[K, V], Generic[K, V]):
|
||||
"""A dict that maintains insertion order with O(1) index access."""
|
||||
|
||||
__slots__ = ("_dict", "_keys", "_key_to_index")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._dict: dict[K, V] = {}
|
||||
self._keys: list[K] = [] # typing: ignore[bad-override]
|
||||
self._key_to_index: dict[K, int] = {}
|
||||
|
||||
def __setitem__(self, key: K, value: V) -> None:
|
||||
if key not in self._dict:
|
||||
self._key_to_index[key] = len(self._keys)
|
||||
self._keys.append(key)
|
||||
self._dict[key] = value
|
||||
|
||||
def __getitem__(self, key: K) -> V:
|
||||
return self._dict[key]
|
||||
|
||||
def __delitem__(self, key: K) -> None:
|
||||
raise NotImplementedError("Deletion not supported for IndexedDict")
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._dict)
|
||||
|
||||
def __iter__(self) -> Iterator[K]:
|
||||
return iter(self._keys)
|
||||
|
||||
def __contains__(self, key: object) -> bool:
|
||||
return key in self._dict
|
||||
|
||||
def next_key(self, key: K) -> K | None:
|
||||
"""Get the next key in insertion order. O(1)."""
|
||||
idx = self._key_to_index.get(key)
|
||||
if idx is not None and idx + 1 < len(self._keys):
|
||||
return self._keys[idx + 1]
|
||||
return None
|
||||
|
||||
def prev_key(self, key: K) -> K | None:
|
||||
"""Get the previous key in insertion order. O(1)."""
|
||||
idx = self._key_to_index.get(key)
|
||||
if idx is not None and idx > 0:
|
||||
return self._keys[idx - 1]
|
||||
return None
|
||||
+496
@@ -0,0 +1,496 @@
|
||||
"""
|
||||
This module is one of the analysis modules - it takes as input a function or graph
|
||||
and some preexisting properties, and returns some data that is useful for deciding
|
||||
how to further proceed with compilation or construct runtime wrappers.
|
||||
|
||||
In particular, the following analyses are provided:
|
||||
1. Refine the view and mutation metadata collected previously - removing duplicate
|
||||
inputs or mapping views to their bases.
|
||||
2. We also analyze the function signature for export graphs.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import itertools
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch import Tensor
|
||||
from torch._C._dynamo.guards import compute_overlapping_tensors
|
||||
from torch._functorch._aot_autograd.schemas import PlainTensorMeta
|
||||
from torch._guards import StorageOverlap
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
from torch.fx.experimental.symbolic_shapes import is_concrete_int
|
||||
|
||||
from .collect_metadata_analysis import coerce_tangent_and_suggest_memory_format
|
||||
from .descriptors import AOTInput, InputMutationAOTOutput, TangentAOTInput
|
||||
from .schemas import (
|
||||
AOTConfig,
|
||||
BackwardSignature,
|
||||
GraphSignature,
|
||||
InputAliasInfo,
|
||||
MemoryFormatMeta,
|
||||
OutputAliasInfo,
|
||||
OutputType,
|
||||
ViewAndMutationMeta,
|
||||
)
|
||||
from .utils import strict_zip
|
||||
|
||||
|
||||
zip = strict_zip
|
||||
|
||||
|
||||
def remove_dupe_metadata(
|
||||
m: ViewAndMutationMeta,
|
||||
keep_arg_mask: list[bool],
|
||||
add_dupe_map: list[int],
|
||||
) -> ViewAndMutationMeta:
|
||||
if len(m.input_info) != len(keep_arg_mask):
|
||||
raise AssertionError(
|
||||
f"len(m.input_info)={len(m.input_info)} != len(keep_arg_mask)={len(keep_arg_mask)}"
|
||||
)
|
||||
# Easy invariant: the first argument should never be a dupe (it will be kept)
|
||||
if len(keep_arg_mask) == 0 or not keep_arg_mask[0]:
|
||||
raise AssertionError(
|
||||
"keep_arg_mask must be non-empty and keep_arg_mask[0] must be True"
|
||||
)
|
||||
|
||||
# Filter dupe'd mutated inputs out of traced_tangents
|
||||
num_data_mutations = len([x for x in m.input_info if x.mutates_data])
|
||||
other_traced_tangents = m.traced_tangents[num_data_mutations:]
|
||||
inp_traced_tangents = m.traced_tangents[:num_data_mutations]
|
||||
other_traced_tangents_descs = m.traced_tangents_descs[num_data_mutations:]
|
||||
inp_traced_tangents_descs = m.traced_tangents_descs[:num_data_mutations]
|
||||
filtered_inp_traced_tangents = [
|
||||
# See Note [Tangents memory format]
|
||||
x
|
||||
for i, x in enumerate(inp_traced_tangents)
|
||||
if keep_arg_mask[m.mutated_inp_runtime_indices[i]]
|
||||
]
|
||||
filtered_inp_traced_tangents_descs = [
|
||||
x_desc
|
||||
for i, x_desc in enumerate(inp_traced_tangents_descs)
|
||||
if keep_arg_mask[m.mutated_inp_runtime_indices[i]]
|
||||
]
|
||||
traced_tangents = filtered_inp_traced_tangents + other_traced_tangents
|
||||
traced_tangents_descs = (
|
||||
filtered_inp_traced_tangents_descs + other_traced_tangents_descs
|
||||
)
|
||||
|
||||
if m.subclass_tangent_meta is None:
|
||||
raise AssertionError("m.subclass_tangent_meta must not be None")
|
||||
subclass_tangent_meta = [
|
||||
PlainTensorMeta(
|
||||
0, memory_format=MemoryFormatMeta(memory_format=torch.contiguous_format)
|
||||
)
|
||||
] * len(filtered_inp_traced_tangents) + m.subclass_tangent_meta[num_data_mutations:]
|
||||
|
||||
return ViewAndMutationMeta(
|
||||
input_info=[x for i, x in enumerate(m.input_info) if keep_arg_mask[i]],
|
||||
# For outputs that are views of inputs, we store the index of the input that the output
|
||||
# was generated from. Need to update that index to account for removed dupes.
|
||||
output_info=[
|
||||
OutputAliasInfo(
|
||||
output_type=o.output_type,
|
||||
raw_type=o.raw_type,
|
||||
dynamic_dims=o.dynamic_dims,
|
||||
base_idx=None if o.base_idx is None else add_dupe_map[o.base_idx],
|
||||
requires_grad=o.requires_grad,
|
||||
requires_grad_for_backward=o.requires_grad_for_backward,
|
||||
view_meta_sequence=o.view_meta_sequence,
|
||||
)
|
||||
for o in m.output_info
|
||||
],
|
||||
num_intermediate_bases=m.num_intermediate_bases,
|
||||
keep_input_mutations=m.keep_input_mutations,
|
||||
traced_tangents=traced_tangents,
|
||||
traced_tangents_descs=traced_tangents_descs,
|
||||
# We are guaranteed not to get here, since dupes are not supported today with subclass inputs.
|
||||
subclass_inp_meta=[],
|
||||
subclass_fw_graph_out_meta=[],
|
||||
subclass_tangent_meta=subclass_tangent_meta,
|
||||
)
|
||||
|
||||
|
||||
# Given our ViewAndMutation metadata, this fn constructs a new set of metadata,
|
||||
# after adding synthetic base arguments to the function.
|
||||
# Most of the work in this fn is slogging through all of the metadata corresponding to inputs,
|
||||
# and updating it with our synthetic base calling convention.
|
||||
#
|
||||
# When config.debug_assert is set, we automatically regenerate the metadata
|
||||
# and compare it to this output for sanity.
|
||||
#
|
||||
# In addition to the updated metadata, also return the list of input indices
|
||||
# that will need to be updated in the synthetic base epilogue
|
||||
def create_synthetic_base_metadata(
|
||||
m: ViewAndMutationMeta,
|
||||
# Maps each outer argument idx to its inner idx (or, if this outer arg is generated from a
|
||||
# synthetic base, you get a tuple of (i, TensorMeta), telling you the base tensor idx, and view metadata)
|
||||
synthetic_base_info: list[int | tuple[int, torch.Tensor]],
|
||||
outer_args: list[Any],
|
||||
inner_args: list[Any],
|
||||
inner_args_desc: list[AOTInput],
|
||||
) -> tuple[ViewAndMutationMeta, list[int]]:
|
||||
# maps inner arg indices to outer arg indices
|
||||
synthetic_base_to_indices: dict[int, list[int]] = {}
|
||||
for inner_idx in range(len(inner_args)):
|
||||
outer_aliased_indices_of_current_base_arg = [
|
||||
outer_idx
|
||||
for outer_idx, inner_idx_or_tuple in enumerate(synthetic_base_info)
|
||||
if (isinstance(inner_idx_or_tuple, int) and inner_idx_or_tuple == inner_idx)
|
||||
or (
|
||||
isinstance(inner_idx_or_tuple, tuple)
|
||||
and inner_idx_or_tuple[0] == inner_idx
|
||||
)
|
||||
]
|
||||
synthetic_base_to_indices[inner_idx] = outer_aliased_indices_of_current_base_arg
|
||||
|
||||
# given the requires_grad info on mutated inputs,
|
||||
# generate the requires_grad info on those same mutated inputs, but after constructing synthetic bases.
|
||||
# pyrefly: ignore [implicit-any]
|
||||
input_infos = []
|
||||
for outer_indices in synthetic_base_to_indices.values():
|
||||
# leaf-ness should be all-or-nothing for aliased tensor.
|
||||
# (aka if "a" and "b" are views, then a.is_leaf == b.is_leaf)
|
||||
any_leaf = any(m.input_info[x].is_leaf for x in outer_indices)
|
||||
all_leaf = all(m.input_info[x].is_leaf for x in outer_indices)
|
||||
if any_leaf != all_leaf:
|
||||
raise AssertionError(
|
||||
f"any_leaf={any_leaf} != all_leaf={all_leaf} for outer_indices={outer_indices}"
|
||||
)
|
||||
|
||||
mutates_data = (
|
||||
True
|
||||
if len(outer_indices) > 1
|
||||
else m.input_info[outer_indices[0]].mutates_data
|
||||
)
|
||||
mutates_metadata = (
|
||||
False
|
||||
if len(outer_indices) > 1
|
||||
else m.input_info[outer_indices[0]].mutates_metadata
|
||||
)
|
||||
requires_grad = any(m.input_info[x].requires_grad for x in outer_indices)
|
||||
mutations_under_no_grad_or_inference_mode = all(
|
||||
m.input_info[x].mutations_under_no_grad_or_inference_mode
|
||||
for x in outer_indices
|
||||
)
|
||||
|
||||
mutation_inductor_storage_resize = all(
|
||||
m.input_info[x].mutation_inductor_storage_resize for x in outer_indices
|
||||
)
|
||||
|
||||
inpt_info = InputAliasInfo(
|
||||
# If len(outer_indices) > 1, then this input is a synthetic base.
|
||||
# The invariant is that to the rest of aot autograd, synthetic bases only show up if
|
||||
# one of their aliases gets a data mutation. And if any of their aliases get metadata
|
||||
# mutations, they will be hidden from the rest of aot autograd.
|
||||
mutates_data=mutates_data,
|
||||
mutates_metadata=mutates_metadata,
|
||||
mutations_hidden_from_autograd=all(
|
||||
m.input_info[x].mutations_hidden_from_autograd for x in outer_indices
|
||||
),
|
||||
mutates_storage_metadata=(
|
||||
False
|
||||
if len(outer_indices) > 1
|
||||
else m.input_info[outer_indices[0]].mutates_storage_metadata
|
||||
),
|
||||
mutations_under_no_grad_or_inference_mode=mutations_under_no_grad_or_inference_mode,
|
||||
mutation_inductor_storage_resize=mutation_inductor_storage_resize,
|
||||
is_leaf=any_leaf,
|
||||
requires_grad=requires_grad,
|
||||
keep_input_mutations=m.keep_input_mutations,
|
||||
)
|
||||
input_infos.append(inpt_info)
|
||||
|
||||
# Find any inputs that fulfill the following criteria:
|
||||
# (1) They are part of a synthetic base (because they alias another input,
|
||||
# and at least one input experiences a data mutation)
|
||||
# (2) They experience a metadata mutation
|
||||
outer_aliased_arg_idx_with_metadata_mutations = [
|
||||
outer_idx
|
||||
for outer_idx, inpt_info in enumerate(m.input_info)
|
||||
if inpt_info.mutates_metadata
|
||||
and not isinstance(synthetic_base_info[outer_idx], int)
|
||||
]
|
||||
|
||||
# grab the original requires grad info on the outputs, except the ones from the mutated inputs
|
||||
input_metadata_output_info = [
|
||||
OutputAliasInfo(
|
||||
output_type=OutputType.alias_of_input,
|
||||
raw_type=FunctionalTensor,
|
||||
dynamic_dims={
|
||||
i
|
||||
for i, s in enumerate(outer_args[outer_idx].shape)
|
||||
if not is_concrete_int(s)
|
||||
},
|
||||
base_idx=synthetic_base_info[outer_idx][0], # type: ignore[index]
|
||||
requires_grad=(requires_grad := outer_args[outer_idx].requires_grad),
|
||||
requires_grad_for_backward=requires_grad,
|
||||
)
|
||||
for outer_idx in outer_aliased_arg_idx_with_metadata_mutations
|
||||
]
|
||||
existing_output_infos = []
|
||||
for o in m.output_info:
|
||||
new_base_idx = (
|
||||
None
|
||||
if o.base_idx is None
|
||||
else (
|
||||
synthetic_base_info[o.base_idx]
|
||||
if isinstance(synthetic_base_info[o.base_idx], int)
|
||||
else synthetic_base_info[o.base_idx][0] # type: ignore[index]
|
||||
)
|
||||
)
|
||||
# If base_idx is changed for OutputType.is_input, we need to update the output type to reflect the change
|
||||
new_output_type = (
|
||||
OutputType.alias_of_input
|
||||
if o.output_type == OutputType.is_input and o.base_idx != new_base_idx
|
||||
else o.output_type
|
||||
)
|
||||
existing_output_infos.append(
|
||||
OutputAliasInfo(
|
||||
output_type=new_output_type,
|
||||
raw_type=o.raw_type,
|
||||
dynamic_dims=o.dynamic_dims,
|
||||
# Map the input idx pre-synthetic-bases to the new idx post-synthetic-bases
|
||||
base_idx=new_base_idx, # type: ignore[arg-type]
|
||||
requires_grad=o.requires_grad,
|
||||
requires_grad_for_backward=o.requires_grad_for_backward,
|
||||
view_meta_sequence=o.view_meta_sequence,
|
||||
)
|
||||
)
|
||||
|
||||
inner_mutated_tangents_and_memory_formats = [
|
||||
# See Note [Tangents memory format]
|
||||
(
|
||||
coerce_tangent_and_suggest_memory_format(x),
|
||||
TangentAOTInput(InputMutationAOTOutput(x_desc)),
|
||||
)
|
||||
for inner_idx, (x, x_desc) in enumerate(zip(inner_args, inner_args_desc))
|
||||
if input_infos[inner_idx].mutates_data and input_infos[inner_idx].requires_grad
|
||||
]
|
||||
inner_mutated_tangents = [
|
||||
x[0][0] for x in inner_mutated_tangents_and_memory_formats
|
||||
]
|
||||
inner_mutated_tangents_descs = [
|
||||
x[1] for x in inner_mutated_tangents_and_memory_formats
|
||||
]
|
||||
inner_mutated_tangents_memory_formats = [
|
||||
x[0][1] for x in inner_mutated_tangents_and_memory_formats
|
||||
]
|
||||
|
||||
output_info = existing_output_infos + input_metadata_output_info
|
||||
# Regenerate traced tangents to include mutated inputs including synthetic bases
|
||||
traced_tangents = (
|
||||
inner_mutated_tangents + m.traced_tangents[len(inner_mutated_tangents) :]
|
||||
)
|
||||
traced_tangents_descs = (
|
||||
inner_mutated_tangents_descs
|
||||
+ m.traced_tangents_descs[len(inner_mutated_tangents) :]
|
||||
)
|
||||
if m.subclass_tangent_meta is None:
|
||||
raise AssertionError("m.subclass_tangent_meta must not be None")
|
||||
subclass_tangent_meta = [
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
PlainTensorMeta(0, memory_format=x)
|
||||
for x in inner_mutated_tangents_memory_formats
|
||||
] + m.subclass_tangent_meta[len(inner_mutated_tangents) :]
|
||||
|
||||
return (
|
||||
ViewAndMutationMeta(
|
||||
input_info=input_infos,
|
||||
output_info=output_info,
|
||||
num_intermediate_bases=m.num_intermediate_bases,
|
||||
keep_input_mutations=m.keep_input_mutations,
|
||||
traced_tangents=traced_tangents,
|
||||
traced_tangents_descs=traced_tangents_descs,
|
||||
# We are guaranteed not to get here, since synthetic_base codepaths are not supported today with subclass inputs.
|
||||
subclass_inp_meta=[],
|
||||
subclass_fw_graph_out_meta=[],
|
||||
subclass_tangent_meta=subclass_tangent_meta,
|
||||
),
|
||||
outer_aliased_arg_idx_with_metadata_mutations,
|
||||
)
|
||||
|
||||
|
||||
def compute_overlapping_inputs(
|
||||
aot_config: AOTConfig, fwd_inputs: list[Any], aliased_input_indices: list[int]
|
||||
) -> set[int]:
|
||||
num_aliases = len(aliased_input_indices)
|
||||
|
||||
shape_env = None
|
||||
maybe_suppress_guards = contextlib.nullcontext
|
||||
tracing_context = torch._guards.TracingContext.try_get()
|
||||
|
||||
if tracing_context is not None:
|
||||
if tracing_context.fake_mode is None:
|
||||
raise AssertionError("tracing_context.fake_mode must not be None")
|
||||
shape_env = tracing_context.fake_mode.shape_env
|
||||
|
||||
# Check whether we can actually get the dynamo sources from within AOTAutograd.
|
||||
if aot_config.aot_autograd_arg_pos_to_source and shape_env is not None:
|
||||
maybe_suppress_guards = shape_env.suppress_guards # type: ignore[assignment]
|
||||
|
||||
# Check whether there are any symbolic values being used.
|
||||
# We do this for 2 reasons:
|
||||
# 1. StorageOverlap guard is only issued whenever dynamic shapes is turned on
|
||||
# 2. Triggers the fast-path for computing storage overlapping
|
||||
symbolic = any(
|
||||
isinstance(x, torch.SymInt)
|
||||
for i in aliased_input_indices
|
||||
for x in [
|
||||
*fwd_inputs[i].shape,
|
||||
*fwd_inputs[i].stride(),
|
||||
fwd_inputs[i].storage_offset(),
|
||||
]
|
||||
)
|
||||
|
||||
if torch._inductor.config.is_fbcode():
|
||||
if symbolic and num_aliases > 400:
|
||||
from torch._subclasses.fake_tensor import (
|
||||
UnsupportedMutationAliasingException,
|
||||
)
|
||||
from torch._utils_internal import justknobs_check
|
||||
|
||||
msg = f"Encountered {num_aliases} dynamic, aliased/mutated inputs, consider setting dynamic=False"
|
||||
|
||||
if justknobs_check(
|
||||
"pytorch/compiler:aliased_inputs_with_mutation_and_dyn_shapes_killswitch",
|
||||
False,
|
||||
):
|
||||
raise UnsupportedMutationAliasingException(msg)
|
||||
|
||||
with maybe_suppress_guards():
|
||||
aliased_fwd_inputs = [fwd_inputs[i] for i in aliased_input_indices]
|
||||
actual_aliased_indices = {
|
||||
aliased_input_indices[i]
|
||||
for i in compute_overlapping_tensors(aliased_fwd_inputs, symbolic=symbolic)
|
||||
}
|
||||
|
||||
# Add the StorageOverlap AOTAutograd guard only if we are actually keeping track of
|
||||
# dynamo sources inside AOTAutograd.
|
||||
if (
|
||||
tracing_context is not None
|
||||
# Make sure dynamic shapes is currently being used.
|
||||
and symbolic
|
||||
# We check that we have more than 1 aliased tensor, which should be true at
|
||||
# this point, anyway.
|
||||
and num_aliases > 1
|
||||
and aot_config.aot_autograd_arg_pos_to_source
|
||||
):
|
||||
no_overlap_indices = list(set(aliased_input_indices) - actual_aliased_indices)
|
||||
|
||||
overlapping_sources = [
|
||||
aot_config.aot_autograd_arg_pos_to_source[i] for i in actual_aliased_indices
|
||||
]
|
||||
non_overlapping_sources = [
|
||||
aot_config.aot_autograd_arg_pos_to_source[i] for i in no_overlap_indices
|
||||
]
|
||||
|
||||
tracing_context.guards_context.aotautograd_guards.append(
|
||||
StorageOverlap(overlapping_sources, non_overlapping_sources)
|
||||
)
|
||||
|
||||
return actual_aliased_indices
|
||||
|
||||
|
||||
def _graph_input_names(gm: torch.fx.GraphModule) -> list[str]:
|
||||
return [node.name for node in gm.graph.find_nodes(op="placeholder")]
|
||||
|
||||
|
||||
def _graph_output_names(gm: torch.fx.GraphModule) -> list[Any]:
|
||||
output_node = next(iter(reversed(gm.graph.nodes)))
|
||||
if output_node.op != "output" or len(output_node.args) != 1:
|
||||
raise AssertionError(
|
||||
f"expected output node with 1 arg, got op={output_node.op}, args={len(output_node.args)}"
|
||||
)
|
||||
return_args = output_node.args[0]
|
||||
return [getattr(return_arg, "name", None) for return_arg in return_args]
|
||||
|
||||
|
||||
def create_graph_signature(
|
||||
fx_g: torch.fx.GraphModule,
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
in_spec: pytree.TreeSpec,
|
||||
out_spec: pytree.TreeSpec,
|
||||
*,
|
||||
user_args_flat: list[Tensor],
|
||||
params_and_buffers_flat: list[Tensor],
|
||||
param_names: list[str],
|
||||
buffer_names: list[str],
|
||||
trace_joint: bool,
|
||||
num_user_fw_outs: int | None,
|
||||
loss_index: int | None,
|
||||
) -> GraphSignature:
|
||||
# Retrieve graph input names
|
||||
graph_input_names = _graph_input_names(fx_g)
|
||||
# Retrieve graph output names
|
||||
graph_output_names = _graph_output_names(fx_g)
|
||||
|
||||
num_params_buffers = len(param_names) + len(buffer_names)
|
||||
num_tokens = len(fw_metadata.tokens)
|
||||
# We have enough restrictions on the graph (no de-duping, synthetic bases, etc),
|
||||
# Such that # graph inps = # user inps + # params + # buffers
|
||||
num_user_args = len(graph_input_names) - num_params_buffers - num_tokens
|
||||
|
||||
if trace_joint:
|
||||
if num_user_fw_outs is None:
|
||||
raise AssertionError(
|
||||
"num_user_fw_outs must not be None when trace_joint=True"
|
||||
)
|
||||
num_fw_outs = num_user_fw_outs + fw_metadata.num_mutated_inp_runtime_indices
|
||||
backward_output_names = graph_output_names[num_fw_outs:]
|
||||
|
||||
grad_index = itertools.count(0)
|
||||
gradients_to_parameters = {
|
||||
backward_output_names[next(grad_index)]: param_names[i]
|
||||
for i, param in enumerate(params_and_buffers_flat)
|
||||
if param.requires_grad
|
||||
}
|
||||
|
||||
gradients_to_user_inputs = {
|
||||
backward_output_names[next(grad_index)]: graph_input_names[
|
||||
i + len(params_and_buffers_flat)
|
||||
]
|
||||
for i, user_input in enumerate(user_args_flat)
|
||||
if user_input.requires_grad
|
||||
}
|
||||
|
||||
if len(gradients_to_parameters) + len(gradients_to_user_inputs) != len(
|
||||
backward_output_names
|
||||
):
|
||||
raise AssertionError(
|
||||
f"len(gradients_to_parameters)={len(gradients_to_parameters)} + "
|
||||
f"len(gradients_to_user_inputs)={len(gradients_to_user_inputs)} != "
|
||||
f"len(backward_output_names)={len(backward_output_names)}"
|
||||
)
|
||||
|
||||
# Check that we have fully accounted for all graph outputs
|
||||
if loss_index is None:
|
||||
raise AssertionError("loss_index must not be None")
|
||||
backward_signature = BackwardSignature(
|
||||
gradients_to_parameters,
|
||||
gradients_to_user_inputs,
|
||||
graph_output_names[loss_index],
|
||||
)
|
||||
else:
|
||||
backward_signature = None
|
||||
num_user_fw_outs = (
|
||||
len(graph_output_names)
|
||||
- fw_metadata.num_mutated_inp_runtime_indices
|
||||
- num_tokens
|
||||
)
|
||||
|
||||
return GraphSignature.from_tracing_metadata(
|
||||
in_spec=in_spec,
|
||||
out_spec=out_spec,
|
||||
graph_input_names=graph_input_names,
|
||||
graph_output_names=graph_output_names,
|
||||
view_mutation_metadata=fw_metadata,
|
||||
named_parameters=param_names,
|
||||
named_buffers=buffer_names,
|
||||
num_user_inputs=num_user_args,
|
||||
num_user_outputs=num_user_fw_outs,
|
||||
trace_joint=trace_joint,
|
||||
loss_index=loss_index,
|
||||
backward_signature=backward_signature,
|
||||
)
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
"""
|
||||
Contains utils for logging in AOTAutograd, including managing the names of the graphs under
|
||||
compilation, capturing user-friendly tracebacks, and debug messages.
|
||||
"""
|
||||
|
||||
import collections
|
||||
from collections.abc import Callable, Generator, Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.fx.traceback as fx_traceback
|
||||
|
||||
from .schemas import AOTConfig
|
||||
|
||||
|
||||
# This is a list since looking forward, we can have this arbitrarily nested.
|
||||
graph_being_compiled: list[str] = []
|
||||
# TODO: It would be nice to reset the numbering every time aot_id goes
|
||||
# up, but this is annoying to do right now (because we don't know if
|
||||
# an aot_id will come back from the dead), so right now this also happens
|
||||
# to be a globally unique number too (at the cost of wobbling if you change
|
||||
# how the graphs compile)
|
||||
nth_graph: int = 0
|
||||
model_name: str = "model"
|
||||
|
||||
|
||||
def set_model_name(name: str) -> None:
|
||||
global model_name
|
||||
model_name = name
|
||||
|
||||
|
||||
def get_aot_compilation_context() -> tuple[list[str], str, int]:
|
||||
return list(graph_being_compiled), model_name, nth_graph
|
||||
|
||||
|
||||
def get_aot_graph_name() -> str:
|
||||
"""
|
||||
Returns the name of the graph being compiled.
|
||||
"""
|
||||
global model_name, graph_being_compiled, nth_graph
|
||||
return f"{model_name}__{'_'.join(graph_being_compiled)}_{nth_graph}"
|
||||
|
||||
|
||||
get_graph_being_compiled = get_aot_graph_name
|
||||
|
||||
|
||||
@contextmanager
|
||||
def track_graph_compiling(
|
||||
aot_config: AOTConfig, graph_name: str
|
||||
) -> Generator[None, None, None]:
|
||||
global graph_being_compiled
|
||||
# TODO: Don't shove the aot_id in here; set it in the context
|
||||
graph_being_compiled = [f"{aot_config.aot_id}_{graph_name}"]
|
||||
old_name = None
|
||||
if tracing_context := torch._guards.TracingContext.try_get():
|
||||
old_name = tracing_context.aot_graph_name
|
||||
tracing_context.aot_graph_name = graph_being_compiled
|
||||
has_tracing_context = True
|
||||
else:
|
||||
has_tracing_context = False
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
global nth_graph
|
||||
nth_graph += 1
|
||||
graph_being_compiled = []
|
||||
if has_tracing_context:
|
||||
if tracing_context := torch._guards.TracingContext.try_get():
|
||||
tracing_context.aot_graph_name = old_name
|
||||
|
||||
|
||||
# Set up hooks so that during backward the fx's stack_trace is properly set
|
||||
callback_set = False
|
||||
|
||||
|
||||
def setup_stacktrace_preservation_hooks(roots: list[torch.autograd.graph.Node]) -> None:
|
||||
def iter_graph(
|
||||
roots: list[torch.autograd.graph.Node],
|
||||
) -> Iterator[torch.autograd.graph.Node]:
|
||||
if not roots:
|
||||
return
|
||||
seen = set()
|
||||
q = collections.deque()
|
||||
for node in roots:
|
||||
if node is not None and node not in seen:
|
||||
seen.add(node)
|
||||
q.append(node)
|
||||
|
||||
while q:
|
||||
node = q.popleft()
|
||||
for fn, _idx in node.next_functions:
|
||||
if fn in seen or fn is None:
|
||||
continue
|
||||
seen.add(fn)
|
||||
q.append(fn)
|
||||
|
||||
yield node
|
||||
|
||||
def get_callback(saved_stack_: list[str]) -> Callable[[], None]:
|
||||
def callback() -> None:
|
||||
global callback_set
|
||||
fx_traceback.set_stack_trace(saved_stack_)
|
||||
callback_set = False
|
||||
|
||||
return callback
|
||||
|
||||
def get_prehook(stack_: list[str], seq_nr: int) -> Callable[[Any], None]:
|
||||
def prehook(grad_output: Any) -> None:
|
||||
global callback_set
|
||||
|
||||
if not callback_set:
|
||||
torch.autograd.variable.Variable._execution_engine.queue_callback( # type: ignore[attr-defined]
|
||||
get_callback(fx_traceback.format_stack())
|
||||
)
|
||||
callback_set = True
|
||||
|
||||
fx_traceback.set_stack_trace(stack_)
|
||||
fx_traceback.set_grad_fn_seq_nr(seq_nr)
|
||||
fx_traceback._mark_autograd_backward()
|
||||
|
||||
return prehook
|
||||
|
||||
def get_posthook(
|
||||
special_stack_: list[str], seq_nr: int
|
||||
) -> Callable[[Any, Any], None]:
|
||||
def posthook(grad_input: Any, grad_output: Any) -> None:
|
||||
fx_traceback.set_stack_trace(special_stack_)
|
||||
fx_traceback.reset_grad_fn_seq_nr()
|
||||
fx_traceback._reset_autograd_backward()
|
||||
|
||||
return posthook
|
||||
|
||||
for node in iter_graph(roots):
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
forward_node_stack = node.metadata.get("traceback_", [])
|
||||
node.register_prehook(get_prehook(forward_node_stack, node._sequence_nr()))
|
||||
|
||||
special_stack = forward_node_stack.copy()
|
||||
special_stack.append(fx_traceback.GRADIENT_ACC_SPECIAL_STACK)
|
||||
node.register_hook(get_posthook(special_stack, node._sequence_nr()))
|
||||
|
||||
|
||||
def setup_stacktrace_preservation_hooks_from_tensors(outputs: Any) -> None:
|
||||
roots = [
|
||||
t.grad_fn
|
||||
for t in (outputs if isinstance(outputs, (list, tuple)) else (outputs,))
|
||||
if isinstance(t, torch.Tensor) and t.grad_fn is not None
|
||||
]
|
||||
if roots:
|
||||
setup_stacktrace_preservation_hooks(roots)
|
||||
|
||||
|
||||
def describe_input(i: int, aot_config: AOTConfig) -> str:
|
||||
if i < aot_config.num_params_buffers:
|
||||
return f"parameter/buffer {i}"
|
||||
else:
|
||||
return f"input {i - aot_config.num_params_buffers}"
|
||||
|
||||
|
||||
def format_guard_bug_msg(aot_config: AOTConfig, expected: str) -> str:
|
||||
return (
|
||||
f"At compilation time, graph {aot_config.aot_id} was compiled under the "
|
||||
f"assumption that {expected}, but at runtime this was not the case. "
|
||||
"This indicates a guard bug in AOTAutograd or Dynamo, please file a bug to PyTorch."
|
||||
)
|
||||
+3254
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,596 @@
|
||||
import operator
|
||||
from typing import Any, TYPE_CHECKING, TypeAlias
|
||||
|
||||
import torch.fx
|
||||
import torch.fx.traceback
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._dynamo.graph_utils import _get_flat_args
|
||||
from torch._dynamo.variables.streams import get_current_stream, new_event
|
||||
from torch.fx.node import map_arg
|
||||
from torch.utils._runtime_estimation import (
|
||||
_FLOAT_TYPES,
|
||||
_IGNORE_OPS,
|
||||
get_compute_time,
|
||||
get_transfer_time,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .schemas import ViewAndMutationMeta # noqa: TC004
|
||||
|
||||
from .indexed_dict import IndexedDict
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
Node: TypeAlias = torch.fx.Node
|
||||
Graph: TypeAlias = torch.fx.Graph
|
||||
|
||||
_SYNC_OPS = (
|
||||
torch.ops.streams.record_event.default,
|
||||
torch.ops.streams.wait_event.default,
|
||||
torch.ops.streams.synchronize_event.default,
|
||||
torch.ops.streams.synchronize_device.default,
|
||||
torch.ops.streams.synchronize_stream.default,
|
||||
)
|
||||
|
||||
|
||||
def get_roofline_estimate(node: Node) -> float:
|
||||
if node.op != "call_function":
|
||||
raise AssertionError(f"non-func node in roofline estimate: {node.op}")
|
||||
|
||||
def map_value(x: Any) -> Any:
|
||||
return x.meta.get("value", x) if isinstance(x, Node) else x
|
||||
|
||||
func = node.target
|
||||
if func in _IGNORE_OPS:
|
||||
return 0.0
|
||||
|
||||
mapped_args = torch.fx.map_arg(node.args, map_value)
|
||||
mapped_kwargs = torch.fx.map_arg(node.kwargs, map_value)
|
||||
flat_args_kwargs = [map_value(x) for x in _get_flat_args(node, {})]
|
||||
flat_outs, _ = pytree.tree_flatten(node.meta.get("value", node))
|
||||
out = node.meta.get("value", node)
|
||||
out_dtypes = {
|
||||
t.dtype
|
||||
for t in flat_outs
|
||||
if isinstance(t, torch.Tensor) and t.dtype in _FLOAT_TYPES
|
||||
}
|
||||
|
||||
return (
|
||||
max(
|
||||
get_transfer_time(flat_args_kwargs, flat_outs),
|
||||
get_compute_time(func, mapped_args, mapped_kwargs, out, out_dtypes),
|
||||
)
|
||||
/ 1e6
|
||||
)
|
||||
|
||||
|
||||
def is_gradient_acc(node: Node) -> bool:
|
||||
return node.meta.get("is_gradient_acc", False)
|
||||
|
||||
|
||||
def is_bwd_node(node: Node) -> bool:
|
||||
tag = node.meta.get("partitioner_tag")
|
||||
return tag == "is_backward" or tag == "must_be_in_backward"
|
||||
|
||||
|
||||
def get_device(node: Node) -> torch.device:
|
||||
return node.meta["val"].device
|
||||
|
||||
|
||||
def get_stream(node: Node) -> int | None:
|
||||
maybe_annotation = node.meta.get("custom", None)
|
||||
if maybe_annotation is not None:
|
||||
return node.meta["custom"].get("stream", None)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_stream_or_current_stream(node: Node) -> int:
|
||||
ind = get_stream(node)
|
||||
if ind is None:
|
||||
ind = get_current_stream(get_device(node))
|
||||
return ind
|
||||
|
||||
|
||||
def set_stream(node: Node, ind: int) -> None:
|
||||
if "custom" in node.meta:
|
||||
node.meta["custom"].update({"stream": ind})
|
||||
else:
|
||||
node.meta["custom"] = {"stream": ind}
|
||||
|
||||
|
||||
def insert_record_event_after_node(graph: Graph, node: Node, event_ind: int) -> Node:
|
||||
with graph.inserting_after(node):
|
||||
node = graph.call_function(
|
||||
torch.ops.streams.record_event.default,
|
||||
(
|
||||
event_ind,
|
||||
get_stream_or_current_stream(node),
|
||||
),
|
||||
)
|
||||
node.meta["partitioner_tag"] = "must_be_in_backward"
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def insert_wait_event_before_node(graph: Graph, node: Node, event_ind: int) -> Node:
|
||||
with graph.inserting_before(node):
|
||||
node = graph.call_function(
|
||||
torch.ops.streams.wait_event.default,
|
||||
(
|
||||
event_ind,
|
||||
get_stream_or_current_stream(node),
|
||||
),
|
||||
)
|
||||
node.meta["partitioner_tag"] = "must_be_in_backward"
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def populate_stream_timeline(
|
||||
stream_to_timeline: dict[int | None, IndexedDict[Node, float]],
|
||||
graph: Graph,
|
||||
stream_index: int | None,
|
||||
) -> IndexedDict[Node, float]:
|
||||
if stream_index not in stream_to_timeline:
|
||||
stream_to_timeline[stream_index] = IndexedDict()
|
||||
total_time = 0.0
|
||||
for node in graph.nodes:
|
||||
# mlazos: not sure if we should include forward here too but don't think it matters
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and is_bwd_node(node)
|
||||
and get_stream(node) == stream_index
|
||||
):
|
||||
total_time += get_roofline_estimate(node)
|
||||
stream_to_timeline[stream_index][node] = (
|
||||
total_time # NB: total time includes the node's runtime
|
||||
)
|
||||
|
||||
return stream_to_timeline[stream_index]
|
||||
|
||||
|
||||
# NB: we start all estimates at 0, estimating the total runtime of each stream with timestamps at each node
|
||||
# we then try and use these timestamps to estimate when to deallocate tensors used in side streams
|
||||
# See https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream
|
||||
# for details on the problem being addressed. Rather than using the automatic memory management approach of record_stream
|
||||
# we attempt to find the point which to deallocate based on the estimated timestamps.
|
||||
def handle_synced_deallocation(
|
||||
graph: Graph,
|
||||
stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]],
|
||||
node: Node,
|
||||
last_usage: Node,
|
||||
) -> None:
|
||||
if not is_bwd_node(node):
|
||||
raise AssertionError(
|
||||
"synced allocations should only be handled on backward nodes"
|
||||
)
|
||||
if not is_bwd_node(last_usage):
|
||||
raise AssertionError(
|
||||
"synced allocations should only be handled on backward nodes"
|
||||
)
|
||||
allocating_stream = get_stream(node)
|
||||
side_stream = get_stream(last_usage)
|
||||
if allocating_stream == side_stream:
|
||||
raise AssertionError(
|
||||
"allocating and side stream should be different for synced deallocations"
|
||||
)
|
||||
if not torch.cuda.is_available():
|
||||
# fallback to record_stream in this case
|
||||
with graph.inserting_after(node):
|
||||
graph.call_function(
|
||||
torch.ops.streams.record_stream.default,
|
||||
(
|
||||
node,
|
||||
get_stream_or_current_stream(last_usage),
|
||||
),
|
||||
{},
|
||||
)
|
||||
node.meta["partitioner_tag"] = "must_be_in_backward"
|
||||
|
||||
allocating_stream_trace = populate_stream_timeline(
|
||||
stream_to_exec_trace, graph, allocating_stream
|
||||
)
|
||||
side_stream_trace = populate_stream_timeline(
|
||||
stream_to_exec_trace, graph, side_stream
|
||||
)
|
||||
|
||||
alloc_ptr = node
|
||||
target_side_stream_time = side_stream_trace[last_usage]
|
||||
# linear search from first usage of tensor to a point in time after the side stream has finished
|
||||
while alloc_ptr is not None:
|
||||
alloc_time = allocating_stream_trace[alloc_ptr]
|
||||
|
||||
if alloc_time >= target_side_stream_time:
|
||||
break
|
||||
elif alloc_time < target_side_stream_time:
|
||||
next_ptr = allocating_stream_trace.next_key(alloc_ptr)
|
||||
if next_ptr is not None:
|
||||
alloc_ptr = next_ptr
|
||||
else:
|
||||
break
|
||||
|
||||
wait_event = new_event()
|
||||
record_node = insert_record_event_after_node(graph, last_usage, wait_event)
|
||||
with graph.inserting_after(max(alloc_ptr, record_node)):
|
||||
graph.call_function(
|
||||
torch.ops.streams.sync_dealloc.default,
|
||||
(wait_event, get_stream_or_current_stream(alloc_ptr), node),
|
||||
{},
|
||||
)
|
||||
node.meta["partitioner_tag"] = "must_be_in_backward"
|
||||
|
||||
|
||||
def insert_sync(
|
||||
graph: Graph,
|
||||
consumer: Node,
|
||||
producer: Node,
|
||||
node_to_wait_event_ind: dict[Node, int],
|
||||
) -> None:
|
||||
if producer not in node_to_wait_event_ind:
|
||||
node_to_wait_event_ind[producer] = new_event()
|
||||
|
||||
insert_record_event_after_node(
|
||||
graph, producer, node_to_wait_event_ind[producer]
|
||||
)
|
||||
insert_wait_event_before_node(graph, consumer, node_to_wait_event_ind[producer])
|
||||
|
||||
|
||||
def assign_backward_streams(gm: torch.fx.GraphModule) -> None:
|
||||
"""Assigns backward streams to gradient accumulation nodes"""
|
||||
|
||||
# NB: iterate in reverse order to more closely match eager
|
||||
# the user node stream will be populated first
|
||||
for node in reversed(list(gm.graph.nodes)):
|
||||
if is_gradient_acc(node):
|
||||
# Accumulation stream selection. Follow the rules from top to bottom to determine the accumulation stream:
|
||||
# 1. Match first stream assignment of the first user with a stream
|
||||
# 2. Match first stream assignment encountered in the args from left to right
|
||||
# This differs from eager in some cases:
|
||||
# Specifically the eager code uses the autograd node to determine the stream,
|
||||
# crucially this does not necessarily correspond to the FX graph node. For example,
|
||||
# in the backward for an add node with a constant we will passthrough and during backward tracing,
|
||||
# no op will be added to the FX graph, so our stream assignment will differ in this case.
|
||||
gradients = _get_flat_args(node, {})
|
||||
users = list(node.users.keys())
|
||||
|
||||
# All gradients will be on same device, they will be coerced if they were not with a .to() node
|
||||
for neighbor in users + gradients:
|
||||
ind = get_stream(neighbor)
|
||||
if ind is not None:
|
||||
set_stream(node, ind)
|
||||
break
|
||||
|
||||
|
||||
def insert_backward_syncs(gm: torch.fx.GraphModule) -> None:
|
||||
"""Inserts stream syncs for backward nodes if consumer and producer are on different streams"""
|
||||
node_to_wait_event_ind: dict[Node, int] = {}
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_function" and is_bwd_node(node):
|
||||
flat_args = _get_flat_args(node, {})
|
||||
cur_node_stream = get_stream(node)
|
||||
|
||||
for arg in flat_args:
|
||||
if arg.op == "call_function" and is_bwd_node(arg):
|
||||
arg_stream = get_stream(arg)
|
||||
if arg_stream != cur_node_stream and get_device(arg).type != "cpu":
|
||||
insert_sync(gm.graph, node, arg, node_to_wait_event_ind)
|
||||
|
||||
|
||||
def sync_deallocations(gm: torch.fx.GraphModule) -> None:
|
||||
"""Handles https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream"""
|
||||
# Note: this is only needed if the last usage of a tensor is on a stream other than
|
||||
# the stream the tensor was allocated on
|
||||
|
||||
# an estimated timestamp from the beginning of graph execution (assuming 0 CPU overhead)
|
||||
# I think this is fine because you should have large tensors if you're using streams
|
||||
# although perhaps I could add a constant 10us per op ahead of the first stream op?
|
||||
# a trace of all the nodes running in a given stream
|
||||
stream_to_exec_trace: dict[int | None, IndexedDict[Node, float]] = {}
|
||||
for node in gm.graph.nodes:
|
||||
if node.op == "call_function" and is_bwd_node(node):
|
||||
allocating_stream = get_stream(node)
|
||||
users = list(node.users.keys())
|
||||
if not users:
|
||||
continue
|
||||
last_user = max(user for user in users)
|
||||
if last_user.op == "output":
|
||||
continue
|
||||
side_stream = get_stream(last_user)
|
||||
if allocating_stream != side_stream:
|
||||
handle_synced_deallocation(
|
||||
gm.graph, stream_to_exec_trace, node, last_user
|
||||
)
|
||||
|
||||
|
||||
def assign_epilogue_copy_streams(gm: torch.fx.GraphModule) -> None:
|
||||
for epi_copy in gm.graph.find_nodes(op="call_function", target=aten.copy_.default):
|
||||
arg_stream = get_stream(epi_copy.args[1])
|
||||
copy_stream = get_stream(epi_copy)
|
||||
if arg_stream != copy_stream:
|
||||
set_stream(epi_copy, get_stream_or_current_stream(epi_copy.args[1]))
|
||||
|
||||
|
||||
def populate_fw_metadata_with_stream_indices(
|
||||
gm: torch.fx.GraphModule, fw_metadata: "ViewAndMutationMeta"
|
||||
) -> None:
|
||||
"""
|
||||
Populates fw_metadata.mutated_inp_stream_indices with stream indices from the compiled graph.
|
||||
|
||||
The forward graph outputs are structured as:
|
||||
(*mutated_inputs, *user_outputs, *intermediate_bases, *saved_tensors, *saved_symints)
|
||||
|
||||
We extract the stream index for each mutated input from the graph's output node.
|
||||
"""
|
||||
|
||||
num_mutated_inps = fw_metadata.num_mutated_inp_runtime_indices
|
||||
if num_mutated_inps == 0:
|
||||
fw_metadata.mutated_inp_stream_indices = []
|
||||
return
|
||||
|
||||
# Find the output node in the graph
|
||||
output_node = None
|
||||
for node in gm.graph.find_nodes(op="output"):
|
||||
output_node = node
|
||||
break
|
||||
|
||||
if output_node is None:
|
||||
raise AssertionError(
|
||||
"No output node found in the graph when extracting stream indices"
|
||||
)
|
||||
|
||||
# The output node's args[0] is a tuple/list of all outputs
|
||||
output_args = output_node.args[0]
|
||||
|
||||
# Extract stream indices for the first num_mutated_inps outputs
|
||||
stream_indices = []
|
||||
for i in range(num_mutated_inps):
|
||||
if i < len(output_args):
|
||||
output_arg = output_args[i]
|
||||
# Get the stream index from the node metadata
|
||||
stream_idx = (
|
||||
get_stream(output_arg)
|
||||
if isinstance(output_arg, torch.fx.Node)
|
||||
else None
|
||||
)
|
||||
stream_indices.append(stream_idx)
|
||||
else:
|
||||
stream_indices.append(None)
|
||||
|
||||
fw_metadata.mutated_inp_stream_indices = stream_indices
|
||||
|
||||
|
||||
def _wrap_sync_node(
|
||||
gm: torch.fx.GraphModule,
|
||||
sync_node: Node,
|
||||
deps_before_sync: list[Node],
|
||||
visited: set[Node],
|
||||
) -> tuple[Node, list[Node]]:
|
||||
"""
|
||||
Core logic: wrap a single sync node in control_deps.
|
||||
|
||||
Returns (control_deps_node, passthrough_getitems) where passthrough_getitems
|
||||
are the getitem nodes that thread dependencies through the control_deps node.
|
||||
``visited`` is the set of nodes at or before the sync node in graph order,
|
||||
used to distinguish pre-sync vs post-sync users.
|
||||
"""
|
||||
from torch._inductor.fx_passes.control_dependencies import (
|
||||
_create_subgraph_for_node,
|
||||
control_deps,
|
||||
get_subgraph_name,
|
||||
)
|
||||
|
||||
graph = gm.graph
|
||||
|
||||
# Use dep.users to find deps with uses after the sync — avoids a forward walk.
|
||||
deps_with_uses_after_sync = [
|
||||
dep
|
||||
for dep in deps_before_sync
|
||||
if any(user not in visited for user in dep.users)
|
||||
]
|
||||
|
||||
# Create subgraph that executes sync and passes through only used dependencies
|
||||
subgraph_module = _create_subgraph_for_node(
|
||||
graph, sync_node, deps_with_uses_after_sync
|
||||
)
|
||||
subgraph_attr_name = get_subgraph_name(gm, sync_node.name)
|
||||
setattr(gm, subgraph_attr_name, subgraph_module)
|
||||
|
||||
# Create control_deps call
|
||||
# Note: sync nodes (record_event/wait_event) only take int args, no Node args.
|
||||
with graph.inserting_before(sync_node):
|
||||
get_subgraph = graph.get_attr(subgraph_attr_name)
|
||||
control_deps_node = graph.call_function(
|
||||
control_deps,
|
||||
args=(
|
||||
tuple(deps_before_sync), # additional_deps (all deps for ordering)
|
||||
get_subgraph, # subgraph
|
||||
*deps_with_uses_after_sync, # only pass through deps that are used
|
||||
),
|
||||
kwargs={},
|
||||
)
|
||||
|
||||
# Mark newly created nodes as visited so subsequent syncs don't
|
||||
# misclassify them as "after the sync" during replacement.
|
||||
visited.add(get_subgraph)
|
||||
visited.add(control_deps_node)
|
||||
|
||||
# The output is (sync_result, *deps_with_uses_after_sync)
|
||||
# Create getitem nodes only for dependencies that have uses after sync
|
||||
replacements: dict[Node, Node] = {}
|
||||
with graph.inserting_after(control_deps_node):
|
||||
for i, dep in enumerate(deps_with_uses_after_sync):
|
||||
getitem_node = graph.call_function(
|
||||
operator.getitem,
|
||||
args=(control_deps_node, i + 1), # +1 because index 0 is sync result
|
||||
)
|
||||
getitem_node.meta.update(dep.meta)
|
||||
replacements[dep] = getitem_node
|
||||
visited.add(getitem_node)
|
||||
|
||||
# Replace uses of dependencies that come after sync_node.
|
||||
# Use map_arg to handle nested structures (e.g. output node's list args).
|
||||
for dep, getitem_node in replacements.items():
|
||||
for user in list(dep.users.keys()):
|
||||
if user is control_deps_node:
|
||||
continue
|
||||
if user in visited:
|
||||
continue
|
||||
# Don't replace forward outputs in the output node — they belong
|
||||
# to the forward partition and must not reference backward nodes.
|
||||
if user.op == "output" and not is_bwd_node(dep):
|
||||
continue
|
||||
|
||||
def _replace(n: Node) -> Node:
|
||||
return getitem_node if n is dep else n
|
||||
|
||||
user.args = map_arg(user.args, _replace)
|
||||
user.kwargs = map_arg(user.kwargs, _replace)
|
||||
|
||||
# Remove original sync node
|
||||
sync_node.replace_all_uses_with(control_deps_node)
|
||||
graph.erase_node(sync_node)
|
||||
return control_deps_node, list(replacements.values())
|
||||
|
||||
|
||||
def wrap_all_sync_nodes_with_control_deps(gm: torch.fx.GraphModule) -> None:
|
||||
"""
|
||||
Single-pass wrap of all sync nodes in control_deps.
|
||||
|
||||
Iterates through the graph once, accumulating per-stream node lists.
|
||||
When a sync node is encountered, it is wrapped using the accumulated deps
|
||||
for that stream, then the deps are reset to the control_deps node
|
||||
(maintaining the ordering chain for subsequent syncs on the same stream).
|
||||
"""
|
||||
graph = gm.graph
|
||||
if len(graph.nodes) == 0:
|
||||
raise RuntimeError("Expected a non-empty graph")
|
||||
stream_to_nodes: dict[int | None, list[Node]] = {}
|
||||
# Maps event_index -> control_deps node that wrapped its record_event,
|
||||
# so the corresponding wait_event/synchronize_event can depend on the record.
|
||||
event_to_ctrl: dict[int, Node] = {}
|
||||
# Maps event_index -> getitem nodes threaded through record_event's control_deps,
|
||||
# so synchronize_event can thread them through to subsequent ops.
|
||||
event_to_passthrough: dict[int, list[Node]] = {}
|
||||
# Maps event_index -> stream that the event was recorded on,
|
||||
# so synchronize_event can infer its stream.
|
||||
event_to_stream: dict[int, int | None] = {}
|
||||
visited: set[Node] = set()
|
||||
found_sync = False
|
||||
|
||||
# Walk the node linked-list manually so we can mutate the graph
|
||||
# (wrapping sync nodes inserts/erases nodes) without losing our place.
|
||||
node = next(iter(graph.nodes))
|
||||
while node.op != "root":
|
||||
next_node = node.next
|
||||
visited.add(node)
|
||||
|
||||
if node.op == "call_function":
|
||||
if node.target in _SYNC_OPS:
|
||||
# synchronize_device and synchronize_stream block the CPU,
|
||||
# so all subsequent kernel launches are host-ordered after
|
||||
# them. Treat both as full barriers across all streams.
|
||||
if node.target in (
|
||||
torch.ops.streams.synchronize_device.default,
|
||||
torch.ops.streams.synchronize_stream.default,
|
||||
):
|
||||
all_stream_deps: list[Node] = [
|
||||
n for nodes in stream_to_nodes.values() for n in nodes
|
||||
]
|
||||
if all_stream_deps:
|
||||
found_sync = True
|
||||
_wrap_sync_node(gm, node, all_stream_deps, visited)
|
||||
stream_to_nodes.clear()
|
||||
node = next_node
|
||||
continue
|
||||
|
||||
event_index: int = node.args[0] # type: ignore[assignment]
|
||||
|
||||
# synchronize_event blocks the CPU thread, so it acts
|
||||
# as a barrier across all streams. Collect deps from every
|
||||
# stream and reset them all afterward. If the event was
|
||||
# recorded externally, thread the graph inputs through so
|
||||
# that any post-sync uses depend on the synchronize.
|
||||
if node.target is torch.ops.streams.synchronize_event.default:
|
||||
sync_stream: int | None = event_to_stream.get(event_index)
|
||||
all_stream_deps: list[Node] = [
|
||||
n for nodes in stream_to_nodes.values() for n in nodes
|
||||
]
|
||||
if event_index not in event_to_stream:
|
||||
placeholders = [n for n in graph.nodes if n.op == "placeholder"]
|
||||
deps_before_sync = [*placeholders, *all_stream_deps]
|
||||
else:
|
||||
deps_before_sync = all_stream_deps
|
||||
else:
|
||||
sync_stream = node.args[1] # type: ignore[assignment]
|
||||
deps_before_sync = list(stream_to_nodes.get(sync_stream, ()))
|
||||
# Nodes without explicit stream annotation (custom.stream=None)
|
||||
# run on the current/default stream. Include them when the sync
|
||||
# op references a stream, since the unannotated nodes are
|
||||
# implicitly on that stream.
|
||||
if None in stream_to_nodes and sync_stream is not None:
|
||||
deps_before_sync.extend(stream_to_nodes[None])
|
||||
|
||||
# For wait_event and synchronize_event, add a cross-event
|
||||
# dependency on the matching record_event's control_deps node
|
||||
# so they cannot be reordered before the record.
|
||||
if (
|
||||
node.target
|
||||
in (
|
||||
torch.ops.streams.wait_event.default,
|
||||
torch.ops.streams.synchronize_event.default,
|
||||
)
|
||||
and event_index in event_to_ctrl
|
||||
):
|
||||
deps_before_sync = [
|
||||
event_to_ctrl[event_index],
|
||||
*deps_before_sync,
|
||||
]
|
||||
|
||||
# For synchronize_event, also include the getitem nodes
|
||||
# threaded through record_event's control_deps. This ensures
|
||||
# subsequent ops that depend on recorded values get rewired
|
||||
# through synchronize_event.
|
||||
if (
|
||||
node.target is torch.ops.streams.synchronize_event.default
|
||||
and event_index in event_to_passthrough
|
||||
):
|
||||
deps_before_sync = [
|
||||
*deps_before_sync,
|
||||
*event_to_passthrough[event_index],
|
||||
]
|
||||
|
||||
if deps_before_sync:
|
||||
found_sync = True
|
||||
ctrl_node, passthrough = _wrap_sync_node(
|
||||
gm, node, deps_before_sync, visited
|
||||
)
|
||||
else:
|
||||
ctrl_node = None
|
||||
passthrough: list[torch.fx.Node] = []
|
||||
|
||||
if node.target is torch.ops.streams.record_event.default:
|
||||
event_to_stream[event_index] = sync_stream
|
||||
if ctrl_node is not None:
|
||||
event_to_ctrl[event_index] = ctrl_node
|
||||
event_to_passthrough[event_index] = passthrough
|
||||
|
||||
# Reset: ops between this sync and the next will accumulate
|
||||
# fresh. Ordering with prior ops is already enforced because
|
||||
# their uses were rewired through getitems from control_deps.
|
||||
if node.target is torch.ops.streams.synchronize_event.default:
|
||||
stream_to_nodes.clear()
|
||||
else:
|
||||
stream_to_nodes[sync_stream] = []
|
||||
if None in stream_to_nodes:
|
||||
stream_to_nodes[None] = []
|
||||
elif "val" in node.meta:
|
||||
stream = get_stream(node)
|
||||
stream_to_nodes.setdefault(stream, []).append(node)
|
||||
|
||||
node = next_node
|
||||
|
||||
if found_sync:
|
||||
gm.recompile()
|
||||
+397
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
Codegen for AOTDispatchSubclassWrapper.
|
||||
|
||||
Generates a Python function that replaces the data-driven
|
||||
runtime_unwrap_tensor_subclasses / wrap_tensor_subclasses loop with
|
||||
a straight-line function where all metadata (indices, attr names,
|
||||
subclass types, symint positions) is baked in at compile time.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import keyword
|
||||
import logging
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import torch
|
||||
from torch import SymInt
|
||||
|
||||
from .schemas import OpaqueMeta, PlainTensorMeta, SubclassCreationMeta
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_symint_placeholder(x: None | int | SymInt) -> bool:
|
||||
"""Check whether a size/stride entry is symbolic and needs a runtime value.
|
||||
|
||||
Works both before make_runtime_safe() (entries are SymInt) and after
|
||||
(symbolic entries replaced with None, nested ints with -1).
|
||||
"""
|
||||
if x is None:
|
||||
return True
|
||||
if isinstance(x, SymInt) and not x.node.is_nested_int():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _compute_placeholders(outer: Iterable[None | int | SymInt]) -> list[bool]:
|
||||
return [_is_symint_placeholder(s) for s in outer]
|
||||
|
||||
|
||||
def _safe_attr_access(var: str, attr: str) -> str:
|
||||
if attr.isidentifier() and not keyword.iskeyword(attr):
|
||||
return f"{var}.{attr}"
|
||||
return f"getattr({var}, {attr!r})"
|
||||
|
||||
|
||||
class _CodegenState:
|
||||
"""Accumulates lines of generated source and global bindings."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.lines: list[str] = []
|
||||
self.globals: dict[str, object] = {}
|
||||
self._name_counter: int = 0
|
||||
|
||||
def emit(self, line: str, indent: int = 1) -> None:
|
||||
self.lines.append(" " * indent + line)
|
||||
|
||||
def fresh_name(self, prefix: str) -> str:
|
||||
name = f"{prefix}_{self._name_counter}"
|
||||
self._name_counter += 1
|
||||
return name
|
||||
|
||||
def add_global(self, name: str, value: object) -> str:
|
||||
self.globals[name] = value
|
||||
return name
|
||||
|
||||
|
||||
def _codegen_unwrap_subclass(
|
||||
state: _CodegenState,
|
||||
meta: SubclassCreationMeta,
|
||||
var: str,
|
||||
indent: int = 1,
|
||||
include_symints: bool = True,
|
||||
) -> None:
|
||||
"""Emit code to recursively unwrap a single subclass input."""
|
||||
for attr, attr_meta in meta.attrs.items():
|
||||
match attr_meta:
|
||||
case PlainTensorMeta() | OpaqueMeta():
|
||||
state.emit(
|
||||
f"unwrapped_args.append({_safe_attr_access(var, attr)})",
|
||||
indent=indent,
|
||||
)
|
||||
case SubclassCreationMeta():
|
||||
inner_var = state.fresh_name("_inner")
|
||||
state.emit(
|
||||
f"{inner_var} = {_safe_attr_access(var, attr)}", indent=indent
|
||||
)
|
||||
_codegen_unwrap_subclass(
|
||||
state,
|
||||
attr_meta,
|
||||
inner_var,
|
||||
indent=indent,
|
||||
include_symints=include_symints,
|
||||
)
|
||||
|
||||
# Emit symint extraction
|
||||
if include_symints:
|
||||
size_placeholders = _compute_placeholders(meta.outer_size)
|
||||
stride_placeholders = _compute_placeholders(meta.outer_stride)
|
||||
has_size_symints = any(size_placeholders)
|
||||
has_stride_symints = any(stride_placeholders)
|
||||
|
||||
if has_size_symints or has_stride_symints:
|
||||
size_var = state.fresh_name("_size")
|
||||
state.emit(f"{size_var} = {var}.size()", indent=indent)
|
||||
for i, is_sym in enumerate(size_placeholders):
|
||||
if is_sym:
|
||||
state.emit(f"unwrapped_args.append({size_var}[{i}])", indent=indent)
|
||||
|
||||
stride_var = state.fresh_name("_stride")
|
||||
state.emit(f"{stride_var} = {var}.stride()", indent=indent)
|
||||
for i, is_sym in enumerate(stride_placeholders):
|
||||
if is_sym:
|
||||
state.emit(
|
||||
f"unwrapped_args.append({stride_var}[{i}])", indent=indent
|
||||
)
|
||||
|
||||
|
||||
def _concrete_value(val: None | int | SymInt) -> int:
|
||||
"""Get the concrete int value for a non-symbolic size/stride entry.
|
||||
|
||||
Used for entries that are NOT symbolic placeholders, meaning they are
|
||||
concrete ints or nested ints (represented as -1 after make_runtime_safe).
|
||||
"""
|
||||
if isinstance(val, int):
|
||||
return val
|
||||
# Before make_runtime_safe: nested ints are SymInts; use -1 as dummy.
|
||||
# After make_runtime_safe: they're already -1.
|
||||
if isinstance(val, SymInt) and val.node.is_nested_int():
|
||||
return -1
|
||||
raise AssertionError(f"Expected concrete int, got {type(val)}: {val}")
|
||||
|
||||
|
||||
def _codegen_wrap_subclass(
|
||||
state: _CodegenState,
|
||||
meta: SubclassCreationMeta,
|
||||
out_idx_ref: list[int],
|
||||
) -> str:
|
||||
"""Emit code to reconstruct one subclass output. Returns the variable name."""
|
||||
inner_dict_var = state.fresh_name("_out_inner")
|
||||
entries: list[str] = []
|
||||
|
||||
for attr, attr_meta in meta.attrs.items():
|
||||
match attr_meta:
|
||||
case PlainTensorMeta() | OpaqueMeta():
|
||||
idx = out_idx_ref[0]
|
||||
out_idx_ref[0] += 1
|
||||
entries.append(f"{attr!r}: unwrapped_outs[{idx}]")
|
||||
case SubclassCreationMeta():
|
||||
nested_var = _codegen_wrap_subclass(state, attr_meta, out_idx_ref)
|
||||
entries.append(f"{attr!r}: {nested_var}")
|
||||
|
||||
state.emit(f"{inner_dict_var} = {{{', '.join(entries)}}}")
|
||||
|
||||
# Reconstruct outer_size and outer_stride
|
||||
size_placeholders = _compute_placeholders(meta.outer_size)
|
||||
stride_placeholders = _compute_placeholders(meta.outer_stride)
|
||||
|
||||
def _build_tuple(
|
||||
outer: Iterable[None | int | SymInt], placeholders: list[bool]
|
||||
) -> str:
|
||||
parts: list[str] = []
|
||||
for val, is_sym in zip(outer, placeholders):
|
||||
if is_sym:
|
||||
idx = out_idx_ref[0]
|
||||
out_idx_ref[0] += 1
|
||||
parts.append(f"unwrapped_outs[{idx}]")
|
||||
else:
|
||||
parts.append(repr(_concrete_value(val)))
|
||||
if len(parts) == 1:
|
||||
return f"({parts[0]},)"
|
||||
return f"({', '.join(parts)})"
|
||||
|
||||
size_expr = _build_tuple(meta.outer_size, size_placeholders)
|
||||
stride_expr = _build_tuple(meta.outer_stride, stride_placeholders)
|
||||
|
||||
type_name = state.add_global(
|
||||
state.fresh_name("_subclass_type"),
|
||||
meta.original_subclass_type or type(meta.original_subclass),
|
||||
)
|
||||
meta_name = state.add_global(state.fresh_name("_meta"), meta.meta)
|
||||
|
||||
result_var = state.fresh_name("_out")
|
||||
state.emit(
|
||||
f"{result_var} = {type_name}.__tensor_unflatten__("
|
||||
f"{inner_dict_var}, {meta_name}, {size_expr}, {stride_expr})"
|
||||
)
|
||||
return result_var
|
||||
|
||||
|
||||
def _emit_output_wrapping(
|
||||
state: _CodegenState,
|
||||
out_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
) -> tuple[list[str], int]:
|
||||
"""Emit wrapping code for output metas.
|
||||
|
||||
Returns (result_exprs, num_args_tallied) where result_exprs are Python
|
||||
expression strings referencing each wrapped output.
|
||||
"""
|
||||
out_idx_ref = [0]
|
||||
result_exprs: list[str] = []
|
||||
num_args_tallied = 0
|
||||
|
||||
for meta in out_metas:
|
||||
if isinstance(meta, PlainTensorMeta):
|
||||
result_exprs.append(f"unwrapped_outs[{meta.unwrapped_idx}]")
|
||||
num_args_tallied += 1
|
||||
out_idx_ref[0] = max(out_idx_ref[0], meta.unwrapped_idx + 1)
|
||||
else:
|
||||
result_var = _codegen_wrap_subclass(state, meta, out_idx_ref)
|
||||
result_exprs.append(result_var)
|
||||
num_args_tallied += meta.arg_count
|
||||
|
||||
return result_exprs, num_args_tallied
|
||||
|
||||
|
||||
def _emit_input_unwrapping(
|
||||
state: _CodegenState,
|
||||
inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
frozen_inp_indices: frozenset[int] = frozenset(),
|
||||
include_symints: bool = True,
|
||||
) -> None:
|
||||
"""Emit unwrapping code for input metas into unwrapped_args.
|
||||
|
||||
Caller must have already emitted ``unwrapped_args = []``.
|
||||
"""
|
||||
for i, meta in enumerate(inp_metas):
|
||||
if isinstance(meta, PlainTensorMeta):
|
||||
state.emit(f"unwrapped_args.append(args[{i}])")
|
||||
elif i in frozen_inp_indices:
|
||||
# Frozen by inductor freezing: constant already baked into graph.
|
||||
state.emit("unwrapped_args.append(None)")
|
||||
else:
|
||||
inp_var = state.fresh_name("_inp")
|
||||
type_name = state.add_global(
|
||||
state.fresh_name("_expected_type"),
|
||||
meta.original_subclass_type or type(meta.original_subclass),
|
||||
)
|
||||
state.emit(f"{inp_var} = args[{i}]")
|
||||
state.emit(
|
||||
f"assert type({inp_var}) is {type_name}, "
|
||||
f"f'expected {{{type_name}}}, got {{type({inp_var})}}'",
|
||||
)
|
||||
_codegen_unwrap_subclass(
|
||||
state, meta, inp_var, indent=1, include_symints=include_symints
|
||||
)
|
||||
|
||||
|
||||
def _codegen_subclass_wrapper_source(
|
||||
inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
out_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
num_fw_outs_saved_for_bw: int | None,
|
||||
frozen_inp_indices: frozenset[int] = frozenset(),
|
||||
act_input_indices: list[int] | None = None,
|
||||
) -> tuple[str, dict[str, object]]:
|
||||
"""Generate source and globals for a subclass wrapper.
|
||||
|
||||
Returns (source, globals_dict). The globals_dict will NOT contain
|
||||
``compiled_fn`` — the caller is responsible for adding it before exec.
|
||||
"""
|
||||
state = _CodegenState()
|
||||
|
||||
state.emit("def inner_fn(args):", indent=0)
|
||||
|
||||
# --- Resolve AsyncCollectiveTensors ---
|
||||
# ACTs are transient eager-mode wrappers for async collective overlap.
|
||||
# Inductor triton kernels bypass __torch_dispatch__, so we must call
|
||||
# trigger_wait() before the compiled graph uses the data.
|
||||
if act_input_indices:
|
||||
for i in act_input_indices:
|
||||
state.emit(f"args[{i}] = args[{i}].trigger_wait()")
|
||||
|
||||
# --- Input unwrapping ---
|
||||
state.emit("unwrapped_args = []")
|
||||
_emit_input_unwrapping(state, inp_metas, frozen_inp_indices=frozen_inp_indices)
|
||||
|
||||
# Pass through any trailing args not covered by inp_metas
|
||||
# (e.g. rng seed/offset added by FunctionalizedRngRuntimeWrapper).
|
||||
num_inp_metas = len(inp_metas)
|
||||
state.emit(f"unwrapped_args.extend(args[{num_inp_metas}:])")
|
||||
state.emit("args.clear()")
|
||||
|
||||
# --- Call compiled function ---
|
||||
state.emit("unwrapped_outs = compiled_fn(unwrapped_args)")
|
||||
|
||||
# --- Output wrapping ---
|
||||
result_exprs, num_args_tallied = _emit_output_wrapping(state, out_metas)
|
||||
result_tuple = f"({', '.join(result_exprs)},)" if result_exprs else "()"
|
||||
if num_fw_outs_saved_for_bw is not None:
|
||||
state.emit(
|
||||
f"return {result_tuple} + tuple(unwrapped_outs[{num_args_tallied}:])"
|
||||
)
|
||||
else:
|
||||
state.emit(f"return {result_tuple}")
|
||||
|
||||
source = "\n".join(state.lines)
|
||||
return source, state.globals
|
||||
|
||||
|
||||
def _codegen_subclass_wrap_source(
|
||||
out_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
) -> tuple[str, dict[str, object]]:
|
||||
"""Generate source for wrapping flat outputs into subclasses.
|
||||
|
||||
Used for the backward epilogue. Shares output-wrapping logic with
|
||||
_codegen_subclass_wrapper_source via _emit_output_wrapping.
|
||||
"""
|
||||
state = _CodegenState()
|
||||
state.emit("def wrap_fn(unwrapped_outs):", indent=0)
|
||||
result_exprs, _ = _emit_output_wrapping(state, out_metas)
|
||||
result_tuple = f"({', '.join(result_exprs)},)" if result_exprs else "()"
|
||||
state.emit(f"return {result_tuple}")
|
||||
source = "\n".join(state.lines)
|
||||
return source, state.globals
|
||||
|
||||
|
||||
def _compile_and_exec_source(
|
||||
source: str,
|
||||
globals_dict: dict[str, object],
|
||||
fn_name: str,
|
||||
artifact_name: str,
|
||||
wrapped_fn: Callable[..., object] | None = None,
|
||||
) -> Callable[..., object]:
|
||||
"""Compile generated source, exec it, and return the named function.
|
||||
|
||||
If wrapped_fn is provided, applies functools.update_wrapper so that
|
||||
__wrapped__ and __dict__ (e.g. _fx_graph_cache_key) propagate to the
|
||||
generated function.
|
||||
"""
|
||||
if log.isEnabledFor(logging.DEBUG):
|
||||
log.debug("Generated %s:\n%s", artifact_name, source)
|
||||
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": artifact_name,
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: source,
|
||||
)
|
||||
|
||||
code = compile(source, f"<{artifact_name}>", "exec")
|
||||
local_dict: dict[str, object] = {}
|
||||
exec(code, globals_dict, local_dict) # noqa: S102
|
||||
fn = local_dict[fn_name]
|
||||
if wrapped_fn is not None:
|
||||
functools.update_wrapper(fn, wrapped_fn) # type: ignore[arg-type]
|
||||
return fn # type: ignore[return-value]
|
||||
|
||||
|
||||
def codegen_backward_subclass_fns(
|
||||
grad_input_metas: list[PlainTensorMeta | SubclassCreationMeta] | None = None,
|
||||
) -> tuple[Callable[..., object], Callable[..., object] | None]:
|
||||
"""Generate codegen'd unwrap and wrap functions for the backward pass.
|
||||
|
||||
Returns (unwrap_fn, wrap_fn). unwrap_fn is used by the backward prologue
|
||||
to unwrap non-tangent subclass inputs (always an identity in AOT dispatch
|
||||
since the compiled forward operates on unwrapped inner tensors). wrap_fn
|
||||
is used by the backward epilogue to wrap flat grad inputs back into
|
||||
subclasses; it is None when grad_input_metas is None.
|
||||
"""
|
||||
source = "def unwrap_fn(args):\n return list(args)"
|
||||
globals_dict: dict[str, object] = {}
|
||||
unwrap_fn = _compile_and_exec_source(
|
||||
source, globals_dict, "unwrap_fn", "backward_subclass_unwrap"
|
||||
)
|
||||
|
||||
wrap_fn = None
|
||||
if grad_input_metas is not None:
|
||||
wrap_source, wrap_globals = _codegen_subclass_wrap_source(grad_input_metas)
|
||||
wrap_fn = _compile_and_exec_source(
|
||||
wrap_source, wrap_globals, "wrap_fn", "backward_subclass_wrapper"
|
||||
)
|
||||
|
||||
return unwrap_fn, wrap_fn
|
||||
|
||||
|
||||
def codegen_subclass_wrapper(
|
||||
compiled_fn: Callable[..., object],
|
||||
inp_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
out_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
num_fw_outs_saved_for_bw: int | None,
|
||||
frozen_inp_indices: frozenset[int] = frozenset(),
|
||||
act_input_indices: list[int] | None = None,
|
||||
) -> Callable[..., object]:
|
||||
"""Generate a specialized wrapper function for subclass unwrap/wrap."""
|
||||
source, globals_dict = _codegen_subclass_wrapper_source(
|
||||
inp_metas,
|
||||
out_metas,
|
||||
num_fw_outs_saved_for_bw,
|
||||
frozen_inp_indices,
|
||||
act_input_indices=act_input_indices,
|
||||
)
|
||||
globals_dict["compiled_fn"] = compiled_fn
|
||||
return _compile_and_exec_source(
|
||||
source, globals_dict, "inner_fn", "subclass_wrapper", wrapped_fn=compiled_fn
|
||||
)
|
||||
+138
@@ -0,0 +1,138 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import itertools
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._library.opaque_object import is_opaque_reference_type
|
||||
from torch._opaque_base import OpaqueBase
|
||||
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
|
||||
|
||||
from .schemas import OpaqueMeta
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
# This is technically very similar to SubclassCreatingMeta
|
||||
# in aot_autograd, but we don't need all the stuff in there
|
||||
# so just recreated a new dataclass.
|
||||
@dataclasses.dataclass
|
||||
class SubclassCreationMeta:
|
||||
start_idx: int
|
||||
num_tensors: int
|
||||
class_type: Any
|
||||
# None means the attr is a plain tensor (base case of recursion)
|
||||
attrs: dict[str, SubclassCreationMeta | OpaqueMeta | None]
|
||||
metadata: Any
|
||||
outer_size: Iterable[None | int | torch.SymInt]
|
||||
outer_stride: Iterable[None | int | torch.SymInt]
|
||||
|
||||
|
||||
class UnwrapTensorSubclass(torch.nn.Module):
|
||||
def forward(self, *tensors) -> torch.Tensor: # type: ignore[no-untyped-def]
|
||||
todo: list[torch.Tensor | OpaqueBase] = list(tensors)
|
||||
|
||||
def _unwrap_tensor_subclasses(subclass_meta, tensors, offset): # type: ignore[no-untyped-def]
|
||||
if subclass_meta is None:
|
||||
return tensors[offset], offset + 1
|
||||
inner_tensors = {}
|
||||
for attr, meta in subclass_meta.attrs.items():
|
||||
if isinstance(meta, OpaqueMeta):
|
||||
inner_tensors[attr] = tensors[offset]
|
||||
offset += 1
|
||||
else:
|
||||
built_tensor, offset = _unwrap_tensor_subclasses(
|
||||
meta, tensors, offset
|
||||
)
|
||||
inner_tensors[attr] = built_tensor
|
||||
rebuilt = subclass_meta.class_type.__tensor_unflatten__(
|
||||
inner_tensors,
|
||||
subclass_meta.metadata,
|
||||
subclass_meta.outer_size,
|
||||
subclass_meta.outer_stride,
|
||||
)
|
||||
return rebuilt, offset
|
||||
|
||||
return _unwrap_tensor_subclasses(self.subclass_meta, todo, 0)[0]
|
||||
|
||||
def right_inverse(self, tensor: torch.Tensor) -> list[torch.Tensor | OpaqueBase]:
|
||||
if type(tensor) is torch.Tensor:
|
||||
raise AssertionError("tensor must be a subclass, not torch.Tensor")
|
||||
plain_tensors: list[torch.Tensor | OpaqueBase] = []
|
||||
|
||||
def _create_subclass_meta(tensor, idx, plain_tensor_container): # type: ignore[no-untyped-def]
|
||||
if type(tensor) is torch.Tensor:
|
||||
plain_tensor_container.append(tensor)
|
||||
return None, idx + 1
|
||||
inner_tensors_attrnames, metadata = tensor.__tensor_flatten__() # type: ignore[attr-defined]
|
||||
new_idx = idx
|
||||
attr_to_meta: dict[str, SubclassCreationMeta | OpaqueMeta | None] = {}
|
||||
for attr in inner_tensors_attrnames:
|
||||
val = getattr(tensor, attr)
|
||||
match val:
|
||||
case OpaqueBase():
|
||||
if not is_opaque_reference_type(type(val)):
|
||||
raise ValueError(
|
||||
f"{type(val).__name__!r} found in tensor attrs of "
|
||||
f"{type(tensor).__name__}.__tensor_flatten__(). "
|
||||
"Only tensors and reference-type opaques are allowed "
|
||||
"in tensor attrs."
|
||||
)
|
||||
attr_to_meta[attr] = OpaqueMeta()
|
||||
plain_tensor_container.append(val)
|
||||
new_idx += 1
|
||||
case torch.Tensor():
|
||||
subclass_meta, new_idx = _create_subclass_meta(
|
||||
val, new_idx, plain_tensor_container
|
||||
)
|
||||
attr_to_meta[attr] = subclass_meta
|
||||
case _:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(val)}"
|
||||
)
|
||||
return (
|
||||
SubclassCreationMeta(
|
||||
start_idx=idx,
|
||||
num_tensors=new_idx - idx,
|
||||
class_type=type(tensor),
|
||||
attrs=attr_to_meta,
|
||||
metadata=metadata,
|
||||
outer_size=tensor.size(),
|
||||
outer_stride=tensor.stride(),
|
||||
),
|
||||
new_idx,
|
||||
)
|
||||
|
||||
self.subclass_meta = _create_subclass_meta(tensor, 0, plain_tensors)[0]
|
||||
return plain_tensors
|
||||
|
||||
|
||||
def unwrap_tensor_subclass_parameters(module: torch.nn.Module) -> torch.nn.Module:
|
||||
"""
|
||||
Model transformation that replaces all the parameters that are subclasses to plain tensors.
|
||||
This reduces runtime overhead of flattening/unflattening the parameters.
|
||||
|
||||
This transformation adds parametrization with `torch.nn.utils.parametrize`.
|
||||
The FQNs of the subclass parameters will be changed and state_dict will become incompatible with the original model.
|
||||
E.g.
|
||||
Original model state_dict: {"p1": torch.testing._internal.TwoTensor}
|
||||
becomes: {"parametrizations.p2.original0": torch.Tensor, "parametrizations.p2.original1": torch.Tensor}
|
||||
|
||||
"""
|
||||
for name, tensor in itertools.chain(
|
||||
list(module.named_parameters(recurse=False)),
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
list(module.named_buffers(recurse=False)),
|
||||
):
|
||||
if is_traceable_wrapper_subclass(tensor):
|
||||
torch.nn.utils.parametrize.register_parametrization(
|
||||
module, name, UnwrapTensorSubclass()
|
||||
)
|
||||
|
||||
for child in module.children():
|
||||
unwrap_tensor_subclass_parameters(child)
|
||||
|
||||
return module
|
||||
+655
@@ -0,0 +1,655 @@
|
||||
"""
|
||||
This file contains utilities for tracing through __torch_dispatch__ based tensor subclasses and modes.
|
||||
AOTAutograd's responsibility is to trace through all pytorch capabilities that live in the pytorch dispatcher,
|
||||
and this includes tensor subclasses that implement __torch_dispatch__.
|
||||
"""
|
||||
|
||||
import collections
|
||||
import typing
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
from typing import Any, TypeGuard, TypeVar
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch import SymInt, Tensor
|
||||
from torch._library.fake_class_registry import maybe_unwrap_fake_script_object
|
||||
from torch._library.opaque_object import is_opaque_reference_type
|
||||
from torch._opaque_base import OpaqueBase
|
||||
from torch._subclasses.fake_tensor import get_plain_tensors
|
||||
from torch.types import IntLikeType
|
||||
from torch.utils._python_dispatch import is_traceable_wrapper_subclass
|
||||
|
||||
from .descriptors import (
|
||||
AOTInput,
|
||||
AOTOutput,
|
||||
DummyAOTInput,
|
||||
SubclassGetAttrAOTInput,
|
||||
SubclassGetAttrAOTOutput,
|
||||
SubclassSizeAOTInput,
|
||||
SubclassSizeAOTOutput,
|
||||
SubclassStrideAOTInput,
|
||||
SubclassStrideAOTOutput,
|
||||
)
|
||||
from .schemas import (
|
||||
FakifiedFlatArgs,
|
||||
FxValue,
|
||||
MutationType,
|
||||
OpaqueMeta,
|
||||
PlainTensorMeta,
|
||||
SubclassCreationMeta,
|
||||
ViewAndMutationMeta,
|
||||
)
|
||||
from .utils import strict_zip
|
||||
|
||||
|
||||
zip = strict_zip
|
||||
|
||||
T = TypeVar("T", bound=torch.Tensor)
|
||||
|
||||
|
||||
def requires_subclass_dispatch(
|
||||
args: FakifiedFlatArgs, fw_metadata: ViewAndMutationMeta
|
||||
) -> bool:
|
||||
args_flattened = pytree.arg_tree_leaves(*args)
|
||||
any_subclass_args = any(
|
||||
is_traceable_wrapper_subclass(x)
|
||||
for x in args_flattened
|
||||
if isinstance(x, Tensor)
|
||||
)
|
||||
from torch._functorch._aot_autograd.schemas import SubclassCreationMeta
|
||||
|
||||
any_subclass_outputs = any(
|
||||
type(x) is SubclassCreationMeta for x in fw_metadata.subclass_fw_graph_out_meta
|
||||
)
|
||||
# This tells us whether or not we need to perform any unwrapping/wrapping of tensor subclasses at runtime.
|
||||
return bool(any_subclass_args or any_subclass_outputs)
|
||||
|
||||
|
||||
from .schemas import MemoryFormatMeta
|
||||
|
||||
|
||||
def maybe_suggest_memory_format(
|
||||
t: Tensor, with_memory_format: bool
|
||||
) -> MemoryFormatMeta | None:
|
||||
if not with_memory_format:
|
||||
return None
|
||||
|
||||
return MemoryFormatMeta.from_tensor(t)
|
||||
|
||||
|
||||
def get_subclass_typing_container(
|
||||
tensor_subclass: torch.Tensor,
|
||||
) -> dict[type[torch.Tensor], list[type[torch.Tensor]]]:
|
||||
"""
|
||||
Given a subclass, returns a recursive dictionary mapping each
|
||||
inner tensors to its' subclass types.
|
||||
"""
|
||||
|
||||
def _get_types_for_subclass(tensor_subclass: torch.Tensor) -> None:
|
||||
if not is_traceable_wrapper_subclass(tensor_subclass):
|
||||
return
|
||||
tracker[type(tensor_subclass)].append(tensor_subclass)
|
||||
inner_keys, _ = tensor_subclass.__tensor_flatten__()
|
||||
for key in inner_keys:
|
||||
match getattr(tensor_subclass, key):
|
||||
case torch.Tensor() as inner_value:
|
||||
_get_types_for_subclass(inner_value)
|
||||
case OpaqueBase():
|
||||
pass
|
||||
case unexpected:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(unexpected)}"
|
||||
)
|
||||
|
||||
tracker: dict[Any, list[Any]] = collections.defaultdict(list)
|
||||
_get_types_for_subclass(tensor_subclass)
|
||||
return tracker
|
||||
|
||||
|
||||
def create_subclass_metadata(
|
||||
a: Any,
|
||||
start_idx: int,
|
||||
count_symints: bool,
|
||||
with_memory_format: bool = False,
|
||||
) -> tuple[Any, int]:
|
||||
if not is_traceable_wrapper_subclass(a):
|
||||
idx = start_idx + 1
|
||||
return (
|
||||
PlainTensorMeta(
|
||||
idx,
|
||||
memory_format=maybe_suggest_memory_format(a, with_memory_format),
|
||||
),
|
||||
idx,
|
||||
)
|
||||
|
||||
inner_keys, metadata = a.__tensor_flatten__()
|
||||
new_start_idx = start_idx
|
||||
attrs: dict[str, SubclassCreationMeta | PlainTensorMeta | OpaqueMeta] = {}
|
||||
|
||||
for key in inner_keys:
|
||||
inner_value = getattr(a, key)
|
||||
match inner_value:
|
||||
case OpaqueBase():
|
||||
# During tracing, opaques are wrapped in FakeScriptObject;
|
||||
# unwrap to check the real type.
|
||||
real_type = type(maybe_unwrap_fake_script_object(inner_value))
|
||||
if not is_opaque_reference_type(real_type):
|
||||
raise RuntimeError(
|
||||
f"{real_type.__name__!r} found in tensor attrs of "
|
||||
f"{type(a).__name__}.__tensor_flatten__(). "
|
||||
"Only tensors and reference-type opaques are allowed "
|
||||
"in tensor attrs."
|
||||
)
|
||||
attrs[key] = OpaqueMeta()
|
||||
new_start_idx += 1
|
||||
case Tensor():
|
||||
new_subclass_meta, new_start_idx = create_subclass_metadata(
|
||||
inner_value,
|
||||
new_start_idx,
|
||||
count_symints=count_symints,
|
||||
with_memory_format=with_memory_format,
|
||||
)
|
||||
attrs[key] = new_subclass_meta
|
||||
case _:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(inner_value)}"
|
||||
)
|
||||
|
||||
# It *must* be because is_traceable_wrapper_subclass() - but mypy is not smart.
|
||||
if not isinstance(a, Tensor):
|
||||
raise AssertionError(f"expected Tensor, got {type(a)}")
|
||||
|
||||
new_start_idx = (
|
||||
new_start_idx
|
||||
+ count_symints * len(enumerate_filter_symints(a.size()))
|
||||
+ count_symints * len(enumerate_filter_symints(a.stride()))
|
||||
)
|
||||
|
||||
return (
|
||||
SubclassCreationMeta(
|
||||
flat_tensor_start_idx=start_idx,
|
||||
arg_count=new_start_idx - start_idx,
|
||||
included_subclass_symints=count_symints,
|
||||
attrs=attrs,
|
||||
meta=metadata,
|
||||
outer_size=a.size(), # type: ignore[attr-defined, arg-type]
|
||||
outer_stride=a.stride(), # type: ignore[arg-type]
|
||||
original_subclass=a,
|
||||
memory_format=maybe_suggest_memory_format(a, with_memory_format),
|
||||
),
|
||||
new_start_idx,
|
||||
)
|
||||
|
||||
|
||||
# Given a flat list of arguments, some of which may be tensor subclasses,
|
||||
# computes metadata about "how to reconstruct the current list of subclasses,
|
||||
# if we were given their flattened dense tensors instead"
|
||||
def create_subclass_meta(
|
||||
curr_args: list[Any] | tuple[Any, ...],
|
||||
*,
|
||||
count_symints: bool = True,
|
||||
with_memory_format: bool = False,
|
||||
) -> list[PlainTensorMeta | SubclassCreationMeta]:
|
||||
idx = 0
|
||||
infos: list[PlainTensorMeta | SubclassCreationMeta] = []
|
||||
for a in curr_args:
|
||||
if is_traceable_wrapper_subclass(a):
|
||||
if not isinstance(a, Tensor):
|
||||
raise AssertionError(
|
||||
f"expected Tensor for traceable wrapper subclass, got {type(a)}"
|
||||
)
|
||||
start_idx = idx
|
||||
subclass_meta, _ = create_subclass_metadata(
|
||||
a,
|
||||
start_idx,
|
||||
count_symints=count_symints,
|
||||
with_memory_format=with_memory_format,
|
||||
)
|
||||
infos.append(subclass_meta)
|
||||
cnt = subclass_meta.arg_count
|
||||
else:
|
||||
infos.append(
|
||||
PlainTensorMeta(
|
||||
idx,
|
||||
memory_format=maybe_suggest_memory_format(a, with_memory_format),
|
||||
)
|
||||
)
|
||||
cnt = 1
|
||||
idx += cnt
|
||||
return infos
|
||||
|
||||
|
||||
def enumerate_filter_symints(lst: Iterable[IntLikeType]) -> list[tuple[int, SymInt]]:
|
||||
# Capture all SymInts from the iterable.
|
||||
def symint_check(s: IntLikeType) -> TypeGuard[SymInt]:
|
||||
return isinstance(s, SymInt) and not s.node.is_nested_int()
|
||||
|
||||
return [(i, s) for i, s in enumerate(lst) if symint_check(s)]
|
||||
|
||||
|
||||
def compute_symint_placeholders(lst: Iterable[None | int | SymInt]) -> list[bool]:
|
||||
# Non-nested symints are replaced with None in `make_runtime_safe()`
|
||||
return [s is None for s in lst]
|
||||
|
||||
|
||||
# Intended to make it easier to define function that is
|
||||
# either (AOTInput -> AOTInput) or (AOTOutput -> AOTOutput)
|
||||
# but not the other combos
|
||||
AOTDescriptor = TypeVar("AOTDescriptor", AOTInput, AOTOutput)
|
||||
|
||||
|
||||
# This function takes in a pytree of arguments and unwraps any tensor
|
||||
# subclasses.
|
||||
#
|
||||
# NOTE: The reason for "append_symints":
|
||||
#
|
||||
# * At compile time: we append extra symint args when unwrapping primals
|
||||
# (but not tangents, because they should always share symints with primals).
|
||||
# We also append extra symints when unwrapping the subclass outputs of the
|
||||
# traced function, so we can return them as extra outputs
|
||||
#
|
||||
# * At runtime: we similarly append subclass sizes when we unwrap subclass
|
||||
# primals (but not tangents) on entry to the forward. See the runtime version of
|
||||
# this function below.
|
||||
def unwrap_tensor_subclasses(
|
||||
wrapped_args: list[FxValue],
|
||||
wrapped_args_descs: Sequence[AOTDescriptor],
|
||||
*,
|
||||
append_symints: bool,
|
||||
) -> tuple[list[FxValue], list[AOTDescriptor]]:
|
||||
def _maybe_fakeify_opaque(v: Any) -> Any:
|
||||
# Registered opaque types need to be wrapped as FakeScriptObject for
|
||||
# compile-time FX tracing (proxy slot tracking, hashability, etc.).
|
||||
if isinstance(v, OpaqueBase):
|
||||
from torch._guards import detect_fake_mode
|
||||
from torch._library.fake_class_registry import maybe_to_fake_obj
|
||||
from torch._library.opaque_object import is_opaque_type
|
||||
|
||||
fake_mode = detect_fake_mode()
|
||||
if fake_mode is not None and is_opaque_type(type(v)):
|
||||
return maybe_to_fake_obj(fake_mode, v)
|
||||
return v
|
||||
|
||||
def flatten_subclass(
|
||||
t: FxValue,
|
||||
desc: AOTDescriptor,
|
||||
*,
|
||||
out: tuple[list[FxValue], list[AOTDescriptor]],
|
||||
) -> None:
|
||||
# unwrap a subclass into plain tensors and their size/stride if "append_symint"
|
||||
# is True
|
||||
if not is_traceable_wrapper_subclass(t):
|
||||
out[0].append(_maybe_fakeify_opaque(t))
|
||||
out[1].append(desc)
|
||||
return
|
||||
|
||||
attrs, _ = t.__tensor_flatten__()
|
||||
|
||||
SubclassGetAttr: Callable[[AOTInput | AOTOutput, str], AOTDescriptor]
|
||||
SubclassSize: Callable[[AOTInput | AOTOutput, int], AOTDescriptor]
|
||||
SubclassStride: Callable[[AOTInput | AOTOutput, int], AOTDescriptor]
|
||||
if isinstance(desc, AOTInput):
|
||||
SubclassGetAttr = SubclassGetAttrAOTInput # type: ignore[bad-assignment]
|
||||
SubclassSize = SubclassSizeAOTInput # type: ignore[bad-assignment]
|
||||
SubclassStride = SubclassStrideAOTInput # type: ignore[bad-assignment]
|
||||
else:
|
||||
SubclassGetAttr = SubclassGetAttrAOTOutput # type: ignore[bad-assignment]
|
||||
SubclassSize = SubclassSizeAOTOutput # type: ignore[bad-assignment]
|
||||
SubclassStride = SubclassStrideAOTOutput # type: ignore[bad-assignment]
|
||||
|
||||
for attr in attrs:
|
||||
inner_value = getattr(t, attr)
|
||||
n_desc: Any = SubclassGetAttr(desc, attr)
|
||||
flatten_subclass(inner_value, n_desc, out=out)
|
||||
|
||||
if append_symints:
|
||||
sizes = enumerate_filter_symints(t.size())
|
||||
strides = enumerate_filter_symints(t.stride())
|
||||
out[0].extend(s for _, s in sizes)
|
||||
out[0].extend(s for _, s in strides)
|
||||
out[1].extend(SubclassSize(desc, i) for i, _ in sizes)
|
||||
out[1].extend(SubclassStride(desc, i) for i, _ in strides)
|
||||
|
||||
xs_inner: list[FxValue] = []
|
||||
descs_inner: list[AOTDescriptor] = []
|
||||
|
||||
for x, desc in zip(wrapped_args, wrapped_args_descs):
|
||||
flatten_subclass(x, desc, out=(xs_inner, descs_inner))
|
||||
|
||||
return xs_inner, descs_inner
|
||||
|
||||
|
||||
# subclass_metas is needed at runtime to compute which indices are symints in
|
||||
# the outer_size/outer_stride
|
||||
def runtime_unwrap_tensor_subclasses(
|
||||
wrapped_args: list[Tensor | int],
|
||||
*,
|
||||
append_symints: bool,
|
||||
subclass_metas: list[PlainTensorMeta | SubclassCreationMeta] | None = None,
|
||||
) -> list[int | Tensor | SymInt | OpaqueBase]:
|
||||
def flatten_subclass(
|
||||
x: Tensor,
|
||||
subclass_meta: PlainTensorMeta | SubclassCreationMeta | OpaqueMeta | None,
|
||||
*,
|
||||
out: list[OpaqueBase | SymInt | Tensor | int],
|
||||
) -> list[OpaqueBase | SymInt | Tensor | int]:
|
||||
if not is_traceable_wrapper_subclass(x):
|
||||
out.append(x)
|
||||
return out
|
||||
|
||||
if not isinstance(x, Tensor):
|
||||
raise AssertionError(f"expected Tensor, got {type(x)}")
|
||||
if not isinstance(subclass_meta, SubclassCreationMeta):
|
||||
raise AssertionError("subclass_meta should be a SubclassCreationMeta")
|
||||
|
||||
attrs, _ = x.__tensor_flatten__()
|
||||
|
||||
for attr in attrs:
|
||||
inner_value = getattr(x, attr)
|
||||
match inner_value:
|
||||
case OpaqueBase():
|
||||
out.append(inner_value)
|
||||
case Tensor():
|
||||
inner_meta = subclass_meta.attrs.get(attr)
|
||||
flatten_subclass(inner_value, inner_meta, out=out)
|
||||
case _:
|
||||
raise AssertionError(
|
||||
f"expected Tensor or OpaqueBase, got {type(inner_value)}"
|
||||
)
|
||||
|
||||
if append_symints:
|
||||
# outer_size
|
||||
size = x.size()
|
||||
symint_placeholders = compute_symint_placeholders(subclass_meta.outer_size)
|
||||
if len(size) != len(symint_placeholders):
|
||||
raise AssertionError(
|
||||
f"size length mismatch: {len(size)} != {len(symint_placeholders)}"
|
||||
)
|
||||
out.extend(
|
||||
[r for (r, is_symint) in zip(size, symint_placeholders) if is_symint]
|
||||
)
|
||||
|
||||
# outer_stride
|
||||
stride = x.stride()
|
||||
symint_placeholders = compute_symint_placeholders(
|
||||
subclass_meta.outer_stride
|
||||
)
|
||||
if len(stride) != len(symint_placeholders):
|
||||
raise AssertionError(
|
||||
f"stride length mismatch: {len(stride)} != {len(symint_placeholders)}"
|
||||
)
|
||||
out.extend(
|
||||
[r for (r, is_symint) in zip(stride, symint_placeholders) if is_symint]
|
||||
)
|
||||
return out
|
||||
|
||||
xs_inner: list[int | Tensor | SymInt | OpaqueBase] = []
|
||||
|
||||
if append_symints:
|
||||
if subclass_metas is None:
|
||||
raise AssertionError(
|
||||
"subclass_metas must not be None when append_symints is True"
|
||||
)
|
||||
|
||||
for idx, x in enumerate(wrapped_args):
|
||||
if not is_traceable_wrapper_subclass(x):
|
||||
xs_inner.append(x)
|
||||
continue
|
||||
|
||||
if subclass_metas is None:
|
||||
get_plain_tensors(typing.cast(Tensor, x), out=xs_inner)
|
||||
else:
|
||||
subclass_meta = subclass_metas[idx]
|
||||
if not isinstance(subclass_meta, SubclassCreationMeta):
|
||||
raise AssertionError(
|
||||
f"expected SubclassCreationMeta, got {type(subclass_meta)}"
|
||||
)
|
||||
flatten_subclass(typing.cast(Tensor, x), subclass_meta, out=xs_inner)
|
||||
|
||||
return xs_inner
|
||||
|
||||
|
||||
def unwrap_tensor_subclasses_with_indices_to_original(
|
||||
wrapped_args: list[Any],
|
||||
) -> tuple[list[Any], list[int]]:
|
||||
ret_unwrapped = []
|
||||
ret_indices_to_original = []
|
||||
for i, a in enumerate(wrapped_args):
|
||||
a_unwrapped, _ = unwrap_tensor_subclasses(
|
||||
[a], [DummyAOTInput(9999)], append_symints=False
|
||||
)
|
||||
ret_unwrapped.extend(a_unwrapped)
|
||||
n = len(a_unwrapped)
|
||||
ret_indices_to_original.extend([i] * n)
|
||||
|
||||
return ret_unwrapped, ret_indices_to_original
|
||||
|
||||
|
||||
def remap_unwrapped_subclass_arg_indices(
|
||||
wrapped_args: list[Any], static_input_indices: list[int]
|
||||
) -> list[int]:
|
||||
static_input_indices_set = set(static_input_indices)
|
||||
new_ind = 0
|
||||
remapped_static_indices = []
|
||||
for i, arg in enumerate(wrapped_args):
|
||||
num_indices = 1
|
||||
if is_traceable_wrapper_subclass(arg):
|
||||
num_indices = (
|
||||
len(get_plain_tensors(typing.cast(Tensor, arg), out=[]))
|
||||
+ len(enumerate_filter_symints(arg.size()))
|
||||
+ len(enumerate_filter_symints(arg.stride()))
|
||||
)
|
||||
|
||||
for _ in range(num_indices):
|
||||
if i in static_input_indices_set:
|
||||
remapped_static_indices.append(new_ind)
|
||||
|
||||
new_ind += 1
|
||||
|
||||
return remapped_static_indices
|
||||
|
||||
|
||||
# Turns a flattened list of tensor arguments into (maybe) subclass tensors.
|
||||
# This function is used both at trace time and runtime, so we have an is_runtime flag telling us which context we're in.
|
||||
def wrap_tensor_subclasses(
|
||||
unwrapped_args: Sequence[Any],
|
||||
*,
|
||||
subclass_metas: list[PlainTensorMeta | SubclassCreationMeta],
|
||||
num_fw_outs_saved_for_bw: int | None = None,
|
||||
included_subclass_symints: bool = False,
|
||||
is_runtime: bool = False,
|
||||
make_subclass_override: Callable[..., Any] | None = None,
|
||||
) -> tuple[Any, ...]:
|
||||
# pyrefly: ignore [implicit-any]
|
||||
wrapped_args = []
|
||||
num_args_tallied = 0
|
||||
for subclass_meta in subclass_metas:
|
||||
if isinstance(subclass_meta, PlainTensorMeta):
|
||||
wrapped_args.append(unwrapped_args[subclass_meta.unwrapped_idx])
|
||||
num_args_tallied += 1
|
||||
else:
|
||||
if not isinstance(subclass_meta, SubclassCreationMeta):
|
||||
raise AssertionError(
|
||||
f"expected SubclassCreationMeta, got {type(subclass_meta)}"
|
||||
)
|
||||
if subclass_meta.included_subclass_symints != included_subclass_symints:
|
||||
raise AssertionError(
|
||||
f"included_subclass_symints mismatch: {subclass_meta.included_subclass_symints} != {included_subclass_symints}"
|
||||
)
|
||||
|
||||
if make_subclass_override:
|
||||
wrapped_args.append(
|
||||
make_subclass_override(subclass_meta, is_runtime, unwrapped_args)
|
||||
)
|
||||
else:
|
||||
wrapped_args.append(
|
||||
subclass_meta.creation_fn(
|
||||
unwrapped_args,
|
||||
is_runtime=is_runtime,
|
||||
)
|
||||
)
|
||||
num_args_tallied += subclass_meta.arg_count
|
||||
|
||||
# Note: [Partitioner handling for Subclasses, Part 2]
|
||||
# At the beginning of AOTAutograd, we collect metadata on the inputs and outputs of the user fw,
|
||||
# to figure out which inputs/outputs are subclasses, and how to reconstruct the subclasses after flattening them.
|
||||
#
|
||||
# When this function is called at runtime in the forward,
|
||||
# we have been passed a list of (flattened) dense-tensor fw-outs, and need to reconstruct any subclass fw outs.
|
||||
#
|
||||
# One reasonable question that you should ask: when should the dense_tensor -> subclass_tensor wrapping happen?
|
||||
# Answer: we do it **inside of our compiled autograd.Function**.
|
||||
# This seems like morally the right place: autograd happens above subclass desugaring,
|
||||
# so autograd should see actual tensor subclasses at runtime, and not flattened dense tensors.
|
||||
#
|
||||
# This causes a tricky interaction though: when we run the min-cut partitioner to divvy up the joint graph
|
||||
# into a forward and backward graph, we end up with some activations that show up as extra outputs
|
||||
# in the compiled forward graph, that are **not** user outputs.
|
||||
# These activations are not visible to the user, and so there's no need for us to wrap them back into subclasses.
|
||||
#
|
||||
# On top of that, when we first computed subclass metadata (in `run_functionalized_fw_and_collect_metadata`),
|
||||
# we computed subclass metadata on every forward output, but this did **not** include activations
|
||||
# created by the partitioner.
|
||||
# as a result, `unwrapped_args` here will correspond to (*unwrapped_user_fw_outs, *activations),
|
||||
# but `subclass_metas` will only correspond to subclass metadata on `user_fw_outs`.
|
||||
# We then need to make sure that we return (*wrapped_user_fw_outs, *activations).
|
||||
if num_fw_outs_saved_for_bw is not None:
|
||||
if len(unwrapped_args) != num_args_tallied + num_fw_outs_saved_for_bw:
|
||||
raise AssertionError(
|
||||
f"Expected the number actual unwrapped-subclass outputs {len(unwrapped_args)} to equal "
|
||||
f"the number of args calculated from subclasses ({num_args_tallied}) plus the number of "
|
||||
f"additional activations saved for the backward pass ({num_fw_outs_saved_for_bw})"
|
||||
)
|
||||
activations = unwrapped_args[num_args_tallied:]
|
||||
if isinstance(wrapped_args, tuple) and isinstance(activations, tuple):
|
||||
return wrapped_args + activations
|
||||
return tuple(list(wrapped_args) + list(activations))
|
||||
else:
|
||||
if len(unwrapped_args) != num_args_tallied:
|
||||
raise AssertionError(
|
||||
f"Expected {len(unwrapped_args)} == {num_args_tallied}"
|
||||
)
|
||||
return tuple(wrapped_args)
|
||||
|
||||
|
||||
# Given a bunch of "dense" tensor arguments, this function (potentially) wraps them into tensor subclasses.
|
||||
# This function carefully handles the inference vs. joint cases:
|
||||
# - when is_joint_structure is True, args is (primals, tangents)
|
||||
# - when is_joint_structure is False, args is [*primals]
|
||||
def wrap_tensor_subclasses_maybe_joint(
|
||||
unwrapped_args: Sequence[Any],
|
||||
*,
|
||||
is_joint_structure: bool,
|
||||
meta: ViewAndMutationMeta,
|
||||
) -> tuple[Any, ...]:
|
||||
# Since this function is reused for both inference and joint graphs,
|
||||
if is_joint_structure:
|
||||
if not (isinstance(unwrapped_args, tuple) and len(unwrapped_args) == 2):
|
||||
unwrapped_len = (
|
||||
len(unwrapped_args)
|
||||
if isinstance(unwrapped_args, (tuple, list))
|
||||
else "N/A"
|
||||
)
|
||||
raise AssertionError(
|
||||
f"expected tuple of length 2 for joint structure, "
|
||||
f"got {type(unwrapped_args)} with length {unwrapped_len}"
|
||||
)
|
||||
if not (
|
||||
isinstance(unwrapped_args[0], (tuple, list))
|
||||
and isinstance(unwrapped_args[1], (tuple, list))
|
||||
):
|
||||
raise AssertionError(
|
||||
f"expected primals and tangents to be tuple or list, got {type(unwrapped_args[0])} and {type(unwrapped_args[1])}"
|
||||
)
|
||||
primals, tangents = unwrapped_args[0], unwrapped_args[1]
|
||||
wrapped_primals = wrap_tensor_subclasses(
|
||||
primals,
|
||||
subclass_metas=meta.subclass_inp_meta,
|
||||
included_subclass_symints=True,
|
||||
)
|
||||
wrapped_tangents = wrap_tensor_subclasses(
|
||||
tangents,
|
||||
subclass_metas=meta.subclass_tangent_meta,
|
||||
included_subclass_symints=False,
|
||||
)
|
||||
return (wrapped_primals, wrapped_tangents)
|
||||
else:
|
||||
wrapped_args = wrap_tensor_subclasses(
|
||||
unwrapped_args,
|
||||
subclass_metas=meta.subclass_inp_meta,
|
||||
included_subclass_symints=True,
|
||||
)
|
||||
return wrapped_args
|
||||
|
||||
|
||||
def compute_inner_mutated_inp_indices_from_subclass_meta(
|
||||
fw_metadata: ViewAndMutationMeta,
|
||||
inner_metadata: ViewAndMutationMeta,
|
||||
) -> list[int]:
|
||||
# Note: [Recomputing subclass mutation handling]
|
||||
#
|
||||
# Generally, if a subclass requires grad, its components will not require grad.
|
||||
# But for the purposes of tracking returned tensors, we should treat those component
|
||||
# tensors as if they require grad.
|
||||
#
|
||||
# For example, if the subclass tensor requires grad and will be mutated in a way that
|
||||
# requires us to handle the mutation outside of the graph, we need to return it
|
||||
# from the forward graph. The inner_meta data won't consider the component tensors
|
||||
# as if they need to be returned, because they don't require grad; but really, we
|
||||
# should handle those tensors the same way we handle the subclass tensor itself; i.e.
|
||||
# if we'd include the subclass tensor as part of the outputs, then we should also
|
||||
# include the component tensors.
|
||||
#
|
||||
# To do this, we patch num_mutated_inp_runtime_indices below by expanding the inputs
|
||||
# from the outer subclass tensors and propagating
|
||||
|
||||
updated_input_info = []
|
||||
inner_idx = 0
|
||||
if not fw_metadata.subclass_inp_meta:
|
||||
# Sometimes we don't have subclass info, e.g. synthetic_base codepaths
|
||||
return inner_metadata.mutated_inp_runtime_indices
|
||||
if len(fw_metadata.subclass_inp_meta) != len(fw_metadata.input_info):
|
||||
raise AssertionError(
|
||||
f"subclass_inp_meta length ({len(fw_metadata.subclass_inp_meta)}) != input_info length ({len(fw_metadata.input_info)})"
|
||||
)
|
||||
for outer_idx, inp_meta in enumerate(fw_metadata.subclass_inp_meta):
|
||||
if isinstance(inp_meta, PlainTensorMeta):
|
||||
if outer_idx >= len(fw_metadata.input_info):
|
||||
raise AssertionError(
|
||||
f"outer_idx ({outer_idx}) >= len(fw_metadata.input_info) ({len(fw_metadata.input_info)})"
|
||||
)
|
||||
if inner_metadata is not None:
|
||||
if inner_idx >= len(inner_metadata.input_info):
|
||||
raise AssertionError(
|
||||
f"inner_idx ({inner_idx}) >= len(inner_metadata.input_info) ({len(inner_metadata.input_info)})"
|
||||
)
|
||||
if (
|
||||
inner_metadata.input_info[inner_idx]
|
||||
!= fw_metadata.input_info[outer_idx]
|
||||
):
|
||||
raise AssertionError(
|
||||
f"input_info mismatch at inner_idx={inner_idx}, outer_idx={outer_idx}: "
|
||||
f"{inner_metadata.input_info[inner_idx]} != {fw_metadata.input_info[outer_idx]}"
|
||||
)
|
||||
updated_input_info.append(fw_metadata.input_info[outer_idx])
|
||||
inner_idx += 1
|
||||
else:
|
||||
if inp_meta.original_subclass is None:
|
||||
raise AssertionError(
|
||||
"inp_meta.original_subclass must not be None for SubclassCreationMeta"
|
||||
)
|
||||
for _ in range(inp_meta.arg_count):
|
||||
updated_input_info.append(fw_metadata.input_info[outer_idx])
|
||||
inner_idx += 1
|
||||
if inner_metadata is not None:
|
||||
if len(inner_metadata.input_info) != len(updated_input_info):
|
||||
raise AssertionError(
|
||||
f"inner_metadata.input_info length ({len(inner_metadata.input_info)}) "
|
||||
f"!= updated_input_info length ({len(updated_input_info)})"
|
||||
)
|
||||
|
||||
return [
|
||||
i
|
||||
for i, inp in enumerate(updated_input_info)
|
||||
if inp.mutation_type == MutationType.MUTATED_OUT_GRAPH
|
||||
]
|
||||
@@ -0,0 +1,835 @@
|
||||
"""
|
||||
Contains various utils for AOTAutograd, including those for handling collections.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import dataclasses
|
||||
import logging
|
||||
import operator
|
||||
import warnings
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import nullcontext
|
||||
from functools import partial, wraps
|
||||
from typing import Any, overload, TYPE_CHECKING
|
||||
from typing_extensions import ParamSpec, TypeVar, TypeVarTuple, Unpack
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._library.fake_class_registry import FakeScriptObject
|
||||
from torch._library.opaque_object import is_opaque_value
|
||||
from torch._logging import getArtifactLogger
|
||||
from torch._subclasses.fake_tensor import FakeTensor
|
||||
from torch._subclasses.functional_tensor import FunctionalTensor
|
||||
from torch.fx.experimental._backward_state import BackwardState
|
||||
from torch.fx.experimental.proxy_tensor import py_sym_types
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
if TYPE_CHECKING:
|
||||
from .schemas import AOTConfig, ViewAndMutationMeta
|
||||
|
||||
|
||||
KNOWN_TYPES = [
|
||||
torch.Tensor,
|
||||
BackwardState,
|
||||
int,
|
||||
str,
|
||||
float,
|
||||
bool,
|
||||
type(None),
|
||||
*py_sym_types,
|
||||
FakeScriptObject,
|
||||
torch.ScriptObject,
|
||||
]
|
||||
|
||||
aot_graphs_effects_log = getArtifactLogger(__name__, "aot_graphs_effects")
|
||||
annotation_log = getArtifactLogger(__name__, "annotation")
|
||||
|
||||
strict_zip = partial(zip, strict=True)
|
||||
|
||||
|
||||
def partial_flatten_asdict(obj: object) -> Any:
|
||||
if dataclasses.is_dataclass(obj):
|
||||
return {
|
||||
field.name: getattr(obj, field.name) for field in dataclasses.fields(obj)
|
||||
}
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
return obj.__class__([partial_flatten_asdict(item) for item in obj])
|
||||
elif isinstance(obj, dict):
|
||||
return {k: partial_flatten_asdict(v) for k, v in obj.items()}
|
||||
else:
|
||||
return obj
|
||||
|
||||
|
||||
@overload
|
||||
def normalize_as_list(x: _T) -> list[_T]: ...
|
||||
@overload
|
||||
def normalize_as_list(x: tuple[_T, ...]) -> list[_T]: ...
|
||||
@overload
|
||||
def normalize_as_list(x: list[_T]) -> list[_T]: ...
|
||||
|
||||
|
||||
def normalize_as_list(x: object) -> list[object]:
|
||||
if isinstance(x, tuple):
|
||||
return list(x)
|
||||
elif isinstance(x, list):
|
||||
return x
|
||||
return [x]
|
||||
|
||||
|
||||
def _get_autocast_states() -> list[Any]:
|
||||
return [
|
||||
torch.is_autocast_enabled("cuda"),
|
||||
torch.is_autocast_enabled("cpu"),
|
||||
torch.get_autocast_dtype("cuda"),
|
||||
torch.get_autocast_dtype("cpu"),
|
||||
torch.is_autocast_cache_enabled(),
|
||||
]
|
||||
|
||||
|
||||
def make_boxed_func(f: Callable[..., Any]) -> Callable[[list[Any]], Any]:
|
||||
@simple_wraps(f)
|
||||
def g(args: list[Any]) -> Any:
|
||||
return f(*args)
|
||||
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
g._boxed_call = True
|
||||
return g
|
||||
|
||||
|
||||
def make_boxed_compiler(
|
||||
compiler: Callable[..., Any],
|
||||
) -> Callable[..., Any]:
|
||||
@wraps(compiler)
|
||||
def f(fx_g: Any, inps: Any) -> Any:
|
||||
out_f = compiler(fx_g, inps)
|
||||
fx_g = make_boxed_func(out_f)
|
||||
return fx_g
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def call_func_at_runtime_with_args(
|
||||
f: Callable[..., Any],
|
||||
args: Sequence[Any],
|
||||
steal_args: bool = False,
|
||||
disable_amp: bool = False,
|
||||
) -> list[Any]:
|
||||
if not steal_args:
|
||||
args = list(args)
|
||||
if not isinstance(args, list):
|
||||
raise AssertionError(f"args must be a list, got {type(args)}")
|
||||
|
||||
context = torch._C._DisableAutocast if disable_amp else nullcontext
|
||||
with context():
|
||||
if getattr(f, "_boxed_call", False):
|
||||
out = normalize_as_list(f(args))
|
||||
else:
|
||||
# TODO: Please remove soon
|
||||
# https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670
|
||||
warnings.warn(
|
||||
"Your compiler for AOTAutograd is returning a function that doesn't take boxed arguments. "
|
||||
"Please wrap it with functorch.compile.make_boxed_func or handle the boxed arguments yourself. "
|
||||
"See https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670 for rationale.",
|
||||
stacklevel=2,
|
||||
)
|
||||
out = normalize_as_list(f(*args))
|
||||
return out
|
||||
|
||||
|
||||
# Inspired by autodidax (thanks!)
|
||||
class PytreeThunk:
|
||||
spec: pytree.TreeSpec | None = None
|
||||
# These are some kinda dumb microoptimizations that save about 3-4 us of overhead.
|
||||
is_simple: bool | None = (
|
||||
None # if the output spec is a tuple/list, we won't bother unflattening it.
|
||||
)
|
||||
is_really_simple: bool | None = None # if the output spec is a LeafSpec
|
||||
|
||||
def set(self, spec: pytree.TreeSpec) -> None:
|
||||
if not (self.spec is None or self.spec == spec):
|
||||
raise AssertionError(f"spec mismatch: existing={self.spec}, new={spec}")
|
||||
if spec is None:
|
||||
raise AssertionError("spec must not be None")
|
||||
self.spec: pytree.TreeSpec = spec
|
||||
if self.spec.type in {tuple, list} and all(
|
||||
child.is_leaf() for child in spec.children()
|
||||
):
|
||||
self.is_simple = True
|
||||
if self.spec.is_leaf():
|
||||
self.is_really_simple = True
|
||||
|
||||
def unflatten(self, x: Sequence[Any]) -> Any:
|
||||
if self.is_really_simple:
|
||||
return x[0]
|
||||
if self.is_simple:
|
||||
return x
|
||||
if self.spec is None:
|
||||
raise AssertionError("spec must be set before calling unflatten")
|
||||
return pytree.tree_unflatten(x, self.spec)
|
||||
|
||||
|
||||
# Creates a function that returns flattened inputs and outputs
|
||||
# Also returns the output tree spec, which is needed to recover the "unflattened"
|
||||
# output tree structure later.
|
||||
def create_tree_flattened_fn(
|
||||
fn: Callable[..., Any],
|
||||
args: Sequence[Any],
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[Callable[..., list[Any]], PytreeThunk]:
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
# Save the args_spec for flat_tensor_args to unflatten while tracing
|
||||
_, tensor_args_spec = pytree.tree_flatten((args, kwargs))
|
||||
out_spec = PytreeThunk()
|
||||
|
||||
def flat_fn(*flat_args: Any) -> list[Any]:
|
||||
# The input are flattened tensor args. Prepare the args in the
|
||||
# order that original function expects. Add static args as well.
|
||||
# They will appear as tensor constants in the traced graph.
|
||||
nonlocal out_spec
|
||||
args, kwargs = pytree.tree_unflatten(flat_args, tensor_args_spec)
|
||||
tree_out = fn(*args, **kwargs)
|
||||
flat_out, spec = pytree.tree_flatten(tree_out)
|
||||
for i in flat_out:
|
||||
is_known_type = isinstance(i, tuple(KNOWN_TYPES)) or is_opaque_value(i)
|
||||
if not is_known_type:
|
||||
raise RuntimeError(
|
||||
f"Found {type(i)} in output, which is not a known type. "
|
||||
"If this type holds tensors, you need to register a pytree for it. "
|
||||
"See https://github.com/pytorch/functorch/issues/475 for a brief "
|
||||
"explanation why. If you don't need to register a pytree, please "
|
||||
"leave a comment explaining your use case and we'll make this more "
|
||||
"ergonomic to deal with"
|
||||
)
|
||||
out_spec.set(spec)
|
||||
return flat_out
|
||||
|
||||
# Can't use functools.wraps here because the wrapper has different
|
||||
# calling convention
|
||||
if hasattr(fn, "_orig_mod"):
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
flat_fn._orig_mod = fn._orig_mod
|
||||
|
||||
return flat_fn, out_spec
|
||||
|
||||
|
||||
# This function takes in a tensor t, and returns one of t, t.view(), or t.clone().
|
||||
# When tracing the joint forward + backward, for any inputs in the graph that are mutated,
|
||||
# we need to clone them first (and similarly for metadata-only mutations, we need to view them first).
|
||||
# The idea is that when we trace the backward, we need to pass in the *original* primals
|
||||
# to autograd.grad(), before they were mutated.
|
||||
# Note: when we have synthetic base inputs, we need to clone them *before* creating views off of them.
|
||||
# This means that "idx" here represents the index of the (potentially) synthetic base.
|
||||
# What we need to do is:
|
||||
# (1) map the current (post-synthetic-base calling convention) input argument index
|
||||
# to int index pre-synthetic-base-calling-convention.
|
||||
# (2) There could be multiple, if this index corresponds to a synthetic base
|
||||
# that has multiple input aliases.
|
||||
# (3) If any of those corresponding inputs get metadata mutations, then we clone the base.
|
||||
def maybe_to_fresh_input(idx: int, t: Any, meta: "ViewAndMutationMeta") -> Any:
|
||||
if not isinstance(t, torch.Tensor):
|
||||
return t
|
||||
if idx in meta.mutated_inp_runtime_indices:
|
||||
# We only need to bother cloning mutated inputs that participate in autograd.
|
||||
if meta.input_info[idx].requires_grad and meta.input_info[idx].mutates_data:
|
||||
# Make sure the primal we pass to autograd.grad()
|
||||
# sees the tensor before the mutation
|
||||
return t.clone()
|
||||
if meta.input_info[idx] and meta.input_info[idx].mutates_metadata:
|
||||
# Make sure the primal we pass to autograd.grad()
|
||||
# sees the tensor before the metadata mutation
|
||||
return t.view(t.shape)
|
||||
return t
|
||||
|
||||
|
||||
def is_with_effects(node: torch.fx.Node) -> bool:
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.higher_order.with_effects
|
||||
):
|
||||
return True
|
||||
elif (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.higher_order.invoke_subgraph
|
||||
):
|
||||
# Check if subgraph has effects by looking in the cache
|
||||
from torch._guards import InvokeSubgraphCache, TracingContext
|
||||
|
||||
tracing_ctx = TracingContext.try_get()
|
||||
if tracing_ctx:
|
||||
invoke_subgraph_cache = tracing_ctx.hop_dispatch_set_cache.get_cache(
|
||||
torch.ops.higher_order.invoke_subgraph
|
||||
)
|
||||
if invoke_subgraph_cache:
|
||||
if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
|
||||
raise AssertionError(
|
||||
f"expected InvokeSubgraphCache, got {type(invoke_subgraph_cache)}"
|
||||
)
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
effects = invoke_subgraph_cache.get_effects(node.args[1])
|
||||
return effects is not None
|
||||
return False
|
||||
|
||||
|
||||
def unlift_tokens(
|
||||
fw_module: torch.fx.GraphModule,
|
||||
fw_metadata: "ViewAndMutationMeta",
|
||||
aot_config: "AOTConfig",
|
||||
bw_module: torch.fx.GraphModule | None = None,
|
||||
) -> None:
|
||||
# Remove the tokens from the inputs/outputs of the graph since inductor does
|
||||
# not want these extra inputs/outputs, and replace them with
|
||||
# _make_token() to create a token, and _sink_tokens() to collect the
|
||||
# tokens. See Note [Side-Effectful Tokens in AOTAutograd]
|
||||
# Logic:
|
||||
# 1. In the case of with_effects:
|
||||
# Before:
|
||||
# ```
|
||||
# def forward(self, token, arg1_1):
|
||||
# with_effects = torch.ops.higher_order.with_effects(token, ...)
|
||||
# getitem = with_effects[0]
|
||||
# getitem_1 = with_effects[0]
|
||||
# return (getitem, getitem_1)
|
||||
# ```
|
||||
#
|
||||
# After:
|
||||
# ```
|
||||
# def forward(self, arg1_1):
|
||||
# _make_token_default = torch.ops.prims._make_token.default()
|
||||
# with_effects = torch.ops.higher_order.with_effects(_make_token_default, ...)
|
||||
# getitem = with_effects[0]
|
||||
# getitem_1 = with_effects[0]
|
||||
# _sink_tokens_default = torch.ops.prims._sink_tokens.default([getitem]);
|
||||
# return (getitem_1,)
|
||||
# ```
|
||||
#
|
||||
# 2. In the case of an invoke_subgraph node, we will use the
|
||||
# InvokeSubgraphCache to determine if the subgraph has effects. Then we will
|
||||
# turn it into a `with_effects` node. This is so that at the toplevel graph,
|
||||
# the nodes will have the correct with_effects threading. We will apply this
|
||||
# pass recursively to submodules so the tokens will be removed from the
|
||||
# subgraph's inputs.
|
||||
#
|
||||
# Before:
|
||||
# ```
|
||||
# def forward(self, token, arg1_1):
|
||||
# repeated_subgraph0 = self.repeated_subgraph0
|
||||
# invoke_subgraph = torch.ops.higher_order.invoke_subgraph(
|
||||
# repeated_subgraph0, 'subgraph_0', token, x, arg1_1)
|
||||
# getitem = invoke_subgraph[0]
|
||||
# getitem_1 = invoke_subgraph[1]
|
||||
# return (getitem, getitem1)
|
||||
# ```
|
||||
#
|
||||
# After:
|
||||
# ```
|
||||
# def forward(self, arg1_1):
|
||||
# _make_token_default = torch.ops.prims._make_token.default()
|
||||
# repeated_subgraph0 = self.repeated_subgraph0
|
||||
# with_effects_1 = torch.ops.higher_order.with_effects(
|
||||
# _make_token_default, torch.ops.higher_order.invoke_subgraph,
|
||||
# repeated_subgraph0, 'subgraph_0', arg1_1)
|
||||
# getitem = with_effects_1[0]
|
||||
# getitem_1 = with_effects_1[1]; with_effects_1 = None
|
||||
# _sink_tokens_default = torch.ops.prims._sink_tokens.default([getitem])
|
||||
# return (getitem_1,)
|
||||
# ```
|
||||
#
|
||||
# 3. The toplevel module should have the following invariants:
|
||||
# forward:
|
||||
# expected_num_erased_inputs == len(fw_metadata.tokens)
|
||||
# expected_num_erased_outputs == len(fw_metadata.tokens)
|
||||
# backward:
|
||||
# expected_num_erased_inputs == fw_metadata.num_backward_tokens
|
||||
# expected_num_erased_outputs == fw_metadata.num_backward_tokens
|
||||
num_forward_tokens = len(fw_metadata.tokens)
|
||||
num_backward_tokens = fw_metadata.num_backward_tokens
|
||||
|
||||
def replace_input_token_with_make_token(
|
||||
module: torch.fx.GraphModule, node: torch.fx.Node
|
||||
) -> None:
|
||||
with module.graph.inserting_before(node):
|
||||
new_token_node = module.graph.call_function(
|
||||
torch.ops.prims._make_token.default, ()
|
||||
)
|
||||
new_token_node.meta["val"] = torch.tensor([])
|
||||
new_token_node.meta["tensor_meta"] = torch.tensor([])
|
||||
node.replace_all_uses_with(new_token_node)
|
||||
module.graph.erase_node(node)
|
||||
|
||||
def get_output_tokens(node: torch.fx.Node) -> set[torch.fx.Node]:
|
||||
output_tokens = set()
|
||||
for user in list(node.users.keys()):
|
||||
# Check if this is a getitem accessing index 0 (the token)
|
||||
if (
|
||||
user.op == "call_function"
|
||||
and user.target is operator.getitem
|
||||
and len(user.args) > 1
|
||||
and user.args[1] == 0
|
||||
):
|
||||
# Check if this getitem is used in an output
|
||||
for user_user in list(user.users.keys()):
|
||||
if user_user.op == "output":
|
||||
output_tokens.add(user)
|
||||
return output_tokens
|
||||
|
||||
def _unlift_tokens_from_module_helper(
|
||||
module: torch.fx.GraphModule,
|
||||
subgraph_str: str,
|
||||
expected_num_erased: int | None,
|
||||
) -> None:
|
||||
input_token_nodes = set()
|
||||
output_token_nodes = set()
|
||||
|
||||
for node in module.graph.nodes:
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.higher_order.with_effects
|
||||
):
|
||||
if node.args[0].op == "placeholder":
|
||||
input_token_nodes.add(node.args[0])
|
||||
replace_input_token_with_make_token(module, node.args[0])
|
||||
|
||||
tokens_from_with_effects = get_output_tokens(node)
|
||||
output_token_nodes = output_token_nodes | tokens_from_with_effects
|
||||
|
||||
elif (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.higher_order.invoke_subgraph
|
||||
):
|
||||
subgraph_node, identifier, *operands = node.args
|
||||
|
||||
# Check if subgraph has effects by looking in the cache
|
||||
from torch._guards import InvokeSubgraphCache, TracingContext
|
||||
|
||||
effects = None
|
||||
tracing_ctx = TracingContext.try_get()
|
||||
if tracing_ctx:
|
||||
invoke_subgraph_cache = (
|
||||
tracing_ctx.hop_dispatch_set_cache.get_cache(
|
||||
torch.ops.higher_order.invoke_subgraph
|
||||
)
|
||||
)
|
||||
if invoke_subgraph_cache:
|
||||
if not isinstance(invoke_subgraph_cache, InvokeSubgraphCache):
|
||||
raise AssertionError(
|
||||
f"expected InvokeSubgraphCache, got {type(invoke_subgraph_cache)}"
|
||||
)
|
||||
effects = invoke_subgraph_cache.get_effects(identifier)
|
||||
|
||||
if effects is not None:
|
||||
# Wrap invoke_subgraph with with_effects
|
||||
# Before: invoke_subgraph(subgraph, id, token, *args) -> (token_out, result)
|
||||
# After: with_effects(token, invoke_subgraph, subgraph, id, *args) -> (token_out, result)
|
||||
#
|
||||
# Note: The subgraph itself will be unlifted separately when we iterate
|
||||
# through named_modules() below.
|
||||
|
||||
num_tokens = len(effects)
|
||||
if num_tokens != 1:
|
||||
raise AssertionError(
|
||||
f"Multiple token subgraph NYI, got {num_tokens} tokens"
|
||||
)
|
||||
token_args = operands[:num_tokens]
|
||||
non_token_args = operands[num_tokens:]
|
||||
|
||||
# Create with_effects wrapper around invoke_subgraph
|
||||
# with_effects(token, op, *args) where op is invoke_subgraph
|
||||
# Pass the subgraph and non-token args to invoke_subgraph
|
||||
with module.graph.inserting_before(node):
|
||||
new_node = module.graph.call_function(
|
||||
torch.ops.higher_order.with_effects,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
(
|
||||
token_args[0], # pyrefly: ignore[bad-argument-type]
|
||||
torch.ops.higher_order.invoke_subgraph,
|
||||
subgraph_node,
|
||||
identifier,
|
||||
*tuple(non_token_args),
|
||||
),
|
||||
)
|
||||
node.replace_all_uses_with(new_node)
|
||||
new_node.meta = node.meta
|
||||
module.graph.erase_node(node)
|
||||
|
||||
for token in token_args:
|
||||
if token.op == "placeholder":
|
||||
input_token_nodes.add(token)
|
||||
replace_input_token_with_make_token(module, token)
|
||||
|
||||
# Get output tokens from the new with_effects node
|
||||
tokens_from_invoke_subgraph = get_output_tokens(new_node)
|
||||
output_token_nodes = (
|
||||
output_token_nodes | tokens_from_invoke_subgraph
|
||||
)
|
||||
|
||||
if not output_token_nodes and not input_token_nodes:
|
||||
return
|
||||
|
||||
output_node = next(reversed(module.graph.find_nodes(op="output")))
|
||||
if output_node is None:
|
||||
raise AssertionError("output node not found in graph")
|
||||
with module.graph.inserting_before(output_node):
|
||||
module.graph.call_function(
|
||||
torch.ops.prims._sink_tokens.default,
|
||||
(list(output_token_nodes),),
|
||||
)
|
||||
new_out_args = tuple(
|
||||
[out for out in output_node.args[0] if out not in output_token_nodes]
|
||||
)
|
||||
output_node.args = (new_out_args,)
|
||||
|
||||
if expected_num_erased:
|
||||
if len(input_token_nodes) != expected_num_erased:
|
||||
raise AssertionError(
|
||||
f"{subgraph_str} num_erased_inputs:{len(input_token_nodes)} "
|
||||
f"{input_token_nodes} != expected {expected_num_erased} \n"
|
||||
f"{fw_module.print_readable(print_output=False)}"
|
||||
)
|
||||
if len(output_token_nodes) != expected_num_erased:
|
||||
raise AssertionError(
|
||||
f"{subgraph_str} num_erased_outs:{len(output_token_nodes)} "
|
||||
f"{output_token_nodes} != expected {expected_num_erased} \n"
|
||||
f"{fw_module.print_readable(print_output=False)}"
|
||||
)
|
||||
|
||||
module.recompile()
|
||||
|
||||
def unlift_tokens_from_module(
|
||||
module: torch.fx.GraphModule, subgraph_str: str, expected_num_erased: int
|
||||
) -> None:
|
||||
for name, m in module.named_modules():
|
||||
if isinstance(m, torch.fx.GraphModule):
|
||||
if name == "":
|
||||
_unlift_tokens_from_module_helper(
|
||||
m, subgraph_str, expected_num_erased
|
||||
)
|
||||
else:
|
||||
# Subgraph -- we may or may not have effects applied
|
||||
_unlift_tokens_from_module_helper(m, f"{subgraph_str}_{name}", None)
|
||||
|
||||
if num_forward_tokens > 0:
|
||||
if aot_config.enable_log:
|
||||
from torch._dynamo.utils import lazy_format_graph_code
|
||||
|
||||
aot_graphs_effects_log.debug(
|
||||
"%s",
|
||||
lazy_format_graph_code(
|
||||
"Forward graph before unlifting tokens",
|
||||
fw_module,
|
||||
aot_config.aot_id,
|
||||
include_stride=True,
|
||||
include_device=True,
|
||||
colored=True,
|
||||
),
|
||||
)
|
||||
unlift_tokens_from_module(
|
||||
fw_module,
|
||||
"forward",
|
||||
num_forward_tokens,
|
||||
)
|
||||
|
||||
if bw_module is not None and num_backward_tokens > 0:
|
||||
if aot_config.enable_log:
|
||||
from torch._dynamo.utils import lazy_format_graph_code
|
||||
|
||||
aot_graphs_effects_log.debug(
|
||||
"%s",
|
||||
lazy_format_graph_code(
|
||||
"Backward graph before unlifting tokens",
|
||||
bw_module,
|
||||
aot_config.aot_id,
|
||||
include_stride=True,
|
||||
include_device=True,
|
||||
colored=True,
|
||||
),
|
||||
)
|
||||
unlift_tokens_from_module(bw_module, "backward", num_backward_tokens)
|
||||
|
||||
# This is sad, but we need to update the metadata to get rid of
|
||||
# the tokens.
|
||||
fw_metadata.tokens = {}
|
||||
fw_metadata.num_backward_tokens = 0
|
||||
|
||||
|
||||
def root_module_when_exporting_non_strict(
|
||||
flat_fn: Callable[..., Any],
|
||||
) -> torch.nn.Module | None:
|
||||
# When exporting in non-strict mode, we wrap the root module in a specific pattern.
|
||||
# See `_aot_export_non_strict` in torch.export._trace.py.
|
||||
# We look for that wrapping pattern here.
|
||||
if hasattr(flat_fn, "_orig_mod") and hasattr(flat_fn._orig_mod, "_export_root"):
|
||||
return flat_fn._orig_mod._export_root
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _is_forward_node_with_seq_nr(node: torch.fx.Node) -> bool:
|
||||
# For now, assume that if nn_module_stack_metadata is populated, this
|
||||
# node is from the forward. Ignore nodes without `seq_nr`.
|
||||
# TODO(future): there is likely a less brittle way to do this by walking
|
||||
# the descendants of graph inputs corresponding to fwd inputs, didn't
|
||||
# seem obvious at first glance on how to partition graph inputs into
|
||||
# fwd vs bwd without relying on string names.
|
||||
return node.meta.get("partitioner_tag") != "is_backward" and "seq_nr" in node.meta
|
||||
|
||||
|
||||
def _is_backward_node_with_seq_nr(node: torch.fx.Node) -> bool:
|
||||
# For now, assume that if nn_module_stack_metadata is not populated,
|
||||
# this node is from the backward. Ignore nodes without `seq_nr`.
|
||||
# TODO(future): there is likely a less brittle way to do this, same
|
||||
# as with the forward.
|
||||
return node.meta.get("partitioner_tag") == "is_backward" and "seq_nr" in node.meta
|
||||
|
||||
|
||||
def _collect_fwd_nodes_from_subgraph(
|
||||
fx_g: torch.fx.GraphModule, fwd_seq_nr_to_node: dict[str, torch.fx.Node]
|
||||
) -> None:
|
||||
"""Collect forward nodes from a single subgraph into the global mapping."""
|
||||
for node in fx_g.graph.nodes:
|
||||
if not _is_forward_node_with_seq_nr(node):
|
||||
continue
|
||||
seq_nr = node.meta["seq_nr"]
|
||||
if seq_nr in fwd_seq_nr_to_node:
|
||||
# If we already saw an op with the current `seq_nr`, that means
|
||||
# that the current op did not create an autograd node, and there
|
||||
# is no corresponding backward node, so we skip.
|
||||
continue
|
||||
fwd_seq_nr_to_node[seq_nr] = node
|
||||
|
||||
|
||||
def _copy_metadata_to_bw_nodes_in_subgraph(
|
||||
fx_g: torch.fx.GraphModule, fwd_seq_nr_to_node: dict[str, torch.fx.Node]
|
||||
) -> None:
|
||||
"""Copy metadata from forward nodes to backward nodes in a single subgraph."""
|
||||
for node in fx_g.graph.nodes:
|
||||
annotation_log.debug("node: %s", node.name)
|
||||
seq_nr = node.meta.get("seq_nr")
|
||||
annotation_log.debug("seq_nr: %s", seq_nr)
|
||||
|
||||
if not _is_backward_node_with_seq_nr(node):
|
||||
continue
|
||||
|
||||
# We exclude gradient accumulation nodes from copying tags
|
||||
if node.meta.get("is_gradient_acc", False):
|
||||
annotation_log.debug("is_gradient_acc")
|
||||
continue
|
||||
|
||||
# fwd_node should always exist, but handle non-existence just in case
|
||||
fwd_node = fwd_seq_nr_to_node.get(node.meta["seq_nr"])
|
||||
if fwd_node is not None:
|
||||
node.meta["fwd_nn_module_stack"] = fwd_node.meta.get("nn_module_stack")
|
||||
node.meta["fwd_source_fn_stack"] = fwd_node.meta.get("source_fn_stack")
|
||||
# TODO: better to change to a specific field of custom?
|
||||
custom = fwd_node.meta.get("custom")
|
||||
if custom is not None:
|
||||
node.meta["custom"] = copy.deepcopy(custom)
|
||||
|
||||
|
||||
def copy_fwd_metadata_to_bw_nodes(fx_g: torch.fx.GraphModule) -> None:
|
||||
"""
|
||||
Input: `fx_g` which contains the joint fwd+bwd FX graph created by
|
||||
aot_autograd.
|
||||
|
||||
This function walks the graph and copies over metadata from forward nodes
|
||||
to backward nodes, using the `seq_nr` field as a one-to-many mapping
|
||||
from forward node to backward node. This metadata is useful for performance
|
||||
profiling and debugging.
|
||||
|
||||
This function supports matching forward and backward nodes across different
|
||||
subgraphs (e.g., in recursive submodules from HOPs), enabling backward nodes
|
||||
in any submodule to match forward nodes in any submodule.
|
||||
"""
|
||||
|
||||
# Build a global mapping of seq_nr to forward nodes across all subgraphs
|
||||
fwd_seq_nr_to_node: dict[str, torch.fx.Node] = {}
|
||||
|
||||
# First pass: collect all forward nodes from all subgraphs
|
||||
for submod in fx_g.modules():
|
||||
if isinstance(submod, torch.fx.GraphModule):
|
||||
_collect_fwd_nodes_from_subgraph(submod, fwd_seq_nr_to_node)
|
||||
|
||||
if annotation_log.isEnabledFor(logging.DEBUG):
|
||||
for k, v in fwd_seq_nr_to_node.items():
|
||||
annotation_log.debug("forward:: key: %s, value: %s", k, v)
|
||||
|
||||
# Second pass: copy metadata to backward nodes in all subgraphs
|
||||
# using the global forward mapping
|
||||
for submod in fx_g.modules():
|
||||
if isinstance(submod, torch.fx.GraphModule):
|
||||
_copy_metadata_to_bw_nodes_in_subgraph(submod, fwd_seq_nr_to_node)
|
||||
|
||||
|
||||
def register_buffer_assignment_hook(
|
||||
mod: torch.nn.Module, assigned_buffers: dict[str, str]
|
||||
) -> Any:
|
||||
"""
|
||||
Register a hook that intercepts buffer assignments.
|
||||
This is used to detect when a buffer is assigned to, and then we can
|
||||
map that buffer to the corresponding proxy node in the graph.
|
||||
"""
|
||||
|
||||
def _map_assigned_buffer_to_proxy(
|
||||
_mod: torch.nn.Module, name: str, buffer: Any
|
||||
) -> Any:
|
||||
# We intercept buffer assignments on the root module through this hook.
|
||||
if _mod._buffers is mod._buffers:
|
||||
# either buffer is a functional tensor, which wraps a fake tensor
|
||||
if isinstance(buffer, FunctionalTensor):
|
||||
buffer = buffer.from_functional()
|
||||
# or buffer is a fake tensor
|
||||
if not isinstance(buffer, FakeTensor):
|
||||
raise AssertionError(f"expected FakeTensor, got {type(buffer)}")
|
||||
# The fake tensor in turn is associated with a proxy node.
|
||||
proxy_mode = torch.fx.experimental.proxy_tensor.get_proxy_mode()
|
||||
if proxy_mode is None:
|
||||
raise AssertionError("proxy_mode must not be None")
|
||||
proxy = torch.fx.experimental.proxy_tensor.get_proxy_slot(
|
||||
buffer, proxy_mode.tracer
|
||||
).proxy.node
|
||||
# We map the assigned buffer to this proxy node.
|
||||
assigned_buffers[name] = proxy.name
|
||||
return buffer
|
||||
|
||||
return torch.nn.modules.module.register_module_buffer_registration_hook(
|
||||
_map_assigned_buffer_to_proxy
|
||||
)
|
||||
|
||||
|
||||
def contain_metadata_mutation_ops(module: torch.fx.GraphModule) -> bool:
|
||||
"""
|
||||
Checks if the module contains any metadata mutation ops.
|
||||
"""
|
||||
for node in module.graph.nodes:
|
||||
if (
|
||||
node.op == "call_function"
|
||||
and hasattr(node.target, "tags")
|
||||
and torch.Tag.inplace_view in node.target.tags
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def get_cuda_generator_meta_val(device_idx: int) -> Any:
|
||||
"""
|
||||
Get a generator value to use as a meta val
|
||||
|
||||
newly cloned generator will not contain tensors. it is only Generators that are
|
||||
registered to a CUDAGraph that contain tensors. since this does not contain Tensor
|
||||
it is fine to use in the meta.
|
||||
"""
|
||||
return torch.cuda.default_generators[device_idx].clone_state()
|
||||
|
||||
|
||||
def top_saved_tensors_hooks() -> Any:
|
||||
return torch._C._autograd._top_saved_tensors_default_hooks(True)
|
||||
|
||||
|
||||
def saved_tensors_hooks_are_inlineable(hooks: Any) -> bool:
|
||||
if not hooks:
|
||||
return False
|
||||
pack, unpack = hooks
|
||||
return isinstance(pack, torch.fx.GraphModule) and isinstance(
|
||||
unpack, torch.fx.GraphModule
|
||||
)
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_T = TypeVar("_T")
|
||||
_S = TypeVar("_S")
|
||||
|
||||
|
||||
def without_output_descs(f: Callable[_P, tuple[_T, _S]]) -> Callable[_P, _T]:
|
||||
@wraps(f)
|
||||
@simple_wraps(f)
|
||||
def inner(*args: _P.args, **kwargs: _P.kwargs) -> _T:
|
||||
return f(*args, **kwargs)[0]
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
_P2 = ParamSpec("_P2")
|
||||
_R = TypeVar("_R")
|
||||
_R2 = TypeVar("_R2")
|
||||
|
||||
|
||||
def simple_wraps(
|
||||
f: Callable[_P, _R],
|
||||
) -> Callable[[Callable[_P2, _R2]], Callable[_P2, _R2]]:
|
||||
# NB: omit ('__module__', '__name__', '__qualname__') for ease of
|
||||
# debugging
|
||||
return wraps(f, assigned=("__doc__", "__annotations__", "__type_params__"))
|
||||
|
||||
|
||||
_Ts = TypeVarTuple("_Ts")
|
||||
|
||||
|
||||
def call_and_expect_output_descs(
|
||||
fn: Callable[[*_Ts], tuple[Any, Any]], args: tuple[Unpack[_Ts]]
|
||||
) -> tuple[Any, Any]:
|
||||
from .descriptors import AOTOutput
|
||||
|
||||
outs_pair = fn(*args)
|
||||
if not (isinstance(outs_pair, tuple) and len(outs_pair) == 2):
|
||||
raise AssertionError(
|
||||
f"expected tuple of length 2, got {type(outs_pair)} with value {outs_pair}"
|
||||
)
|
||||
outs, outs_descs = outs_pair
|
||||
# The Tensor tests protects against the test when there are no outputs
|
||||
out_vals, out_spec = pytree.tree_flatten(outs)
|
||||
out_desc_vals, out_desc_spec = pytree.tree_flatten(outs_descs)
|
||||
if out_spec != out_desc_spec:
|
||||
raise AssertionError(
|
||||
f"output spec mismatch: {fn_wrappers(fn)}, outs={outs}, outs_descs={outs_descs}, "
|
||||
f"out_spec={out_spec}, out_desc_spec={out_desc_spec}"
|
||||
)
|
||||
if any(isinstance(x, AOTOutput) for x in out_vals):
|
||||
raise AssertionError(
|
||||
f"unexpected AOTOutput in out_vals: {fn_wrappers(fn)}, outs={outs}, "
|
||||
f"outs_descs={outs_descs}, out_vals={out_vals}"
|
||||
)
|
||||
if not all(
|
||||
isinstance(d, AOTOutput)
|
||||
for (x, d) in zip(out_vals, out_desc_vals)
|
||||
if isinstance(x, (torch.Tensor, torch.SymInt)) or type(x) is int
|
||||
):
|
||||
raise AssertionError(
|
||||
f"expected all descriptors to be AOTOutput: {fn_wrappers(fn)}, outs={outs}, "
|
||||
f"outs_descs={outs_descs}, out_vals={out_vals}, out_desc_vals={out_desc_vals}"
|
||||
)
|
||||
return outs_pair
|
||||
|
||||
|
||||
def fn_wrappers(fn: Callable[..., Any]) -> list[Callable[..., Any]]:
|
||||
fns = [fn]
|
||||
f = fn
|
||||
while hasattr(f, "__wrapped__"):
|
||||
f = f.__wrapped__
|
||||
fns.append(f)
|
||||
return fns
|
||||
|
||||
|
||||
def _is_primal(node: torch.fx.Node) -> bool:
|
||||
return (
|
||||
node.op == "placeholder"
|
||||
and "tangents" not in str(node.target)
|
||||
and not _is_bwd_seed_offset(node)
|
||||
and not _is_fwd_seed_offset(node)
|
||||
)
|
||||
|
||||
|
||||
def _is_tangent(node: torch.fx.Node) -> bool:
|
||||
return node.op == "placeholder" and "tangents" in str(node.target)
|
||||
|
||||
|
||||
def _is_bwd_seed_offset(node: torch.fx.Node) -> bool:
|
||||
return node.op == "placeholder" and (
|
||||
"bwd_seed" in str(node.target) or "bwd_base_offset" in str(node.target)
|
||||
)
|
||||
|
||||
|
||||
def _is_fwd_seed_offset(node: torch.fx.Node) -> bool:
|
||||
return node.op == "placeholder" and (
|
||||
"fwd_seed" in str(node.target) or "fwd_base_offset" in str(node.target)
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,484 @@
|
||||
# NOTE: We allow Dynamo to see this file (via torch/_dynamo/trace_rules.py) so that it can
|
||||
# trace through functorch transforms.
|
||||
# Currently, we can't allow Dynamo to see `eager_transforms.py`/`vmap.py` as that break a lot of thing
|
||||
# and there isn't a mechanism to selectively expose only some functions (eg. grad) from a file
|
||||
# to Dynamo.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
from torch._functorch.utils import argnums_t, exposed_in
|
||||
from torch._functorch.vmap import (
|
||||
_check_out_dims_is_int_or_int_pytree,
|
||||
_check_randomness_arg,
|
||||
_chunked_vmap,
|
||||
_process_batched_inputs,
|
||||
Callable,
|
||||
in_dims_t,
|
||||
out_dims_t,
|
||||
vmap_impl,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
# vmap(func)(inputs) wraps all Tensor inputs to be batched in BatchedTensors,
|
||||
# sends those into func, and then unwraps the output BatchedTensors. Operations
|
||||
# on BatchedTensors perform the batched operations that the user is asking for.
|
||||
#
|
||||
# vmap's randomness behavior differs from JAX's, which would require a PRNG key
|
||||
# to be passed everywhere.
|
||||
|
||||
|
||||
@exposed_in("torch.func")
|
||||
def vmap(
|
||||
func: Callable[_P, _R],
|
||||
in_dims: in_dims_t = 0,
|
||||
out_dims: out_dims_t = 0,
|
||||
randomness: str = "error",
|
||||
*,
|
||||
chunk_size: int | None = None,
|
||||
) -> Callable[_P, _R]:
|
||||
"""
|
||||
vmap is the vectorizing map; ``vmap(func)`` returns a new function that
|
||||
maps ``func`` over some dimension of the inputs. Semantically, vmap
|
||||
pushes the map into PyTorch operations called by ``func``, effectively
|
||||
vectorizing those operations.
|
||||
|
||||
vmap is useful for handling batch dimensions: one can write a function
|
||||
``func`` that runs on examples and then lift it to a function that can
|
||||
take batches of examples with ``vmap(func)``. vmap can also be used to
|
||||
compute batched gradients when composed with autograd.
|
||||
|
||||
.. note::
|
||||
:func:`torch.vmap` is aliased to :func:`torch.func.vmap` for
|
||||
convenience. Use whichever one you'd like.
|
||||
|
||||
Args:
|
||||
func (function): A Python function that takes one or more arguments.
|
||||
Must return one or more Tensors.
|
||||
in_dims (int or nested structure): Specifies which dimension of the
|
||||
inputs should be mapped over. ``in_dims`` should have a
|
||||
structure like the inputs. If the ``in_dim`` for a particular
|
||||
input is None, then that indicates there is no map dimension.
|
||||
Default: 0.
|
||||
out_dims (int or Tuple[int]): Specifies where the mapped dimension
|
||||
should appear in the outputs. If ``out_dims`` is a Tuple, then
|
||||
it should have one element per output. Default: 0.
|
||||
randomness (str): Specifies whether the randomness in this
|
||||
vmap should be the same or different across batches. If 'different',
|
||||
the randomness for each batch will be different. If 'same', the
|
||||
randomness will be the same across batches. If 'error', any calls to
|
||||
random functions will error. Default: 'error'. WARNING: this flag
|
||||
only applies to random PyTorch operations and does not apply to
|
||||
Python's random module or numpy randomness.
|
||||
chunk_size (None or int): If None (default), apply a single vmap over inputs.
|
||||
If not None, then compute the vmap :attr:`chunk_size` samples at a time.
|
||||
Note that :attr:`chunk_size=1` is equivalent to computing the vmap with a for-loop.
|
||||
If you run into memory issues computing the vmap, please try a non-None chunk_size.
|
||||
|
||||
Returns:
|
||||
Returns a new "batched" function. It takes the same inputs as
|
||||
``func``, except each input has an extra dimension at the index
|
||||
specified by ``in_dims``. It takes returns the same outputs as
|
||||
``func``, except each output has an extra dimension at the index
|
||||
specified by ``out_dims``.
|
||||
|
||||
.. warning:
|
||||
:func:`vmap` works best with functional-style code. Please do not
|
||||
perform any side-effects in ``func``, with the exception of
|
||||
in-place PyTorch operations. Examples of side-effects include mutating
|
||||
Python data structures and assigning values to variables not captured
|
||||
in ``func``.
|
||||
|
||||
One example of using :func:`vmap` is to compute batched dot products. PyTorch
|
||||
doesn't provide a batched ``torch.dot`` API; instead of unsuccessfully
|
||||
rummaging through docs, use :func:`vmap` to construct a new function.
|
||||
|
||||
>>> torch.dot # [D], [D] -> []
|
||||
>>> batched_dot = torch.func.vmap(torch.dot) # [N, D], [N, D] -> [N]
|
||||
>>> x, y = torch.randn(2, 5), torch.randn(2, 5)
|
||||
>>> batched_dot(x, y)
|
||||
|
||||
:func:`vmap` can be helpful in hiding batch dimensions, leading to a simpler
|
||||
model authoring experience.
|
||||
|
||||
>>> batch_size, feature_size = 3, 5
|
||||
>>> weights = torch.randn(feature_size, requires_grad=True)
|
||||
>>>
|
||||
>>> def model(feature_vec):
|
||||
>>> # Very simple linear model with activation
|
||||
>>> return feature_vec.dot(weights).relu()
|
||||
>>>
|
||||
>>> examples = torch.randn(batch_size, feature_size)
|
||||
>>> result = torch.vmap(model)(examples)
|
||||
|
||||
:func:`vmap` can also help vectorize computations that were previously difficult
|
||||
or impossible to batch. One example is higher-order gradient computation.
|
||||
The PyTorch autograd engine computes vjps (vector-Jacobian products).
|
||||
Computing a full Jacobian matrix for some function f: R^N -> R^N usually
|
||||
requires N calls to ``autograd.grad``, one per Jacobian row. Using :func:`vmap`,
|
||||
we can vectorize the whole computation, computing the Jacobian in a single
|
||||
call to ``autograd.grad``.
|
||||
|
||||
>>> # Setup
|
||||
>>> N = 5
|
||||
>>> f = lambda x: x**2
|
||||
>>> x = torch.randn(N, requires_grad=True)
|
||||
>>> y = f(x)
|
||||
>>> I_N = torch.eye(N)
|
||||
>>>
|
||||
>>> # Sequential approach
|
||||
>>> jacobian_rows = [torch.autograd.grad(y, x, v, retain_graph=True)[0]
|
||||
>>> for v in I_N.unbind()]
|
||||
>>> jacobian = torch.stack(jacobian_rows)
|
||||
>>>
|
||||
>>> # vectorized gradient computation
|
||||
>>> def get_vjp(v):
|
||||
>>> return torch.autograd.grad(y, x, v)
|
||||
>>> jacobian = torch.vmap(get_vjp)(I_N)
|
||||
|
||||
:func:`vmap` can also be nested, producing an output with multiple batched dimensions
|
||||
|
||||
>>> torch.dot # [D], [D] -> []
|
||||
>>> batched_dot = torch.vmap(
|
||||
... torch.vmap(torch.dot)
|
||||
... ) # [N1, N0, D], [N1, N0, D] -> [N1, N0]
|
||||
>>> x, y = torch.randn(2, 3, 5), torch.randn(2, 3, 5)
|
||||
>>> batched_dot(x, y) # tensor of size [2, 3]
|
||||
|
||||
If the inputs are not batched along the first dimension, ``in_dims`` specifies
|
||||
the dimension that each inputs are batched along as
|
||||
|
||||
>>> torch.dot # [N], [N] -> []
|
||||
>>> batched_dot = torch.vmap(torch.dot, in_dims=1) # [N, D], [N, D] -> [D]
|
||||
>>> x, y = torch.randn(2, 5), torch.randn(2, 5)
|
||||
>>> batched_dot(
|
||||
... x, y
|
||||
... ) # output is [5] instead of [2] if batched along the 0th dimension
|
||||
|
||||
If there are multiple inputs each of which is batched along different dimensions,
|
||||
``in_dims`` must be a tuple with the batch dimension for each input as
|
||||
|
||||
>>> torch.dot # [D], [D] -> []
|
||||
>>> batched_dot = torch.vmap(torch.dot, in_dims=(0, None)) # [N, D], [D] -> [N]
|
||||
>>> x, y = torch.randn(2, 5), torch.randn(5)
|
||||
>>> batched_dot(
|
||||
... x, y
|
||||
... ) # second arg doesn't have a batch dim because in_dim[1] was None
|
||||
|
||||
If the input is a Python struct, ``in_dims`` must be a tuple containing a struct
|
||||
matching the shape of the input:
|
||||
|
||||
>>> f = lambda dict: torch.dot(dict["x"], dict["y"])
|
||||
>>> x, y = torch.randn(2, 5), torch.randn(5)
|
||||
>>> input = {"x": x, "y": y}
|
||||
>>> batched_dot = torch.vmap(f, in_dims=({"x": 0, "y": None},))
|
||||
>>> batched_dot(input)
|
||||
|
||||
By default, the output is batched along the first dimension. However, it can be batched
|
||||
along any dimension by using ``out_dims``
|
||||
|
||||
>>> f = lambda x: x**2
|
||||
>>> x = torch.randn(2, 5)
|
||||
>>> batched_pow = torch.vmap(f, out_dims=1)
|
||||
>>> batched_pow(x) # [5, 2]
|
||||
|
||||
For any function that uses kwargs, the returned function will not batch the kwargs but will
|
||||
accept kwargs
|
||||
|
||||
>>> x = torch.randn([2, 5])
|
||||
>>> def fn(x, scale=4.):
|
||||
>>> return x * scale
|
||||
>>>
|
||||
>>> batched_pow = torch.vmap(fn)
|
||||
>>> assert torch.allclose(batched_pow(x), x * 4)
|
||||
>>> batched_pow(x, scale=x) # scale is not batched, output has shape [2, 2, 5]
|
||||
|
||||
.. note::
|
||||
vmap does not provide general autobatching or handle variable-length
|
||||
sequences out of the box.
|
||||
"""
|
||||
from torch.compiler import is_compiling
|
||||
|
||||
_check_randomness_arg(randomness)
|
||||
if not (chunk_size is None or chunk_size > 0):
|
||||
raise ValueError(
|
||||
f"vmap: chunk_size should be None or greater than 0. (got {chunk_size})"
|
||||
)
|
||||
|
||||
def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
return vmap_impl(
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
func,
|
||||
in_dims,
|
||||
out_dims,
|
||||
randomness,
|
||||
chunk_size,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not is_compiling():
|
||||
wrapped = functools.wraps(func)(wrapped)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def chunk_vmap(
|
||||
func: Callable[_P, _R],
|
||||
in_dims: in_dims_t = 0,
|
||||
out_dims: out_dims_t = 0,
|
||||
randomness: str = "error",
|
||||
chunks: int = 2,
|
||||
) -> Callable[_P, _R]:
|
||||
"""
|
||||
chunk_vmap is the vectorizing map (vmap) using chunks of input data. It is a mix of vmap (which vectorizes
|
||||
everything) and map (which executes things sequentially). ``chunk_vmap`` vectorizes the input with number of
|
||||
chunks at a time. For more details about vectorizing map, see :func:`vmap`.
|
||||
|
||||
.. note::
|
||||
Please use :func:`vmap` with ``chunk_size`` argument instead of this API.
|
||||
|
||||
Args:
|
||||
func (function): A Python function that takes one or more arguments.
|
||||
Must return one or more Tensors.
|
||||
in_dims (int or nested structure): Specifies which dimension of the
|
||||
inputs should be mapped over. ``in_dims`` should have a
|
||||
structure like the inputs. If the ``in_dim`` for a particular
|
||||
input is None, then that indicates there is no map dimension.
|
||||
Default: 0.
|
||||
out_dims (int or Tuple[int]): Specifies where the mapped dimension
|
||||
should appear in the outputs. If ``out_dims`` is a Tuple, then
|
||||
it should have one element per output. Default: 0.
|
||||
randomness (str): Specifies whether the randomness in this
|
||||
vmap should be the same or different across batches. If 'different',
|
||||
the randomness for each batch will be different. If 'same', the
|
||||
randomness will be the same across batches. If 'error', any calls to
|
||||
random functions will error. Default: 'error'. WARNING: this flag
|
||||
only applies to random PyTorch operations and does not apply to
|
||||
Python's random module or numpy randomness.
|
||||
chunks (int): Number of chunks to use to split the input data. Default is 2.
|
||||
If equals to 1 then :func:`vmap` is called.
|
||||
|
||||
Returns:
|
||||
Returns a new "batched" function. It takes the same inputs as
|
||||
``func``, except each input has an extra dimension at the index
|
||||
specified by ``in_dims``. It takes returns the same outputs as
|
||||
``func``, except each output has an extra dimension at the index
|
||||
specified by ``out_dims``.
|
||||
"""
|
||||
_check_randomness_arg(randomness)
|
||||
|
||||
if chunks == 1:
|
||||
return vmap(func, in_dims=in_dims, out_dims=out_dims, randomness=randomness)
|
||||
|
||||
def _get_chunk_flat_args(
|
||||
flat_args_: Iterable[Any],
|
||||
flat_in_dims_: Iterable[int | None],
|
||||
chunks_: int,
|
||||
) -> Iterable[Any]:
|
||||
flat_args_chunks = tuple(
|
||||
t.chunk(chunks_, dim=in_dim)
|
||||
if in_dim is not None
|
||||
else [
|
||||
t,
|
||||
]
|
||||
* chunks_
|
||||
for t, in_dim in zip(flat_args_, flat_in_dims_)
|
||||
)
|
||||
# transpose chunk dim and flatten structure
|
||||
# chunks_flat_args is a list of flatten args
|
||||
chunks_flat_args = zip(*flat_args_chunks)
|
||||
return chunks_flat_args
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapped_with_chunks(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
_check_out_dims_is_int_or_int_pytree(out_dims, func)
|
||||
_, flat_in_dims, flat_args, args_spec = _process_batched_inputs(
|
||||
in_dims, args, func
|
||||
)
|
||||
# Chunk flat arguments
|
||||
chunks_flat_args = _get_chunk_flat_args(flat_args, flat_in_dims, chunks)
|
||||
|
||||
# Apply vmap on chunks
|
||||
return _chunked_vmap(
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
func,
|
||||
flat_in_dims,
|
||||
chunks_flat_args,
|
||||
args_spec,
|
||||
out_dims,
|
||||
randomness,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return wrapped_with_chunks
|
||||
|
||||
|
||||
# TODO: Improve the return type of this function
|
||||
@exposed_in("torch.func")
|
||||
def grad(
|
||||
func: Callable[_P, Any], argnums: argnums_t = 0, has_aux: bool = False
|
||||
) -> Callable[_P, Any]:
|
||||
"""``grad`` operator helps computing gradients of ``func`` with respect to the
|
||||
input(s) specified by ``argnums``. This operator can be nested to
|
||||
compute higher-order gradients.
|
||||
|
||||
Args:
|
||||
func (Callable): A Python function that takes one or more arguments.
|
||||
Must return a single-element Tensor. If specified ``has_aux`` equals ``True``,
|
||||
function can return a tuple of single-element Tensor and other auxiliary objects:
|
||||
``(output, aux)``.
|
||||
argnums (int or Tuple[int]): Specifies arguments to compute gradients with respect to.
|
||||
``argnums`` can be single integer or tuple of integers. Default: 0.
|
||||
has_aux (bool): Flag indicating that ``func`` returns a tensor and other
|
||||
auxiliary objects: ``(output, aux)``. Default: False.
|
||||
|
||||
Returns:
|
||||
Function to compute gradients with respect to its inputs. By default, the output of
|
||||
the function is the gradient tensor(s) with respect to the first argument.
|
||||
If specified ``has_aux`` equals ``True``, tuple of gradients and output auxiliary objects
|
||||
is returned. If ``argnums`` is a tuple of integers, a tuple of output gradients with
|
||||
respect to each ``argnums`` value is returned.
|
||||
|
||||
Example of using ``grad``:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> from torch.func import grad
|
||||
>>> x = torch.randn([])
|
||||
>>> cos_x = grad(lambda x: torch.sin(x))(x)
|
||||
>>> assert torch.allclose(cos_x, x.cos())
|
||||
>>>
|
||||
>>> # Second-order gradients
|
||||
>>> neg_sin_x = grad(grad(lambda x: torch.sin(x)))(x)
|
||||
>>> assert torch.allclose(neg_sin_x, -x.sin())
|
||||
|
||||
When composed with ``vmap``, ``grad`` can be used to compute per-sample-gradients:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> from torch.func import grad, vmap
|
||||
>>> batch_size, feature_size = 3, 5
|
||||
>>>
|
||||
>>> def model(weights, feature_vec):
|
||||
>>> # Very simple linear model with activation
|
||||
>>> assert feature_vec.dim() == 1
|
||||
>>> return feature_vec.dot(weights).relu()
|
||||
>>>
|
||||
>>> def compute_loss(weights, example, target):
|
||||
>>> y = model(weights, example)
|
||||
>>> return ((y - target) ** 2).mean() # MSELoss
|
||||
>>>
|
||||
>>> weights = torch.randn(feature_size, requires_grad=True)
|
||||
>>> examples = torch.randn(batch_size, feature_size)
|
||||
>>> targets = torch.randn(batch_size)
|
||||
>>> inputs = (weights, examples, targets)
|
||||
>>> grad_weight_per_example = vmap(grad(compute_loss), in_dims=(None, 0, 0))(
|
||||
... *inputs
|
||||
... )
|
||||
|
||||
Example of using ``grad`` with ``has_aux`` and ``argnums``:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> from torch.func import grad
|
||||
>>> def my_loss_func(y, y_pred):
|
||||
>>> loss_per_sample = (0.5 * y_pred - y) ** 2
|
||||
>>> loss = loss_per_sample.mean()
|
||||
>>> return loss, (y_pred, loss_per_sample)
|
||||
>>>
|
||||
>>> fn = grad(my_loss_func, argnums=(0, 1), has_aux=True)
|
||||
>>> y_true = torch.rand(4)
|
||||
>>> y_preds = torch.rand(4, requires_grad=True)
|
||||
>>> out = fn(y_true, y_preds)
|
||||
>>> # > output is ((grads w.r.t y_true, grads w.r.t y_preds), (y_pred, loss_per_sample))
|
||||
|
||||
.. note::
|
||||
Using PyTorch ``torch.no_grad`` together with ``grad``.
|
||||
|
||||
Case 1: Using ``torch.no_grad`` inside a function:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> def f(x):
|
||||
>>> with torch.no_grad():
|
||||
>>> c = x ** 2
|
||||
>>> return x - c
|
||||
|
||||
In this case, ``grad(f)(x)`` will respect the inner ``torch.no_grad``.
|
||||
|
||||
Case 2: Using ``grad`` inside ``torch.no_grad`` context manager:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> with torch.no_grad():
|
||||
>>> grad(f)(x)
|
||||
|
||||
In this case, ``grad`` will respect the inner ``torch.no_grad``, but not the
|
||||
outer one. This is because ``grad`` is a "function transform": its result
|
||||
should not depend on the result of a context manager outside of ``f``.
|
||||
|
||||
"""
|
||||
# To avoid cyclical dependency.
|
||||
import torch._functorch.eager_transforms as eager_transforms
|
||||
from torch.compiler import is_compiling
|
||||
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> tuple[Any, torch.Tensor]:
|
||||
return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs)
|
||||
|
||||
if not is_compiling():
|
||||
wrapper = functools.wraps(func)(wrapper)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# TODO: Improve the return type of this function
|
||||
@exposed_in("torch.func")
|
||||
def grad_and_value(
|
||||
func: Callable[_P, Any], argnums: argnums_t = 0, has_aux: bool = False
|
||||
) -> Callable[_P, tuple[Any, Any]]:
|
||||
"""
|
||||
Returns a function to compute a tuple of the gradient and primal, or
|
||||
forward, computation.
|
||||
|
||||
Args:
|
||||
func (Callable): A Python function that takes one or more arguments.
|
||||
Must return a single-element Tensor. If specified ``has_aux``
|
||||
equals ``True``, function can return a tuple of single-element
|
||||
Tensor and other auxiliary objects: ``(output, aux)``.
|
||||
argnums (int or Tuple[int]): Specifies arguments to compute gradients
|
||||
with respect to. ``argnums`` can be single integer or tuple of
|
||||
integers. Default: 0.
|
||||
has_aux (bool): Flag indicating that ``func`` returns a tensor and
|
||||
other auxiliary objects: ``(output, aux)``. Default: False.
|
||||
|
||||
Returns:
|
||||
Function to compute a tuple of gradients with respect to its inputs
|
||||
and the forward computation. By default, the output of the function is
|
||||
a tuple of the gradient tensor(s) with respect to the first argument
|
||||
and the primal computation. If specified ``has_aux`` equals
|
||||
``True``, tuple of gradients and tuple of the forward computation with
|
||||
output auxiliary objects is returned. If ``argnums`` is a tuple of
|
||||
integers, a tuple of a tuple of the output gradients with respect to
|
||||
each ``argnums`` value and the forward computation is returned.
|
||||
|
||||
See :func:`grad` for examples
|
||||
"""
|
||||
from torch._functorch import eager_transforms
|
||||
from torch.compiler import is_compiling
|
||||
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> tuple[Any, torch.Tensor]:
|
||||
return eager_transforms.grad_and_value_impl(
|
||||
func, argnums, has_aux, args, kwargs
|
||||
)
|
||||
|
||||
if not is_compiling():
|
||||
wrapper = functools.wraps(func)(wrapper)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,909 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, NamedTuple, TYPE_CHECKING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._C._functorch import (
|
||||
_unwrap_for_grad,
|
||||
_wrap_for_grad,
|
||||
current_level,
|
||||
TransformType,
|
||||
)
|
||||
from torch._functorch.apis import vmap
|
||||
from torch._functorch.utils import enable_single_level_autograd_function
|
||||
from torch._functorch.vmap import (
|
||||
_add_batch_dim,
|
||||
_broadcast_to_and_flatten,
|
||||
restore_vmap,
|
||||
unwrap_batched,
|
||||
wrap_batched,
|
||||
)
|
||||
from torch._ops import HigherOrderOperator
|
||||
from torch.autograd.forward_ad import _set_fwd_grad_enabled
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
|
||||
from torch._functorch.pyfunctorch import FuncTorchInterpreter, VmapInterpreter
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
# autograd.Function technically runs before the regular PyTorch dispatcher.
|
||||
# This is how features like autocast and torch_dispatch (e.g. PythonTLSSnapshot)
|
||||
# work with it. One day we might decide to change this, but until then,
|
||||
# we need to give the illusion that autograd.Function runs before those things.
|
||||
#
|
||||
# We do this by using creating a custom HigherOrderOperator that only functorch
|
||||
# dispatches specially.
|
||||
class CustomFunctionHigherOrderOperator(HigherOrderOperator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("custom_function_call")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs,
|
||||
) -> Any:
|
||||
# When custom_function_call is done dispatching through functorch,
|
||||
# it should just invoke the autograd.Function. This is consistent
|
||||
# with the autograd.Function behavior of being invoked before the
|
||||
# PyTorch dispatcher.
|
||||
#
|
||||
# This will lead us into trouble later down the line, but this is
|
||||
# pre-existing. There is an invariant that a function traced by
|
||||
# make_fx should have the same behavior when provided the same
|
||||
# Tensor. However, make_fx sees autograd.Function as a composite
|
||||
# (because autograd.Function happens before the Python dispatch key)
|
||||
# and only traces the forward pass.
|
||||
if torch._C._are_functorch_transforms_active():
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return super().__call__(autograd_function, *args, **kwargs)
|
||||
return autograd_function.apply(*args, **kwargs)
|
||||
|
||||
|
||||
# "custom_function_call"
|
||||
# This is the mechanism for an autograd.Function that works with functorch transforms.
|
||||
# It wraps an autograd.Function; interactions with functorch transforms are defined
|
||||
# via PyDispatcher and HigherOrderOperator rather than through the traditional PyTorch
|
||||
# dispatcher.
|
||||
custom_function_call = CustomFunctionHigherOrderOperator()
|
||||
|
||||
|
||||
# The grad rule for custom_function_call is to construct a new _SingleLevelFunction
|
||||
# (autograd.Function that only works with a single layer (level) of functorch) that:
|
||||
# - unwraps the inputs
|
||||
# - redispatches to custom_function_call
|
||||
# - wraps the outputs
|
||||
# and whose backward pass calls the original autograd.Function's backward.
|
||||
#
|
||||
# Why do we need to redispatch to custom_function_call?
|
||||
# -----------------------------------------------------
|
||||
# This is consistent with how ATen operators work with functorch's grad transform:
|
||||
# they always redispatch to the original operator.
|
||||
# Consider torch.sin, and let's say we do grad0(grad1(torch.sin))(x)
|
||||
#
|
||||
# grad1 will:
|
||||
# - set up the autograd graph
|
||||
# - unwrap the inputs
|
||||
# - redispatch to at::sin (*)
|
||||
# - rewrap the outputs on the return
|
||||
#
|
||||
# On the redispatch in (*), grad0 will:
|
||||
# - set up the autograd graph
|
||||
# - unwrap the inputs
|
||||
# - redispatch to at::sin
|
||||
# - rewrap the outputs on the return
|
||||
#
|
||||
# To "set up the autograd graph", we generate a _SingleLevelFunction
|
||||
# and apply it.
|
||||
@custom_function_call.py_impl(TransformType.Grad)
|
||||
@custom_function_call.py_impl(TransformType.Jvp)
|
||||
def custom_function_call_grad(
|
||||
interpreter: FuncTorchInterpreter,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
*operands: Any,
|
||||
) -> Any:
|
||||
Generated = generate_single_level_function(interpreter, autograd_function)
|
||||
with enable_single_level_autograd_function():
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
flat_out = Generated.apply(*operands)
|
||||
return flat_out
|
||||
|
||||
|
||||
def generate_single_level_function(
|
||||
interpreter: FuncTorchInterpreter,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
) -> type[torch.autograd.function._SingleLevelFunction]:
|
||||
level = interpreter.level()
|
||||
|
||||
def forward(*operands: Any) -> Any:
|
||||
unwrapped_operands = pytree.tree_map_only(
|
||||
torch.Tensor, lambda x: _unwrap_for_grad(x, level), operands
|
||||
)
|
||||
# Both enable_grad() and _set_fwd_grad_enabled() are necessary no matter
|
||||
# the transform. _SingleLevelFunction will turn off both fwd and bwd
|
||||
# gradient computation and we need to turn it back on here.
|
||||
with torch.enable_grad(), _set_fwd_grad_enabled(True), interpreter.lower():
|
||||
unwrapped_output = custom_function_call(
|
||||
autograd_function, *unwrapped_operands
|
||||
)
|
||||
|
||||
# See NOTE [mark_dirty object identity check]
|
||||
def wrap_fn(output: torch.Tensor) -> torch.Tensor:
|
||||
return _wrap_for_grad(output, level)
|
||||
|
||||
return wrap_outputs_maintaining_identity(
|
||||
unwrapped_output, unwrapped_operands, operands, wrap_fn
|
||||
)
|
||||
|
||||
def setup_context(ctx: Any, inputs: Any, output: Any) -> Any:
|
||||
return autograd_function.setup_context(ctx, inputs, output)
|
||||
|
||||
# backward is only used if the transform is TransformType.Grad
|
||||
def backward(ctx: Any, *grads: Any) -> Any:
|
||||
result = autograd_function.backward(ctx, *grads)
|
||||
return result
|
||||
|
||||
# jvp is only used if the transform is TransformType.Jvp
|
||||
def jvp(ctx: Any, *tangents: Any) -> Any:
|
||||
result = autograd_function.jvp(ctx, *tangents)
|
||||
return result
|
||||
|
||||
# This is the sequence of magic words to dynamically generate a Subclass with
|
||||
# a given name. A Tensor's .grad_fn field has a class name that is the original
|
||||
# autograd.Function's name + Backward, so we do this to generate some
|
||||
# meaningful name.
|
||||
name = f"{autograd_function.__name__}Generated"
|
||||
Generated = type(
|
||||
name,
|
||||
(torch.autograd.function._SingleLevelFunction,),
|
||||
{
|
||||
"forward": staticmethod(forward),
|
||||
"backward": staticmethod(backward),
|
||||
"jvp": staticmethod(jvp),
|
||||
"setup_context": staticmethod(setup_context),
|
||||
},
|
||||
)
|
||||
return Generated
|
||||
|
||||
|
||||
# wrap_outputs_maintaining_identity handles outputs from the vmap,
|
||||
# backward (vjp), and jvp staticmethod. The way it distinguishes
|
||||
# between the vmap case and the {backward, jvp} case is if the out_dims
|
||||
# are specified or not.
|
||||
#
|
||||
# NB: we cannot use out_dims=None as the deciding factor. This because
|
||||
# out_dims=None can still happen in the vmap staticmethod! What the
|
||||
# user is saying in that case is that their output does not have a
|
||||
# dimension that is being vmapped over, which is valid.
|
||||
NO_OUT_DIMS = "not specified"
|
||||
|
||||
|
||||
# NOTE [mark_dirty object identity check]
|
||||
# autograd.Function's ctx.mark_dirty expect a returned input
|
||||
# to have the same object identity as the input.
|
||||
# Mode-only functorch will greatly simplify this logic.
|
||||
def wrap_outputs_maintaining_identity(
|
||||
outputs: Any,
|
||||
unwrapped_inputs: Any,
|
||||
orig_inputs: Any,
|
||||
wrap_fn: Callable[..., Any],
|
||||
out_dims: Any = NO_OUT_DIMS,
|
||||
) -> Any:
|
||||
flat_unwrapped_inputs = pytree.arg_tree_leaves(*unwrapped_inputs)
|
||||
flat_orig_inputs = pytree.arg_tree_leaves(*orig_inputs)
|
||||
|
||||
unwrapped_input_to_orig_input = {
|
||||
id(unwrapped): orig
|
||||
for unwrapped, orig in zip(flat_unwrapped_inputs, flat_orig_inputs)
|
||||
}
|
||||
|
||||
flat_outputs, spec = pytree.tree_flatten(outputs)
|
||||
result = []
|
||||
|
||||
out_dims_specified = out_dims != NO_OUT_DIMS
|
||||
|
||||
flat_out_dims = None
|
||||
if out_dims_specified:
|
||||
flat_out_dims = _broadcast_to_and_flatten(out_dims, spec)
|
||||
# _broadcast_to_and_flatten returns None if it is unable to broadcast.
|
||||
# TODO: update following link from master to stable once that's out
|
||||
if flat_out_dims is None:
|
||||
raise RuntimeError(
|
||||
f"The autograd.Function's vmap staticmethod returned an "
|
||||
f"incompatible (output, out_dims) tuple. "
|
||||
f"Expected out_dims={out_dims} "
|
||||
f"to be compatible with the structure of `output`. "
|
||||
f"out_dims has structure {pytree.tree_flatten(out_dims)[1]} "
|
||||
f"but output has structure {spec}. "
|
||||
f"For more details, please see "
|
||||
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
||||
)
|
||||
|
||||
for i, output in enumerate(flat_outputs):
|
||||
if not isinstance(output, torch.Tensor):
|
||||
result.append(output)
|
||||
continue
|
||||
if id(output) in unwrapped_input_to_orig_input:
|
||||
result.append(unwrapped_input_to_orig_input[id(output)])
|
||||
continue
|
||||
if out_dims_specified:
|
||||
if flat_out_dims is None:
|
||||
raise AssertionError(
|
||||
"flat_out_dims must not be None when out_dims is specified"
|
||||
)
|
||||
result.append(wrap_fn(output, flat_out_dims[i]))
|
||||
else:
|
||||
result.append(wrap_fn(output))
|
||||
|
||||
return pytree.tree_unflatten(result, spec)
|
||||
|
||||
|
||||
# NOTE: [functorch vjp and autograd interaction]
|
||||
# There's an edge case with the functorch vjp and autograd interaction
|
||||
# that will eventually be fixed by mode-only functorch.
|
||||
# The TL;DR is that there's no way to unwrap a dead GradTensorWrapper,
|
||||
# so we (the framework) need to do it manually. Regular PyTorch operators
|
||||
# automatically do so this is consistent.
|
||||
#
|
||||
# class MyExp(torch.autograd.Function):
|
||||
# @staticmethod
|
||||
# def forward(x):
|
||||
# return x.exp()
|
||||
#
|
||||
# @staticmethod
|
||||
# def setup_context(ctx, inputs, output):
|
||||
# y = output
|
||||
# ctx.save_for_backward(y)
|
||||
#
|
||||
# @staticmethod
|
||||
# def backward(gy):
|
||||
# y, = ctx.saved_tensors()
|
||||
# return MyMul.apply(gy, y)
|
||||
#
|
||||
# x = torch.randn([], requires_grad=True)
|
||||
# gy = torch.randn([], requires_grad=True)
|
||||
# _, vjp_fn = vjp(MySin.apply, x)
|
||||
# result = vjp_fn(gy)
|
||||
#
|
||||
# MyMul is an autograd.Function that is not shown here.
|
||||
# It saves a `y` for backward (since gy requires grad).
|
||||
#
|
||||
# in vjp_fn(gy), we get:
|
||||
# > MyMul.apply(gy, GradTensorWrapper(y, level=dead))
|
||||
# Because the y that is saved for backward by MyExp is a GradTensorWrapper
|
||||
# but is now dead since we are outside the vjp context.
|
||||
#
|
||||
# PyTorch dispatcher operations, upon seeing a dead GradTensorWrapper,
|
||||
# will automatically unwrap the GradTensorWrapper when applied.
|
||||
# But since autograd.Function technically sits above the regular PyTorch
|
||||
# dispatcher, it doesn't get this treatment. So we manually do
|
||||
# the unwrapping to be consistent with regular PyTorch dispatcher operations.
|
||||
|
||||
|
||||
class VmapInfo(NamedTuple):
|
||||
batch_size: int
|
||||
randomness: str
|
||||
|
||||
|
||||
def has_overridden_vmap_rule(
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
) -> bool:
|
||||
return autograd_function.vmap is not torch.autograd.Function.vmap
|
||||
|
||||
|
||||
def validate_vmap_returns_tuple_of_two_elements(result: Any) -> None:
|
||||
base_error_msg = (
|
||||
"Expected the vmap staticmethod to have two returns, an output "
|
||||
"and out_dims with pytree structure compatible with the output. "
|
||||
)
|
||||
if not isinstance(result, tuple):
|
||||
raise RuntimeError(base_error_msg + f"Got a {type(result)} instead")
|
||||
if not len(result) == 2:
|
||||
raise RuntimeError(base_error_msg + f"Got {len(result)} returns instead")
|
||||
|
||||
|
||||
@custom_function_call.py_impl(TransformType.Vmap)
|
||||
def custom_function_call_vmap(
|
||||
interpreter: VmapInterpreter,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
*operands: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if any(
|
||||
isinstance(val, torch.Tensor)
|
||||
for val in torch.utils._pytree.tree_flatten(kwargs)[0]
|
||||
):
|
||||
raise NotImplementedError(
|
||||
f"Run vmap on autograd.Function with kwarg-only Tensor args. "
|
||||
f"Please do not pass kwarg-only Tensors to autograd.Function. "
|
||||
f"Got: {kwargs}"
|
||||
)
|
||||
|
||||
if autograd_function.generate_vmap_rule:
|
||||
if has_overridden_vmap_rule(autograd_function):
|
||||
# TODO: Update link to stable once that's out
|
||||
# https://github.com/pytorch/pytorch/issues/92029
|
||||
raise RuntimeError(
|
||||
f"You tried to vmap over {autograd_function.__name__}, but "
|
||||
f"it has both generate_vmap_rule=True and an overridden vmap "
|
||||
f"staticmethod. Please set generate_vmap_rule=False or delete "
|
||||
f"the overridden vmap staticmethod to avoid ambiguity. "
|
||||
f"For more details, please see "
|
||||
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
||||
)
|
||||
return custom_function_call_vmap_generate_rule(
|
||||
interpreter, autograd_function, *operands
|
||||
)
|
||||
|
||||
if not has_overridden_vmap_rule(autograd_function):
|
||||
# TODO: Update link to stable once that's out
|
||||
# https://github.com/pytorch/pytorch/issues/92029
|
||||
raise RuntimeError(
|
||||
f"You tried to vmap over {autograd_function.__name__}, but "
|
||||
f"it does not have vmap support. Please override and implement the "
|
||||
f"vmap staticmethod or set generate_vmap_rule=True. "
|
||||
f"For more details, please see "
|
||||
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
||||
)
|
||||
|
||||
return custom_function_call_vmap_helper(
|
||||
interpreter, autograd_function.vmap, autograd_function, *operands, **kwargs
|
||||
)
|
||||
|
||||
|
||||
def custom_function_call_vmap_helper(
|
||||
interpreter: VmapInterpreter,
|
||||
vmap_function: Callable[..., Any],
|
||||
op: Any,
|
||||
*operands: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
current_level = interpreter.level()
|
||||
info = VmapInfo(
|
||||
batch_size=interpreter.batch_size(),
|
||||
randomness=interpreter.randomness(),
|
||||
)
|
||||
# We're either in the autograd.Function case (vmap staticmethod)
|
||||
# or the torch.library.register_vmap case.
|
||||
autograd_function_case = isinstance(op, torch.autograd.function.FunctionMeta)
|
||||
|
||||
def lower_to_next() -> Any:
|
||||
if autograd_function_case:
|
||||
return interpreter.lower()
|
||||
else:
|
||||
return torch._C._ExcludeDispatchKeyGuard(
|
||||
torch._C.DispatchKeySet(torch._C.DispatchKey.FuncTorchBatched)
|
||||
)
|
||||
|
||||
unwrapped_operands, in_dims = unwrap_batched(operands, current_level)
|
||||
# If none of the tensors are batched at the current level, then we skip the
|
||||
# current level. This saves the user from needing to handle this case in
|
||||
# their vmap staticmethod (and is consistent with our C++ batching rule API)
|
||||
if pytree.tree_all(lambda dim: dim is None, in_dims):
|
||||
with lower_to_next():
|
||||
if autograd_function_case:
|
||||
return custom_function_call(op, *operands)
|
||||
else:
|
||||
return op(*operands, **kwargs)
|
||||
|
||||
with lower_to_next():
|
||||
result = vmap_function(info, in_dims, *unwrapped_operands, **kwargs)
|
||||
validate_vmap_returns_tuple_of_two_elements(result)
|
||||
unwrapped_output, out_dims = result
|
||||
|
||||
# See NOTE [mark_dirty object identity check]
|
||||
def wrap_fn(output: torch.Tensor, out_dim: int | None) -> torch.Tensor:
|
||||
return (
|
||||
output
|
||||
if out_dim is None
|
||||
else _add_batch_dim(output, out_dim, current_level)
|
||||
)
|
||||
|
||||
return wrap_outputs_maintaining_identity(
|
||||
unwrapped_output, unwrapped_operands, operands, wrap_fn, out_dims=out_dims
|
||||
)
|
||||
|
||||
|
||||
def unpack_outputs(outputs: tuple[Any, ...]) -> tuple[Any, Any]:
|
||||
out_dims = outputs[-1]
|
||||
if isinstance(out_dims, tuple):
|
||||
outputs = outputs[:-1]
|
||||
else:
|
||||
outputs = outputs[0]
|
||||
return outputs, out_dims
|
||||
|
||||
|
||||
def custom_function_call_vmap_generate_rule(
|
||||
interpreter: VmapInterpreter,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
*operands: Any,
|
||||
) -> Any:
|
||||
unwrapped_operands, in_dims = unwrap_batched(operands, interpreter.level())
|
||||
vmapped_function = vmapify_autograd_function(
|
||||
autograd_function,
|
||||
in_dims,
|
||||
interpreter.batch_size(),
|
||||
interpreter.randomness(),
|
||||
)
|
||||
with interpreter.lower():
|
||||
outputs = custom_function_call(vmapped_function, *unwrapped_operands)
|
||||
|
||||
if not isinstance(outputs, tuple):
|
||||
raise AssertionError(f"expected outputs to be a tuple, got {type(outputs)}")
|
||||
outputs, out_dims = unpack_outputs(outputs)
|
||||
return wrap_batched(outputs, out_dims, interpreter.level())
|
||||
|
||||
|
||||
@custom_function_call.py_impl(TransformType.Functionalize)
|
||||
def custom_function_call_functionalize(
|
||||
interpreter: FuncTorchInterpreter,
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
generate_vmap_rule: bool,
|
||||
*operands: Any,
|
||||
) -> Any:
|
||||
raise RuntimeError("NYI: Functionalize rule for custom_function_call")
|
||||
|
||||
|
||||
def vmapify_autograd_function(
|
||||
autograd_function: type[torch.autograd.Function],
|
||||
in_dims: Any,
|
||||
batch_size: int,
|
||||
randomness: str,
|
||||
) -> type[torch.autograd.Function]:
|
||||
def forward(*operands: Any) -> Any:
|
||||
outputs, out_dims = restore_vmap(
|
||||
autograd_function.forward, in_dims, batch_size, randomness
|
||||
)(*operands)
|
||||
if isinstance(outputs, torch.Tensor):
|
||||
return outputs, out_dims
|
||||
else:
|
||||
return *outputs, out_dims
|
||||
|
||||
def setup_context(ctx: Any, inputs: Any, outputs: Any) -> None:
|
||||
outputs, out_dims = unpack_outputs(outputs)
|
||||
key = id(Generated)
|
||||
|
||||
def inner(inputs: Any, outputs: Any) -> None:
|
||||
# wrapped_ctx.save_for_backward will:
|
||||
# - unwrap batchedtensors into (tensor, bdim)
|
||||
# - save_for_backward(*unwrapped_tensors)
|
||||
# - assign the bdims to wrapped_ctx._pt_saved_tensors_bdims
|
||||
wrapped_ctx = CtxCustomSave(ctx, current_level())
|
||||
autograd_function.setup_context(wrapped_ctx, inputs, outputs)
|
||||
|
||||
# input_shapes are used for reductify later to reduce expanded gradients
|
||||
# to the correct shape.
|
||||
# See NOTE: [Why can't we rely on autograd to reduce expanded gradients?]
|
||||
# for more details
|
||||
input_shapes = tuple(
|
||||
inp.shape if isinstance(inp, torch.Tensor) else None for inp in inputs
|
||||
)
|
||||
if not hasattr(ctx, "_pt_input_shapes"):
|
||||
# pyrefly: ignore [implicit-any]
|
||||
ctx._pt_input_shapes = {}
|
||||
ctx._pt_input_shapes.update({key: input_shapes})
|
||||
|
||||
if not hasattr(ctx, "_pt_saved_tensors_bdims_stack"):
|
||||
# pyrefly: ignore [implicit-any]
|
||||
ctx._pt_saved_tensors_bdims_stack = {}
|
||||
ctx._pt_saved_tensors_bdims_stack.update(
|
||||
{key: (wrapped_ctx._pt_saved_tensors_bdims)}
|
||||
)
|
||||
|
||||
# See NOTE: [Why do we need to run setup_context under a vmap?]
|
||||
restore_vmap(
|
||||
inner,
|
||||
(in_dims, out_dims),
|
||||
batch_size,
|
||||
randomness,
|
||||
)(inputs, outputs)
|
||||
|
||||
if not hasattr(ctx, "_pt_out_dims"):
|
||||
# pyrefly: ignore [implicit-any]
|
||||
ctx._pt_out_dims = {}
|
||||
ctx._pt_out_dims.update({key: out_dims})
|
||||
|
||||
def jvp(ctx: Any, *tangents: Any) -> Any:
|
||||
key = id(Generated)
|
||||
|
||||
def jvp_no_context(saved_tensors: Any, tangents: Any) -> Any:
|
||||
wrapped_ctx = CtxWithSavedTensors(ctx, saved_tensors)
|
||||
return autograd_function.jvp(wrapped_ctx, *tangents)
|
||||
|
||||
tangent_in_dims = get_tangents_in_dims(in_dims, tangents)
|
||||
out_tangents, out_tangents_dims = restore_vmap(
|
||||
jvp_no_context,
|
||||
(ctx._pt_saved_tensors_bdims_stack[key], tangent_in_dims),
|
||||
batch_size,
|
||||
randomness,
|
||||
)(ctx.saved_tensors, tangents)
|
||||
|
||||
result = reductify(
|
||||
out_tangents, out_tangents_dims, ctx._pt_out_dims[key], batch_size
|
||||
)
|
||||
if isinstance(result, torch.Tensor):
|
||||
return result, None
|
||||
else:
|
||||
return *result, None
|
||||
|
||||
def backward(ctx: Any, *grad_outputs: Any) -> Any:
|
||||
key = id(Generated)
|
||||
grad_outputs_ = grad_outputs[:-1]
|
||||
grad_outputs_in_dims = ctx._pt_out_dims[key]
|
||||
|
||||
if not isinstance(grad_outputs_in_dims, tuple):
|
||||
grad_outputs_in_dims = (grad_outputs_in_dims,)
|
||||
|
||||
grad_outputs_in_dims = tuple(
|
||||
in_dim if grad_output is not None else None
|
||||
for grad_output, in_dim in zip(grad_outputs_, grad_outputs_in_dims)
|
||||
)
|
||||
|
||||
def backward_no_context(inputs: Any) -> Any:
|
||||
saved_tensors, grad_outputs = inputs
|
||||
wrapped_ctx = CtxWithSavedTensors(ctx, saved_tensors)
|
||||
return autograd_function.backward(wrapped_ctx, *grad_outputs)
|
||||
|
||||
grad_ins, grad_ins_dims = restore_vmap(
|
||||
backward_no_context,
|
||||
((ctx._pt_saved_tensors_bdims_stack[key], grad_outputs_in_dims),),
|
||||
batch_size,
|
||||
randomness,
|
||||
)((ctx.saved_tensors, grad_outputs_))
|
||||
result = reductify(
|
||||
grad_ins, grad_ins_dims, in_dims, batch_size, ctx._pt_input_shapes[key]
|
||||
)
|
||||
return result
|
||||
|
||||
name = f"Vmapped{autograd_function.__name__}"
|
||||
Generated = type(
|
||||
name,
|
||||
(torch.autograd.Function,),
|
||||
{
|
||||
"forward": staticmethod(forward),
|
||||
"backward": staticmethod(backward),
|
||||
"jvp": staticmethod(jvp),
|
||||
"setup_context": staticmethod(setup_context),
|
||||
"generate_vmap_rule": True,
|
||||
},
|
||||
)
|
||||
|
||||
return Generated
|
||||
|
||||
|
||||
# tangents might be None, so we need to replace
|
||||
# the corresponding in_dims with None.
|
||||
def get_tangents_in_dims(input_dims: Any, tangents: tuple[Any, ...]) -> Any:
|
||||
flat_in_dims, spec = pytree.tree_flatten(input_dims)
|
||||
flat_tangents = pytree.arg_tree_leaves(*tangents)
|
||||
result = [
|
||||
None if tangent is None else in_dim
|
||||
for in_dim, tangent in zip(flat_in_dims, flat_tangents)
|
||||
]
|
||||
return pytree.tree_unflatten(result, spec)
|
||||
|
||||
|
||||
# NOTE: [Why do we need to run setup_context under a vmap?]
|
||||
# Consider the following autograd.Function
|
||||
#
|
||||
# class Sum(torch.autograd.Function):
|
||||
# @staticmethod
|
||||
# def forward(x):
|
||||
# return x.sum()
|
||||
# @staticmethod
|
||||
# def setup_context(ctx, inputs, outputs):
|
||||
# ctx.x_shape = inputs[0]
|
||||
# @staticmethod
|
||||
# def backward(ctx, gy):
|
||||
# return gy.expand(ctx.x_shape)
|
||||
#
|
||||
# x = torch.randn(B, 4)
|
||||
# in_dims = 0
|
||||
# vmap(Sum.apply, in_dims)(x)
|
||||
#
|
||||
# Let's assume for a moment that we didn't vmap setup_context in VmappedSum:
|
||||
#
|
||||
# class VmappedSum(torch.autograd.Function):
|
||||
# @staticmethod
|
||||
# def forward(x):
|
||||
# return vmap(Sum.forward, in_dims)(x)
|
||||
#
|
||||
# @staticmethod
|
||||
# def setup_context(ctx, inputs, outputs):
|
||||
# Sum.setup_context(ctx, inputs, outputs)
|
||||
#
|
||||
# @staticmethod
|
||||
# def backward(ctx, gy):
|
||||
# def backward_no_context(gy):
|
||||
# return gy.expand(ctx.x_shape)
|
||||
#
|
||||
# dims = (0,)
|
||||
# gx = vmap(backward_no_context, dims)(gy)
|
||||
# return gx
|
||||
#
|
||||
# We end up saving [B, 4] as x_shape. In the backward, gy has shape [B],
|
||||
# and we're doing:
|
||||
#
|
||||
# def backward_no_context(gy):
|
||||
# return gy.expand([B, 4])
|
||||
#
|
||||
# gx = vmap(backward_no_context, dims)(gy: "Tensor[B]")
|
||||
#
|
||||
# This gives us the wrong result (gx has shape [B, B, 4], but it should
|
||||
# have shape [4]). Performing vmap over setup_context means the shape
|
||||
# saved has shape [4] and leads to a correct result shape for gx.
|
||||
|
||||
|
||||
# Wraps a ctx object. Forwards all attr accesses to the underlying object
|
||||
# except for the attrs in _pt_attrs
|
||||
class WrappedCtx:
|
||||
_pt_reserved_attrs: tuple[str, ...] = ("_pt_reserved_attrs", "_pt_inner_ctx")
|
||||
|
||||
def __init__(self, ctx: Any) -> None:
|
||||
if not isinstance(ctx, WrappedCtx):
|
||||
reserved_attrs = type(self)._pt_reserved_attrs
|
||||
for name in reserved_attrs:
|
||||
if not hasattr(ctx, name):
|
||||
continue
|
||||
raise RuntimeError(
|
||||
f"PyTorch reserves the {reserved_attrs} field on ctx. "
|
||||
"Please name your fields on ctx something else to avoid name "
|
||||
"collision."
|
||||
)
|
||||
self._pt_inner_ctx = ctx
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._pt_inner_ctx, name)
|
||||
|
||||
def __setattr__(self, name: str, value: Any) -> None:
|
||||
if name in type(self)._pt_reserved_attrs:
|
||||
self.__dict__[name] = value
|
||||
return
|
||||
return setattr(self._pt_inner_ctx, name, value)
|
||||
|
||||
|
||||
# Wraps ctx to create a new ctx object that overrides saved_tensors.
|
||||
class CtxWithSavedTensors(WrappedCtx):
|
||||
_pt_reserved_attrs = ("_pt_new_saved_tensors", *WrappedCtx._pt_reserved_attrs)
|
||||
|
||||
def __init__(self, ctx: Any, new_saved_tensors: Sequence[torch.Tensor]) -> None:
|
||||
super().__init__(ctx)
|
||||
self._pt_new_saved_tensors = new_saved_tensors
|
||||
|
||||
@property
|
||||
def saved_tensors(self) -> Sequence[torch.Tensor]:
|
||||
return self._pt_new_saved_tensors
|
||||
|
||||
|
||||
class CtxCustomSave(WrappedCtx):
|
||||
_pt_reserved_attrs = (
|
||||
"_pt_saved_tensors_bdims",
|
||||
"_pt_current_level",
|
||||
*WrappedCtx._pt_reserved_attrs,
|
||||
)
|
||||
|
||||
def __init__(self, ctx: Any, current_level: int) -> None:
|
||||
super().__init__(ctx)
|
||||
self._pt_saved_tensors_bdims: tuple[Any, ...] = ()
|
||||
self._pt_current_level = current_level
|
||||
|
||||
def save_for_backward(self, *tensors: torch.Tensor) -> None:
|
||||
unwrapped_tensors, bdims = unwrap_batched(tensors, self._pt_current_level)
|
||||
self._pt_inner_ctx.save_for_backward(*unwrapped_tensors)
|
||||
self._pt_saved_tensors_bdims = bdims
|
||||
|
||||
def save_for_forward(self, *tensors: torch.Tensor) -> None:
|
||||
unwrapped_tensors, bdims = unwrap_batched(tensors, self._pt_current_level)
|
||||
self._pt_inner_ctx.save_for_forward(*unwrapped_tensors)
|
||||
self._pt_saved_tensors_bdims = bdims
|
||||
|
||||
|
||||
def reductify(
|
||||
grad_input: torch.Tensor | tuple[torch.Tensor, ...],
|
||||
grad_input_bdim: int | tuple[int, ...],
|
||||
input_bdim: int | tuple[int, ...],
|
||||
batch_size: int,
|
||||
target_shape_without_bdim_to_reduce_to: Any = None,
|
||||
) -> tuple[Any, ...]:
|
||||
if not isinstance(grad_input, tuple):
|
||||
grad_input = (grad_input,)
|
||||
if not isinstance(grad_input_bdim, tuple):
|
||||
grad_input_bdim = (grad_input_bdim,)
|
||||
if not isinstance(input_bdim, tuple):
|
||||
input_bdim = (input_bdim,)
|
||||
|
||||
if target_shape_without_bdim_to_reduce_to is None:
|
||||
target_shape_without_bdim_to_reduce_to = len(grad_input) * (None,)
|
||||
result = tuple(
|
||||
reductify_leaf(gi, gi_bdim, i_bdim, batch_size, maybe_ishape)
|
||||
for gi, gi_bdim, i_bdim, maybe_ishape in zip(
|
||||
grad_input,
|
||||
grad_input_bdim,
|
||||
input_bdim,
|
||||
target_shape_without_bdim_to_reduce_to,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def reductify_leaf(
|
||||
grad_input: torch.Tensor | None,
|
||||
grad_input_bdim: int | None,
|
||||
input_bdim: int | None,
|
||||
batch_size: int,
|
||||
target_shape_without_bdim_to_reduce_to: Any = None,
|
||||
) -> torch.Tensor | None:
|
||||
if grad_input is None:
|
||||
return None
|
||||
|
||||
if grad_input_bdim is None and input_bdim is None:
|
||||
return grad_input
|
||||
|
||||
if grad_input_bdim is not None and input_bdim is None:
|
||||
return grad_input.sum(grad_input_bdim)
|
||||
|
||||
# NOTE: [Why can't we rely on autograd to reduce expanded gradients?]
|
||||
# For reverse-mode AD,
|
||||
# given a grad_input and input, it is valid for the user to return a
|
||||
# grad_input that has a broadcasted shape when compared to the input.
|
||||
# In this situation, autograd automatically reduces the grad_input to
|
||||
# the shape of the input.
|
||||
#
|
||||
# However, when input_bdim is not None, we have problems.
|
||||
#
|
||||
# [example 1]
|
||||
# grad_input: Tensor[3, 4], input: Tensor[B, 4]
|
||||
# We can expand grad_input to Tensor[B, 3, 4], but that isn't broadcastable
|
||||
# from [B, 4].
|
||||
#
|
||||
# [example 2]
|
||||
# grad_input: Tensor[3, B, 4], input: Tensor[B, 4]
|
||||
# We can swizzle grad_input to Tensor[B, 3, 4], but that isn't broadcastable
|
||||
# from [B, 4].
|
||||
#
|
||||
# This means that we need to also reduce the grad_input to the shape of the
|
||||
# input. This behavior is controlled by the `target_shape_without_bdim_to_reduce_to` flag;
|
||||
# if not-None then we do the reducing manually, otherwise, we do not do a reduction.
|
||||
if input_bdim is None:
|
||||
raise AssertionError("input_bdim must not be None")
|
||||
|
||||
if grad_input_bdim is None:
|
||||
grad_input = grad_input.unsqueeze(input_bdim)
|
||||
new_shape = list(grad_input.shape)
|
||||
new_shape[input_bdim] = batch_size
|
||||
grad_input = grad_input.expand(new_shape)
|
||||
grad_input_bdim = input_bdim
|
||||
|
||||
if target_shape_without_bdim_to_reduce_to is not None:
|
||||
return vmap(
|
||||
torch.Tensor.sum_to_size,
|
||||
in_dims=(grad_input_bdim, None),
|
||||
out_dims=input_bdim,
|
||||
)(grad_input, target_shape_without_bdim_to_reduce_to)
|
||||
|
||||
if input_bdim != grad_input_bdim:
|
||||
grad_input = grad_input.movedim(grad_input_bdim, input_bdim)
|
||||
return grad_input
|
||||
|
||||
|
||||
def autograd_function_forward_rewritten(
|
||||
original_forward: Callable[_P, _R],
|
||||
original_setup_context: Callable[..., Any],
|
||||
) -> Callable[..., _R]:
|
||||
def new_forward(ctx: Any, *args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
output = original_forward(*args, **kwargs)
|
||||
original_setup_context(ctx, args, output)
|
||||
return output
|
||||
|
||||
return new_forward
|
||||
|
||||
|
||||
class AutogradFunctionApply(HigherOrderOperator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("autograd_function_apply")
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
fwd: torch.fx.GraphModule,
|
||||
bwd: torch.fx.GraphModule,
|
||||
*fwd_args: Any,
|
||||
**fwd_kwargs: Any,
|
||||
) -> Any:
|
||||
saved_values: Iterable[Any] | None = None
|
||||
non_differentiable_idx = fwd_kwargs["non_differentiable_idx"]
|
||||
saved_for_backward_idx = fwd_kwargs["saved_for_backward_idx"]
|
||||
|
||||
class ApplyTemplate(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(*args: Any, **kwargs: Any) -> Any:
|
||||
nonlocal saved_values
|
||||
|
||||
# The Interpreter here is required to propagate metadata
|
||||
# from the dynamo graph body to the local_map graph body.
|
||||
# This is required for fx_traceback.annotate for work.
|
||||
output, saved_values = torch.fx.Interpreter(fwd).run(*args)
|
||||
|
||||
# See Note [Activations with no version counter checks in eager]
|
||||
# Mark tensors that came from ctx.save_for_backward with metadata.
|
||||
# This allows AOT autograd to distinguish between tensors saved via
|
||||
# save_for_backward vs those stashed directly on ctx (e.g., ctx.x = x).
|
||||
from torch.fx.experimental.proxy_tensor import _get_proxies
|
||||
|
||||
for idx, t in enumerate(saved_values):
|
||||
if idx not in saved_for_backward_idx:
|
||||
for proxy in _get_proxies(t):
|
||||
proxy.node.meta["saved_tensor_with_no_vc_check"] = True
|
||||
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
|
||||
# If users call ctx.mark_non_differentiable() in the original fwd function.
|
||||
if len(non_differentiable_idx) > 0:
|
||||
non_differentiable_output = []
|
||||
for i, x in enumerate(output):
|
||||
if i in non_differentiable_idx:
|
||||
non_differentiable_output.append(x)
|
||||
ctx.mark_non_differentiable(*non_differentiable_output)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: Any, *grad: Any) -> Any:
|
||||
# The Interpreter here is required to propagate metadata
|
||||
# from the dynamo graph body to the local_map graph body.
|
||||
# This is required for fx_traceback.annotate for work.
|
||||
|
||||
if saved_values is None:
|
||||
raise AssertionError("saved_values must not be None")
|
||||
return torch.fx.Interpreter(bwd).run(*grad, *saved_values)
|
||||
|
||||
return ApplyTemplate.apply(*fwd_args)
|
||||
|
||||
|
||||
autograd_function_apply = AutogradFunctionApply()
|
||||
|
||||
|
||||
class DynamoAutogradFunctionTraceHelper:
|
||||
@staticmethod
|
||||
def fwd_trace_helper(orig_fwd: Callable[_P, Any]) -> Callable[_P, Any]:
|
||||
# autograd.Function forward does more than just running the forward method. Most
|
||||
# of this logic is in C++. Here, we rewrite that functionality in python and let
|
||||
# Dynamo trace it.
|
||||
def inner(*args: _P.args, **kwargs: _P.kwargs) -> Any:
|
||||
with torch.no_grad():
|
||||
outs = orig_fwd(*args, **kwargs)
|
||||
|
||||
# Handle the case where if the input is passed on directly to the output, we call view_as
|
||||
# Refer to https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/custom_function.cpp#L254
|
||||
tensor_args = {arg for arg in args if isinstance(arg, torch.Tensor)}
|
||||
if isinstance(outs, torch.Tensor):
|
||||
if outs in tensor_args:
|
||||
return outs.view_as(outs)
|
||||
else:
|
||||
return outs
|
||||
|
||||
new_outs = []
|
||||
for out in outs:
|
||||
if isinstance(out, torch.Tensor):
|
||||
if out in tensor_args:
|
||||
new_outs.append(out.view_as(out))
|
||||
else:
|
||||
new_outs.append(out)
|
||||
else:
|
||||
new_outs.append(out)
|
||||
return tuple(new_outs)
|
||||
|
||||
# TODO - there is missing functionality here, where
|
||||
# autograd.Function overwrites the requires_grad_ of the output
|
||||
# tensors depending on the `mark_non_differentiable`. Currently,
|
||||
# this is handled hackily in Dynamo, where we just overwrite the
|
||||
# variable trackers requires_grad. Refer to the function -
|
||||
# overwrite_tensor_vt_requires_grad
|
||||
|
||||
return inner
|
||||
@@ -0,0 +1,27 @@
|
||||
import torch.nn as nn
|
||||
from torch._functorch.utils import exposed_in
|
||||
|
||||
|
||||
def batch_norm_without_running_stats(module: nn.Module) -> None:
|
||||
if (
|
||||
isinstance(module, nn.modules.batchnorm._BatchNorm)
|
||||
and module.track_running_stats
|
||||
):
|
||||
module.running_mean = None
|
||||
module.running_var = None
|
||||
module.num_batches_tracked = None
|
||||
module.track_running_stats = False
|
||||
|
||||
|
||||
@exposed_in("torch.func")
|
||||
def replace_all_batch_norm_modules_(root: nn.Module) -> nn.Module:
|
||||
"""
|
||||
In place updates :attr:`root` by setting the ``running_mean`` and ``running_var`` to be None and
|
||||
setting track_running_stats to be False for any nn.BatchNorm module in :attr:`root`
|
||||
"""
|
||||
# base case
|
||||
batch_norm_without_running_stats(root)
|
||||
|
||||
for obj in root.modules():
|
||||
batch_norm_without_running_stats(obj)
|
||||
return root
|
||||
@@ -0,0 +1,245 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import operator
|
||||
import os
|
||||
import time
|
||||
from contextlib import AbstractContextManager
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
import torch
|
||||
from torch.profiler import profile, ProfilerActivity
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
|
||||
_R = TypeVar("_R")
|
||||
|
||||
|
||||
def synchronize() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def dump_chrome_trace(
|
||||
f: Callable[[tuple[Any, ...]], _R],
|
||||
input_: tuple[Any, ...],
|
||||
trace_filename: str,
|
||||
optimize_ctx: AbstractContextManager[Any],
|
||||
activities: Sequence[ProfilerActivity],
|
||||
num_runs: int = 1,
|
||||
devices: list[str] | None = None,
|
||||
kwargs_for_f: dict[str, Any] | None = None,
|
||||
kwargs_for_profiler: dict[str, Any] | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
Output the chrome trace of running f(input_, **kwargs_for_f) with [optimize_ctx]
|
||||
[num_runs] times to [trace_filename].
|
||||
|
||||
[activities] are the activities that the profiler will record, e.g. ProfilerActivity.CUDA.
|
||||
Return total runtime without the profiler
|
||||
|
||||
Outputs to trace_filename
|
||||
"""
|
||||
|
||||
if devices is None:
|
||||
devices = ["cuda"]
|
||||
|
||||
global synchronize
|
||||
if devices != ["cpu"] and torch.cuda.is_available():
|
||||
synchronize = torch.cuda.synchronize
|
||||
|
||||
if kwargs_for_f is None:
|
||||
kwargs_for_f = {}
|
||||
if kwargs_for_profiler is None:
|
||||
kwargs_for_profiler = {}
|
||||
|
||||
with optimize_ctx:
|
||||
torch.manual_seed(1337)
|
||||
for _ in range(5): # warmup runs
|
||||
f(input_, **kwargs_for_f)
|
||||
synchronize()
|
||||
torch.manual_seed(1337)
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(num_runs):
|
||||
f(input_, **kwargs_for_f)
|
||||
synchronize()
|
||||
t1 = time.perf_counter()
|
||||
timing = t1 - t0
|
||||
|
||||
with profile(activities=activities, **kwargs_for_profiler) as prof:
|
||||
with optimize_ctx:
|
||||
synchronize()
|
||||
torch.manual_seed(1337)
|
||||
for _ in range(num_runs):
|
||||
f(input_, **kwargs_for_f)
|
||||
synchronize()
|
||||
prof.export_chrome_trace(trace_filename)
|
||||
|
||||
return timing
|
||||
|
||||
|
||||
def get_chrome_trace_events(filename: str) -> list[dict[str, Any]]:
|
||||
with open(filename) as f:
|
||||
data = json.load(f)
|
||||
events = data["traceEvents"]
|
||||
return events
|
||||
|
||||
|
||||
def is_gpu_compute_event(event: dict[str, Any]) -> bool:
|
||||
global gpu_pids
|
||||
return (
|
||||
"pid" in event
|
||||
and event["pid"] in gpu_pids
|
||||
and "ph" in event
|
||||
and event["ph"] == "X"
|
||||
)
|
||||
|
||||
|
||||
def get_sorted_gpu_events(events: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
sorted_gpu_events: list[dict[str, Any]] = []
|
||||
for event in events:
|
||||
if not is_gpu_compute_event(event):
|
||||
continue
|
||||
sorted_gpu_events.append(event)
|
||||
return sorted(sorted_gpu_events, key=operator.itemgetter("ts"))
|
||||
|
||||
|
||||
def get_duration(sorted_gpu_events: list[dict[str, Any]]) -> int:
|
||||
if len(sorted_gpu_events) == 0:
|
||||
return 0
|
||||
event = sorted_gpu_events[0]
|
||||
current_end_time = event["ts"] + event["dur"]
|
||||
total_duration = event["dur"]
|
||||
for event in sorted_gpu_events[1:]:
|
||||
start_time = max(event["ts"], current_end_time)
|
||||
end_time = event["ts"] + event["dur"]
|
||||
total_duration = total_duration + max(end_time - start_time, 0)
|
||||
current_end_time = max(current_end_time, end_time)
|
||||
return total_duration
|
||||
|
||||
|
||||
def get_sorted_gpu_mm_conv_events(events: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
def is_mm_conv_event(event: dict[str, Any]) -> bool:
|
||||
return "name" in event and (
|
||||
"gemm" in event["name"]
|
||||
or "conv" in event["name"]
|
||||
or "cutlass" in event["name"]
|
||||
or "wgrad" in event["name"]
|
||||
)
|
||||
|
||||
gpu_events = get_sorted_gpu_events(events)
|
||||
sorted_events: list[dict[str, Any]] = []
|
||||
for event in gpu_events:
|
||||
if not is_mm_conv_event(event):
|
||||
continue
|
||||
sorted_events.append(event)
|
||||
return sorted_events
|
||||
|
||||
|
||||
gpu_pids: list[Any] = []
|
||||
|
||||
|
||||
def compute_utilization(filename: str, total_length: float) -> tuple[float, float]:
|
||||
"""
|
||||
Process the chrome traces outputs by the pytorch profiler to compute GPU Utilization
|
||||
and percent of times spent on matmul and convolution
|
||||
|
||||
Args:
|
||||
filename(str): Name of chrome traces file produced by pytorch profiler
|
||||
|
||||
total_length(float): total length of the process without profiler in second
|
||||
|
||||
Return:
|
||||
tuple: (GPU Utilization, percent of time spent on matmul and convolution)
|
||||
"""
|
||||
events = get_chrome_trace_events(filename)
|
||||
|
||||
# get pids of GPU events
|
||||
global gpu_pids
|
||||
gpu_pids = []
|
||||
for event in events:
|
||||
if "name" not in event:
|
||||
continue
|
||||
if event["name"] == "process_labels" and "GPU" in event["args"]["labels"]:
|
||||
gpu_pids.append(event["pid"])
|
||||
|
||||
total_length = total_length * 1e6
|
||||
sorted_gpu_events = get_sorted_gpu_events(events)
|
||||
utilization = get_duration(sorted_gpu_events) / total_length
|
||||
|
||||
sorted_gpu_mm_conv_events = get_sorted_gpu_mm_conv_events(events)
|
||||
mm_conv_utilization = get_duration(sorted_gpu_mm_conv_events) / total_length
|
||||
|
||||
return utilization, mm_conv_utilization
|
||||
|
||||
|
||||
def benchmark_utilization(
|
||||
f: Callable[[tuple[Any, ...]], _R],
|
||||
input_: tuple[Any, ...],
|
||||
trace_folder: str,
|
||||
optimize_ctx: AbstractContextManager[Any] | None = None,
|
||||
trace_file_name: str = "tmp_chrome_trace",
|
||||
num_runs: int = 1,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Benchmark the GPU Utilization and percent of time spent on matmul and convolution operations of
|
||||
running f(input_, **kwargs_for_f) with [optimize_ctx] [num_runs] times.
|
||||
It will produce a chrome trace file in trace_folder/trace_file_name.json
|
||||
|
||||
Example:
|
||||
|
||||
```
|
||||
def f(a):
|
||||
return a.sum()
|
||||
|
||||
|
||||
a = torch.rand(2**20, device="cuda")
|
||||
utilization, mm_conv_utilization = benchmark_utilization(
|
||||
f, a, "tmp", trace_file_name="tmp_chrome_trace"
|
||||
)
|
||||
```
|
||||
|
||||
Args:
|
||||
f: function to benchmark
|
||||
|
||||
input_: input to :attr:`f`
|
||||
|
||||
trace_folder: name of the folder to store the chrome trace
|
||||
|
||||
optimize_ctx: the context in which f will run
|
||||
|
||||
trace_file_name: name of the dumped chrome trace file, default to "tmp_chrome_trace"
|
||||
|
||||
num_runs: number of times to run f, excluding the warm-up runs, default to 1.
|
||||
|
||||
Return:
|
||||
tuple: (GPU Utilization, percent of time spent on matmul and convolution)
|
||||
|
||||
"""
|
||||
isExist = os.path.exists(trace_folder)
|
||||
if not isExist:
|
||||
os.makedirs(trace_folder)
|
||||
print("create folder " + trace_folder)
|
||||
|
||||
if optimize_ctx is None:
|
||||
optimize_ctx = contextlib.nullcontext()
|
||||
|
||||
chrome_trace_file_name = os.path.join(trace_folder, trace_file_name + ".json")
|
||||
total_length = dump_chrome_trace(
|
||||
f,
|
||||
input_,
|
||||
chrome_trace_file_name,
|
||||
optimize_ctx,
|
||||
[ProfilerActivity.CUDA],
|
||||
num_runs=num_runs,
|
||||
devices=["cuda"],
|
||||
)
|
||||
utilization, mm_conv_utilization = compute_utilization(
|
||||
chrome_trace_file_name, total_length
|
||||
)
|
||||
|
||||
return utilization, mm_conv_utilization
|
||||
@@ -0,0 +1,229 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils import _pytree as pytree
|
||||
from torch.utils._pytree import tree_flatten
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch._ops import OpOverloadPacket
|
||||
from torch.utils._pytree import TreeSpec
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
def get_aten_target(node: fx.Node) -> OpOverloadPacket | Callable[..., Any] | str:
|
||||
if hasattr(node.target, "overloadpacket"):
|
||||
return node.target.overloadpacket
|
||||
return node.target
|
||||
|
||||
|
||||
rand_ops = [
|
||||
aten.dropout,
|
||||
aten._fused_dropout,
|
||||
aten._standard_gamma,
|
||||
aten.bernoulli,
|
||||
aten.multinomial,
|
||||
aten.native_dropout,
|
||||
aten.normal,
|
||||
aten.poisson,
|
||||
aten.binomial,
|
||||
aten.rrelu,
|
||||
aten.rand_like,
|
||||
aten.rand,
|
||||
aten.randint,
|
||||
aten.randn,
|
||||
aten.randperm,
|
||||
]
|
||||
|
||||
|
||||
# return a new copy of torch.fx.graph.Graph with CSE applied to the input graph
|
||||
def fx_graph_cse(fx_g: torch.fx.graph.Graph) -> fx.Graph:
|
||||
new_graph = fx.Graph()
|
||||
env: dict[
|
||||
fx.Node, fx.Node
|
||||
] = {} # map from node in the old graph to node in the new graph
|
||||
hash_env: dict[
|
||||
tuple[str, int], fx.Node
|
||||
] = {} # map from hash to a node in the new graph
|
||||
token_map: dict[tuple[str, int], dict[str, Any]] = {} # map from hash to token
|
||||
|
||||
from torch._inductor.pattern_matcher import (
|
||||
compute_mutation_region_ids,
|
||||
same_mutation_regions,
|
||||
)
|
||||
|
||||
compute_mutation_region_ids(fx_g) # type: ignore[arg-type]
|
||||
|
||||
# Make a set of separate storages returned from the output, which will be preserved
|
||||
# when pruning. This prevents us from deduplicating returned tensors which have
|
||||
# experienced identical operations, but are separate data structures in eager mode.
|
||||
output_node: fx.Node = list(fx_g.nodes)[-1]
|
||||
if output_node.op != "output":
|
||||
raise AssertionError(
|
||||
f"expected output_node.op to be 'output', got '{output_node.op}'"
|
||||
)
|
||||
|
||||
def checkable_node(node: fx.Node) -> bool:
|
||||
"""We can evaluate only nodes that represent tensors with defined storage."""
|
||||
if "val" not in node.meta or not isinstance(node.meta["val"], torch.Tensor):
|
||||
return False
|
||||
|
||||
try:
|
||||
node.meta["val"].untyped_storage()
|
||||
except NotImplementedError:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
output_storages = {
|
||||
StorageWeakRef(n.meta["val"].untyped_storage())
|
||||
for n in output_node.all_input_nodes
|
||||
if checkable_node(n)
|
||||
}
|
||||
nodes_that_alias_outputs = {
|
||||
n
|
||||
for n in fx_g.nodes
|
||||
if checkable_node(n)
|
||||
and StorageWeakRef(n.meta["val"].untyped_storage()) in output_storages
|
||||
}
|
||||
|
||||
for n in fx_g.nodes:
|
||||
# The placeholder, output, and get_attr nodes are copied to the new graph without change
|
||||
# do not CSE away random operations
|
||||
if (
|
||||
n.op == "placeholder"
|
||||
or n.op == "output"
|
||||
or n.op == "get_attr"
|
||||
or get_aten_target(n) in rand_ops
|
||||
# aten.empty is non-deterministic, so don't CSE it.
|
||||
# Also, aten.empty is almost always fusible into its consumer,
|
||||
# so it's not worth CSEing.
|
||||
or get_aten_target(n) is aten.empty
|
||||
or n in nodes_that_alias_outputs
|
||||
# This CSE pass currently doesn't handle re-propagation of unbacked
|
||||
# meta where it'll sometimes eliminate a _local_scalar_dense but not
|
||||
# replace the meta of downstream users. eg. one bug we've seen is:
|
||||
#
|
||||
# _local_scalar_dense_11: "Sym(u14)" = torch.ops.aten._local_scalar_dense.default(select_10);
|
||||
# sym_sum_2: "Sym(u19 + u20 + u21)" = torch.sym_sum((_local_scalar_dense_11, _local_scalar_dense_12, _local_scalar_dense_13)) # noqa: B950
|
||||
#
|
||||
# Notice how _local_scalar_dense_11 is u14 but sym_sum_2's meta is incorrectly the old
|
||||
# pre-cse value of u19.
|
||||
or (
|
||||
"val" in n.meta
|
||||
and isinstance(n.meta["val"], sympy.Symbol)
|
||||
and free_unbacked_symbols(n.meta["val"])
|
||||
)
|
||||
):
|
||||
new_node = new_graph.node_copy(n, lambda x: env[x])
|
||||
env[n] = new_node
|
||||
else: # n.op == 'call_function', should never see n.op == 'call_module' or 'call_method'
|
||||
# substitute args and kwargs members to their mapping in env if exists
|
||||
# specs can be used to reconstruct nested list/dictionaries
|
||||
def substitute(
|
||||
arg_list: list[Any] | tuple[Any, ...],
|
||||
) -> tuple[tuple[Any, ...], TreeSpec]:
|
||||
arg_list, spec = tree_flatten(arg_list)
|
||||
for i in range(len(arg_list)):
|
||||
v = arg_list[i]
|
||||
if isinstance(v, torch.fx.node.Node) and v in env:
|
||||
arg_list[i] = env[v]
|
||||
if isinstance(v, (torch.SymBool, torch.SymInt, torch.SymFloat)):
|
||||
arg_list[i] = v.node
|
||||
return tuple(arg_list), spec
|
||||
|
||||
args, args_spec = substitute(n.args)
|
||||
kwargs, kwargs_spec = substitute(n.kwargs)
|
||||
|
||||
# each token corresponds to a unique node
|
||||
# nodes with the same token can be substituted
|
||||
token = {
|
||||
"target": n.target,
|
||||
"args": args,
|
||||
"args_spec": args_spec,
|
||||
"kwargs": kwargs,
|
||||
"kwargs_spec": kwargs_spec,
|
||||
}
|
||||
|
||||
# hash substituted args to a number, do not hash specs because specs are not hashable
|
||||
# We need to add type into hash to avoid situations like:
|
||||
# hash((primals_2, 1.0)) == hash((primals_2, 1))
|
||||
hash_arg = hash(
|
||||
(tuple((a, type(a)) for a in args), tuple((a, type(a)) for a in kwargs))
|
||||
)
|
||||
hash_val = (n.target, hash_arg)
|
||||
|
||||
# check if a node has a substitute and can be eliminated
|
||||
hash_val_in_hash_env = hash_val in hash_env
|
||||
overwrite_due_to_mutation = False
|
||||
if hash_val_in_hash_env and token_map[hash_val] == token:
|
||||
duplicate_n_prev = hash_env[hash_val]
|
||||
if same_mutation_regions(n, duplicate_n_prev):
|
||||
env[n] = duplicate_n_prev
|
||||
continue
|
||||
else:
|
||||
# any futures duplicates should replace with n, not duplicate_n_prev
|
||||
overwrite_due_to_mutation = True
|
||||
|
||||
new_node = new_graph.node_copy(n, lambda x: env[x])
|
||||
env[n] = new_node
|
||||
if overwrite_due_to_mutation or not hash_val_in_hash_env:
|
||||
hash_env[hash_val] = new_node
|
||||
token_map[hash_val] = token
|
||||
|
||||
return new_graph
|
||||
|
||||
|
||||
def raise_getitems(gm: fx.GraphModule) -> fx.GraphModule:
|
||||
# Pre-create a list of nodes to iterate over, as modifying the node order
|
||||
# during the loop can lead to infinite loops if not handled properly.
|
||||
getitem_nodes = list(
|
||||
gm.graph.find_nodes(op="call_function", target=operator.getitem)
|
||||
)
|
||||
|
||||
# loop through getitem nodes in the graph and raise them to the parent node
|
||||
# in reverse order to preserve their original relative order
|
||||
for node in reversed(getitem_nodes):
|
||||
if len(node.all_input_nodes) != 1:
|
||||
raise AssertionError(
|
||||
f"expected node {node.name} to have 1 input node, got {len(node.all_input_nodes)}"
|
||||
)
|
||||
parent = node.all_input_nodes[0]
|
||||
parent.append(node)
|
||||
|
||||
gm.recompile()
|
||||
return gm
|
||||
|
||||
|
||||
def strip_overloads(gm: fx.GraphModule) -> None:
|
||||
"""
|
||||
Modifies the target of graph nodes in :attr:`gm` to strip overloads.
|
||||
|
||||
Args:
|
||||
gm(fx.GraphModule): The input Fx graph module to be modified
|
||||
"""
|
||||
for node in gm.graph.nodes:
|
||||
if isinstance(node.target, torch._ops.OpOverload):
|
||||
node.target = node.target.overloadpacket
|
||||
gm.recompile()
|
||||
|
||||
|
||||
def get_placeholders(graph: fx.Graph) -> fx.graph._node_list:
|
||||
return graph.find_nodes(op="placeholder")
|
||||
|
||||
|
||||
def get_outputs(graph: fx.Graph) -> list[fx.Node]:
|
||||
for node in graph.find_nodes(op="output"):
|
||||
return pytree.tree_leaves(node.args[0])
|
||||
raise AssertionError("No output node found")
|
||||
@@ -0,0 +1,501 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import random
|
||||
from contextlib import contextmanager
|
||||
from functools import partial
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
import torch.nn as nn
|
||||
import torch.utils._pytree as pytree
|
||||
from torch import SymInt
|
||||
from torch._decomp import get_decompositions
|
||||
from torch.fx.experimental.symbolic_shapes import bind_symbols
|
||||
|
||||
from .aot_autograd import aot_function, aot_module, make_boxed_compiler
|
||||
from .compile_utils import strip_overloads
|
||||
from .partitioners import (
|
||||
default_partition,
|
||||
draw_graph,
|
||||
min_cut_rematerialization_partition,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Generator, Sequence
|
||||
|
||||
from torch.fx.node import Node
|
||||
from torch.types import IntLikeType
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# These canonicalization are needed here (and not decompositions), as the ops
|
||||
# we're trying to canonicalize to CompositeImplicitAutograd.
|
||||
def _canonicalize(fx_g: fx.GraphModule) -> fx.GraphModule:
|
||||
for node in fx_g.graph.find_nodes(
|
||||
op="call_function", target=torch.ops.aten._to_copy
|
||||
):
|
||||
node.target = torch.ops.aten.to
|
||||
fx_g.recompile()
|
||||
return fx_g
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _disable_jit_autocast() -> Generator[None, None, None]:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
old_jit_autocast_flag = torch._C._jit_set_autocast_mode(False)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
torch._C._jit_set_autocast_mode(old_jit_autocast_flag)
|
||||
|
||||
|
||||
@make_boxed_compiler
|
||||
def ts_compile(fx_g: fx.GraphModule, inps: Sequence[Any]) -> torch.jit.ScriptModule:
|
||||
"""
|
||||
Compiles the :attr:`fx_g` with Torchscript compiler.
|
||||
|
||||
.. warning::
|
||||
This API is experimental and likely to change.
|
||||
|
||||
Args:
|
||||
fx_g(fx.GraphModule): The input Fx graph module to be compiled.
|
||||
|
||||
Returns:
|
||||
Torch scripted model.
|
||||
"""
|
||||
|
||||
with _disable_jit_autocast():
|
||||
strip_overloads(fx_g)
|
||||
|
||||
for node in fx_g.graph.find_nodes(
|
||||
op="call_function", target=torch.ops.aten._to_copy
|
||||
):
|
||||
if len(node.args) == 1 and len(node.kwargs) == 1 and "dtype" in node.kwargs:
|
||||
node.target = torch.ops.aten.to
|
||||
|
||||
for node in fx_g.graph.nodes:
|
||||
new_kwargs = {}
|
||||
for k, v in node.kwargs.items():
|
||||
if isinstance(v, torch.device):
|
||||
v = v.type
|
||||
new_kwargs[k] = v
|
||||
node.kwargs = new_kwargs
|
||||
|
||||
fx_g.graph.lint()
|
||||
|
||||
fx_g.recompile()
|
||||
|
||||
f = torch.jit.script(fx_g)
|
||||
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
torch._C._jit_pass_remove_mutation(f.graph)
|
||||
|
||||
f = torch.jit.freeze(f.eval())
|
||||
f = torch.jit.optimize_for_inference(f)
|
||||
if not any(isinstance(t, torch._subclasses.FakeTensor) for t in inps):
|
||||
f(*inps)
|
||||
return f
|
||||
|
||||
|
||||
def _draw_graph_compile(
|
||||
fx_g: fx.GraphModule, _: Any, name: str, clear_meta: bool = True
|
||||
) -> fx.GraphModule:
|
||||
print(fx_g.code)
|
||||
draw_graph(fx_g, name, clear_meta=clear_meta)
|
||||
return fx_g
|
||||
|
||||
|
||||
def draw_graph_compile(
|
||||
name: str,
|
||||
) -> Callable[[fx.GraphModule, list[Any]], fx.GraphModule]:
|
||||
return make_boxed_compiler(partial(_draw_graph_compile, name=name))
|
||||
|
||||
|
||||
@make_boxed_compiler
|
||||
def nop(fx_g: fx.GraphModule, _: Any) -> fx.GraphModule:
|
||||
"""
|
||||
Returns the :attr:`fx_g` Fx graph module as it is. This is a no-op compiler
|
||||
and can be used to check accuracy.
|
||||
|
||||
.. warning::
|
||||
This API is experimental and likely to change.
|
||||
|
||||
"""
|
||||
return fx_g
|
||||
|
||||
|
||||
class DebugInterpreter(fx.Interpreter):
|
||||
def run(
|
||||
self,
|
||||
*args: Any,
|
||||
initial_env: dict[Node, Any] | None = None,
|
||||
enable_io_processing: bool = True,
|
||||
) -> Any:
|
||||
self.symbol_mapping = bind_symbols(
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
self.module,
|
||||
*args,
|
||||
)
|
||||
return super().run(
|
||||
*args, initial_env=initial_env, enable_io_processing=enable_io_processing
|
||||
)
|
||||
|
||||
def run_node(self, n: Node) -> Any:
|
||||
def subst_symint(ni: IntLikeType) -> int:
|
||||
if not isinstance(ni, SymInt):
|
||||
return ni
|
||||
r = sympy.expand(ni.node.expr.xreplace(self.symbol_mapping))
|
||||
if not r.is_number:
|
||||
raise AssertionError(f"expected r to be a number, got {r}")
|
||||
return int(r)
|
||||
|
||||
def subst_symint_tuple(nis: tuple[IntLikeType, ...]) -> tuple[int, ...]:
|
||||
return tuple(subst_symint(ni) for ni in nis)
|
||||
|
||||
def check_significant_strides(a: torch.Tensor, b: torch.Tensor) -> bool:
|
||||
if subst_symint(a.numel()) > 0:
|
||||
for idx in range(a.ndim):
|
||||
if (
|
||||
subst_symint(a.stride(idx)) != b.stride(idx)
|
||||
and subst_symint(a.size(idx)) > 1
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
def check(nv: torch.Tensor, rv: torch.Tensor, desc: Callable[[], str]) -> None:
|
||||
if not callable(desc):
|
||||
raise AssertionError(f"expected desc to be callable, got {type(desc)}")
|
||||
if nv.dtype != rv.dtype:
|
||||
raise AssertionError(f"{desc()}: {nv.dtype} != {rv.dtype}")
|
||||
if subst_symint_tuple(nv.size()) != rv.size():
|
||||
raise AssertionError(
|
||||
f"{desc()}: {nv.size()} aka {subst_symint_tuple(nv.size())} != {rv.size()}"
|
||||
)
|
||||
same_strides = check_significant_strides(nv, rv)
|
||||
if not same_strides:
|
||||
raise AssertionError(
|
||||
f"{desc()}: {nv.stride()} aka {subst_symint_tuple(nv.stride())} != {rv.stride()}"
|
||||
)
|
||||
|
||||
r = super().run_node(n)
|
||||
if "val" in n.meta:
|
||||
n_vals, _n_spec = pytree.tree_flatten(n.meta["val"])
|
||||
r_vals, _r_spec = pytree.tree_flatten(r)
|
||||
# TODO: There is some sort of problem where we record that an
|
||||
# operator returned a tuple/list, and then later it turns out the
|
||||
# real version of the operator returned a list/tuple. Need to
|
||||
# figure out what's actually going on here, the error itself is
|
||||
# harmless enough as we only getitem out the outputs.
|
||||
# assert n_spec == r_spec, f"{n_spec} != {r_spec}"
|
||||
if len(n_vals) != len(r_vals):
|
||||
raise AssertionError(f"{len(n_vals)} != {len(r_vals)}")
|
||||
for i, nv, rv in zip(range(len(n_vals)), n_vals, r_vals):
|
||||
if not isinstance(rv, torch.Tensor):
|
||||
continue
|
||||
check(nv, rv, lambda: f"output {i} where {self.symbol_mapping}")
|
||||
return r
|
||||
|
||||
|
||||
@make_boxed_compiler
|
||||
def debug_nop(
|
||||
fx_g: fx.GraphModule, _: Any
|
||||
) -> Callable[[DebugInterpreter, Any, dict[Node, Any] | None, bool], Any]:
|
||||
"""
|
||||
Returns a (slow) interpreter over the FX graph module that also checks
|
||||
various debugging properties (e.g., that tracing strides matched real
|
||||
strides.)
|
||||
"""
|
||||
return DebugInterpreter(fx_g).run
|
||||
|
||||
|
||||
@make_boxed_compiler
|
||||
def simple_ts_compile(fx_g: fx.GraphModule, _: Any) -> torch.jit.ScriptModule:
|
||||
strip_overloads(fx_g)
|
||||
f = torch.jit.script(fx_g)
|
||||
f = torch.jit.freeze(f.eval())
|
||||
return f
|
||||
|
||||
|
||||
def nnc_jit(f: Callable[..., Any]) -> Callable[..., Any]:
|
||||
return aot_function(f, simple_ts_compile)
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
default_decompositions = {
|
||||
aten.detach,
|
||||
aten.gelu_backward,
|
||||
aten.leaky_relu_backward,
|
||||
aten.sigmoid_backward,
|
||||
aten.threshold_backward,
|
||||
aten.hardtanh_backward,
|
||||
aten.hardsigmoid_backward,
|
||||
aten.hardswish_backward,
|
||||
aten.tanh_backward,
|
||||
aten.silu_backward,
|
||||
aten.elu_backward,
|
||||
aten.cudnn_batch_norm,
|
||||
aten.cudnn_batch_norm_backward,
|
||||
aten.masked_fill.Scalar,
|
||||
aten.masked_fill.Tensor,
|
||||
aten.elu,
|
||||
aten.leaky_relu,
|
||||
aten.hardtanh,
|
||||
aten.hardswish,
|
||||
aten.hardsigmoid,
|
||||
aten.conj_physical,
|
||||
aten.is_same_size,
|
||||
}
|
||||
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
default_decompositions = get_decompositions(default_decompositions)
|
||||
|
||||
|
||||
@make_boxed_compiler
|
||||
def print_compile(fx_g: fx.GraphModule, _: Any) -> fx.GraphModule:
|
||||
print(fx_g.code)
|
||||
return fx_g
|
||||
|
||||
|
||||
def memory_efficient_fusion(
|
||||
fn: Callable[_P, _R] | nn.Module,
|
||||
**kwargs: Any,
|
||||
) -> Callable[_P, _R] | nn.Module:
|
||||
"""
|
||||
Wrapper function over :func:`aot_function` and :func:`aot_module` to perform
|
||||
memory efficient fusion. It uses the
|
||||
:func:`min_cut_rematerialization_partition` partitioner to perform efficient
|
||||
recomputation. It uses NVFuser to compile the generated forward and backward
|
||||
graphs.
|
||||
|
||||
.. warning::
|
||||
This API is experimental and likely to change.
|
||||
|
||||
Args:
|
||||
fn (Union[Callable, nn.Module]): A Python function or a ``nn.Module``
|
||||
that takes one or more arguments. Must return one or more Tensors.
|
||||
**kwargs: Any other overrides you want to make to the settings
|
||||
|
||||
Returns:
|
||||
Returns a ``Callable`` or ``nn.Module`` that retains the eager behavior
|
||||
of the original :attr:`fn`, but whose forward and backward graphs have
|
||||
gone through recomputation optimizations, and the graphs have been
|
||||
compiled with nvfuser.
|
||||
|
||||
"""
|
||||
config = {
|
||||
"fw_compiler": ts_compile,
|
||||
"bw_compiler": ts_compile,
|
||||
"partition_fn": min_cut_rematerialization_partition,
|
||||
"decompositions": default_decompositions,
|
||||
}
|
||||
config.update(kwargs)
|
||||
if isinstance(fn, torch.nn.Module):
|
||||
return aot_module(fn, **config) # pyrefly: ignore[bad-argument-type]
|
||||
else:
|
||||
return aot_function(fn, **config) # pyrefly: ignore[bad-argument-type]
|
||||
|
||||
|
||||
def debug_compile(
|
||||
fx_g: fx.GraphModule, inps: Sequence[torch.Tensor]
|
||||
) -> torch.jit.ScriptModule:
|
||||
fx_g.to_folder("foo")
|
||||
print(
|
||||
f"""
|
||||
##############################################################
|
||||
# To minimize FX graph, copy and paste the below and run it #
|
||||
##############################################################
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from functorch.compile import minifier, check_nvfuser_subprocess, check_nvfuser_correctness_subprocess
|
||||
|
||||
inps = {[(i.shape, i.dtype) for i in inps]}
|
||||
inps = [torch.ones(shape, dtype=dtype, device='cuda') for (shape, dtype) in inps]
|
||||
from foo import FxModule
|
||||
mod = FxModule().cuda()
|
||||
|
||||
with torch.jit.fuser("fuser2"):
|
||||
# check_nvfuser_subprocess can be replaced with check_nvfuser_correctness_subprocess
|
||||
minifier(fx.symbolic_trace(mod), inps, check_nvfuser_subprocess)
|
||||
"""
|
||||
)
|
||||
# pyrefly: ignore[missing-import, missing-module-attribute]
|
||||
from foo import FxModule
|
||||
|
||||
FxModule().cuda()(*inps)
|
||||
|
||||
return ts_compile(fx_g, inps)
|
||||
|
||||
|
||||
graph_index: int = 0
|
||||
|
||||
|
||||
def get_inputs(input_data_path: str) -> list[torch.Tensor]:
|
||||
"""
|
||||
Return a random input for the given inputs meta generated from _save_fx_default.
|
||||
"""
|
||||
inputs: list[torch.Tensor] = []
|
||||
with open(input_data_path, "rb") as f:
|
||||
inputs_meta = pickle.load(f)
|
||||
inputs = []
|
||||
for meta in inputs_meta:
|
||||
if len(meta) == 1:
|
||||
type = meta
|
||||
input_ = type(random.random())
|
||||
else:
|
||||
type, shape, _stride, dtype, device = meta
|
||||
if dtype in {
|
||||
torch.int,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.bool,
|
||||
torch.int,
|
||||
torch.uint8,
|
||||
int,
|
||||
float,
|
||||
}:
|
||||
input_ = torch.randint(0, 1, shape, dtype=dtype, device=device)
|
||||
else:
|
||||
input_ = torch.rand(shape, dtype=dtype, device=device)
|
||||
inputs.append(input_)
|
||||
return inputs
|
||||
|
||||
|
||||
def _save_fx_default(
|
||||
current_name: str,
|
||||
folder_name: str,
|
||||
dump_example_input: bool,
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: list[torch.Tensor],
|
||||
) -> nn.Module:
|
||||
"""
|
||||
The forward, backward, and joint computation graph will be stored in
|
||||
{folder_name}/{current_name}/{current_name}_forward_{graph_index},
|
||||
{folder_name}/{current_name}/{current_name}_backward_{graph_index}, and
|
||||
{folder_name}/{current_name}/{current_name}_joint_{graph_index} respectively.
|
||||
The input shape of the graphs will be stored in the .input files.
|
||||
These files can be loaded with pickle,
|
||||
and is a list of format (type, shape, stride, dtype, device).
|
||||
In the case of type = int or float, it is just (type,).
|
||||
For joint graph input, it is a nested list [[],[]]
|
||||
where the two inner lists have the same format.
|
||||
If dump_example_input is True, example_inputs will be stored in .pt file.
|
||||
Since each function might produce multiple graphs,
|
||||
the graph_index is used to distinguish difference graphs
|
||||
"""
|
||||
from functorch.compile import aot_module_simplified
|
||||
|
||||
def get_input_meta(args: Any) -> list[Any]:
|
||||
input_meta = []
|
||||
if len(args) > 0 and isinstance(args[0], tuple): # joint input
|
||||
input_meta += get_input_meta(args[0])
|
||||
input_meta += get_input_meta(args[1])
|
||||
return input_meta
|
||||
for arg in args:
|
||||
if type(arg) is int or type(arg) is float:
|
||||
input_meta.append((type(arg),))
|
||||
else:
|
||||
input_meta.append(
|
||||
(type(arg), arg.shape, arg.stride(), arg.dtype, arg.device)
|
||||
)
|
||||
return input_meta
|
||||
|
||||
def graph_saver_helper(
|
||||
gm_to_save: fx.GraphModule, args: Any, type_name: str
|
||||
) -> None:
|
||||
global graph_index
|
||||
if len(gm_to_save.graph.nodes) == 0:
|
||||
log.log(
|
||||
logging.WARNING,
|
||||
"No nodes in graph {%s}_{%s}_{%s}.",
|
||||
current_name,
|
||||
type_name,
|
||||
graph_index,
|
||||
)
|
||||
return
|
||||
|
||||
gm = copy.deepcopy(gm_to_save)
|
||||
gm.graph.set_codegen(torch.fx.graph.CodeGen()) # remove codegen
|
||||
gm.recompile()
|
||||
|
||||
input_meta = get_input_meta(args)
|
||||
|
||||
os.makedirs(f"{folder_name}/{current_name}", exist_ok=True)
|
||||
gm.to_folder(
|
||||
f"{folder_name}/{current_name}/{current_name}_{type_name}_{graph_index}"
|
||||
)
|
||||
with open(
|
||||
f"{folder_name}/{current_name}/{current_name}_{type_name}_{graph_index}/{current_name}_{type_name}_{graph_index}.input",
|
||||
"wb",
|
||||
) as f:
|
||||
pickle.dump(input_meta, f)
|
||||
if dump_example_input:
|
||||
torch.save(
|
||||
args,
|
||||
f"{folder_name}/{current_name}/{current_name}_{type_name}_{graph_index}/{current_name}_{type_name}_{graph_index}.pt", # noqa: B950
|
||||
) # noqa: E501
|
||||
|
||||
def graph_saver_forward(
|
||||
gm: fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> fx.GraphModule:
|
||||
graph_saver_helper(gm, example_inputs, "forward")
|
||||
return gm
|
||||
|
||||
def graph_saver_backward(
|
||||
gm: fx.GraphModule, example_inputs: list[torch.Tensor]
|
||||
) -> fx.GraphModule:
|
||||
graph_saver_helper(gm, example_inputs, "backward")
|
||||
global graph_index
|
||||
graph_index += 1
|
||||
return gm
|
||||
|
||||
def graph_saver_joint(
|
||||
gm: fx.GraphModule, joint_args: list[torch.Tensor]
|
||||
) -> tuple[fx.GraphModule, fx.GraphModule]:
|
||||
graph_saver_helper(gm, joint_args, "joint")
|
||||
return default_partition(gm, joint_args) # pyrefly: ignore[missing-argument]
|
||||
|
||||
# pyrefly: ignore[bad-return]
|
||||
return aot_module_simplified(
|
||||
gm,
|
||||
example_inputs,
|
||||
fw_compiler=graph_saver_forward, # pyrefly: ignore[bad-argument-type]
|
||||
bw_compiler=graph_saver_backward, # pyrefly: ignore[bad-argument-type]
|
||||
partition_fn=graph_saver_joint,
|
||||
decompositions=default_decompositions, # pyrefly: ignore[bad-argument-type]
|
||||
)
|
||||
|
||||
|
||||
# WARNING: This isn't tested anywhere!!
|
||||
def graph_dumper_aot(
|
||||
current_name: str, folder_name: str, dump_example_input: bool = False
|
||||
) -> Callable[[bool, nn.Module], Any]:
|
||||
"""
|
||||
Dump the forward, backward, and joint computation graph.
|
||||
Example Usage:
|
||||
save_fx_func = graph_dumper_aot(current_name, folder_name, dump_example_input = False)
|
||||
optimize_ctx = torchdynamo.optimize(
|
||||
save_fx_func
|
||||
)
|
||||
with torch.enable_grad():
|
||||
with optimize_ctx:
|
||||
result = forward_and_backward_pass(model, example_inputs)
|
||||
"""
|
||||
global graph_index
|
||||
graph_index = 0
|
||||
return partial(_save_fx_default, current_name, folder_name, dump_example_input)
|
||||
@@ -0,0 +1,444 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
|
||||
"""
|
||||
Global flags for aot autograd
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from typing import Literal, TYPE_CHECKING
|
||||
|
||||
from torch.utils._config_module import Config, install_config_module
|
||||
|
||||
|
||||
# [@compile_ignored: debug]
|
||||
_save_config_ignore = [
|
||||
# callable not serializable
|
||||
"joint_custom_pass",
|
||||
# callable configs with uuid() for caching, or raw callables
|
||||
"activation_memory_budget_runtime_estimator",
|
||||
"activation_memory_budget_solver",
|
||||
]
|
||||
|
||||
|
||||
# Converts torch rng ops to their functional philox rng equivalents. Note that
|
||||
# we functionalize only CUDA rng ops today.
|
||||
functionalize_rng_ops = False
|
||||
|
||||
# can be useful for debugging if we are incorrectly creating meta fake tensors
|
||||
fake_tensor_allow_meta = os.environ.get("FAKE_ALLOW_META", "1") != "0"
|
||||
|
||||
# Enables optional asserts in hotpath code to check for errors. If
|
||||
# you are seeing weird accuracy problems, try turning this on.
|
||||
# This is currently off by default as it will harm tracing time,
|
||||
# but it is on by default for aot_eager.
|
||||
debug_assert = False
|
||||
|
||||
debug_partitioner = os.environ.get("AOT_PARTITIONER_DEBUG", "0") != "0"
|
||||
|
||||
# See # NOTE [Export custom triton op]
|
||||
decompose_custom_triton_ops = True
|
||||
|
||||
static_weight_shapes = True
|
||||
|
||||
# See https://github.com/pytorch/pytorch/issues/141881
|
||||
# Tells partitioner that parameters are free to save for backward.
|
||||
treat_parameters_as_free_to_save = True
|
||||
|
||||
# Applies CSE to the graph before partitioning
|
||||
cse = True
|
||||
|
||||
from torch._environment import is_fbcode
|
||||
|
||||
|
||||
enable_autograd_cache: bool = Config(
|
||||
justknob="pytorch/remote_cache:enable_local_autograd_cache",
|
||||
env_name_force="TORCHINDUCTOR_AUTOGRAD_CACHE",
|
||||
default=True,
|
||||
)
|
||||
|
||||
autograd_cache_allow_custom_autograd_functions: bool = Config(
|
||||
env_name_force="TORCHINDUCTOR_AUTOGRAD_CACHE_ALLOW_CUSTOM_AUTOGRAD", default=False
|
||||
)
|
||||
|
||||
# For now, this is just for enabling unit testing in test_aot_autograd_cache.py
|
||||
# We will either make this the default with AOTAutogradCache, or
|
||||
# we'll just use it in the precompile flow. So there's no
|
||||
# need to add env vars or make it configurable
|
||||
bundled_autograd_cache: bool = False
|
||||
|
||||
bypass_autograd_cache_key: bool = False
|
||||
|
||||
# Whether or not to normalize placeholder names in graphs
|
||||
# from dynamo in AOTAutogradCache
|
||||
autograd_cache_normalize_inputs = not is_fbcode()
|
||||
|
||||
# Enable debug mode at first invocation to check if custom ops are valid.
|
||||
# When enabled, this checks that custom operators don't violate aliasing constraints.
|
||||
#
|
||||
# check_custom_op_aliasing: Controls whether to run the custom op aliasing check at all.
|
||||
# - When True: The check runs on first invocation of compiled functions.
|
||||
# - When False: The check is skipped entirely.
|
||||
#
|
||||
# error_on_custom_op_aliasing: Controls behavior when a violation is detected.
|
||||
# Only has effect when check_custom_op_aliasing is True.
|
||||
# - When True: Raises RuntimeError on aliasing violations.
|
||||
# - When False: Emits UserWarning on aliasing violations.
|
||||
#
|
||||
# Deprecated: Custom ops returning aliased outputs is deprecated and will
|
||||
# become an error in PyTorch 2.12. Currently error_on_custom_op_aliasing
|
||||
# is True only in CI.
|
||||
check_custom_op_aliasing = True
|
||||
error_on_custom_op_aliasing = bool(os.getenv("CI"))
|
||||
|
||||
|
||||
def remote_autograd_cache_default() -> bool | None:
|
||||
if os.environ.get("TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE") == "1":
|
||||
return True
|
||||
if os.environ.get("TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE") == "0":
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
enable_remote_autograd_cache = remote_autograd_cache_default()
|
||||
|
||||
|
||||
# When AOTAutograd regenerates aliased graph outputs,
|
||||
# attempt to use functionalization's view-replay logic
|
||||
# before falling back to the autograd engine's view replay or as_strided.
|
||||
# This can have some perf implications
|
||||
# (although for many models this will not matter).
|
||||
# (1) If you have many view ops chained together, replaying all of them
|
||||
# at runtime can have more overhead compared to a single as_strided call
|
||||
# (2) If you are doing training, AsStridedBackward is quite slow,
|
||||
# and the individual view op backward formulas will likely be faster.
|
||||
# (3) Some backends like XLA do not support as_strided
|
||||
|
||||
# Temporary hack: disable this flag for internal
|
||||
# (needed to fix an internal issue while avoiding bumping XLA pin)
|
||||
# eventually: either default this config to false completely
|
||||
# once XLA pin update works,
|
||||
# or default config to true and fix relevant bugs
|
||||
|
||||
|
||||
# View replay is currently not compatible with AOTAutogradCache, since
|
||||
# FunctionalTensors are not serializable. We'll need to make them
|
||||
# serializable before enabling warm cache with this config turned on.
|
||||
view_replay_for_aliased_outputs = not is_fbcode()
|
||||
|
||||
# Restricts the amount of computation AOTAutograd can do.
|
||||
# NB: We have essentially disabled this heuristic now. However, this is kept
|
||||
# here for now in case it's useful. Setting it low can artificially reduce the
|
||||
# amount of recomputation AOTAutograd performs, although not in any kind of
|
||||
# principled way.
|
||||
max_dist_from_bw = 1000
|
||||
|
||||
|
||||
# Bans recomputation of nodes that are reading from nodes that are far before
|
||||
# the current node
|
||||
ban_recompute_used_far_apart = True
|
||||
# Breaks up long chain of fusible ops, as otherwise we can have an arbitrarily
|
||||
# long chain of recomputation in the backwards pass.
|
||||
ban_recompute_long_fusible_chains = True
|
||||
# Bans recomputation of nodes that must be materialized in the backwards pass
|
||||
# (used by a non-fusible node)
|
||||
ban_recompute_materialized_backward = True
|
||||
# Chooses to ban recomputation of nodes based off an allowlist. Setting it to
|
||||
# False changes it to use a denylist. Main change is on operators like
|
||||
# sort/pool/stuff that isn't cheap enough to be fusible for free but also isn't
|
||||
# that expensive
|
||||
ban_recompute_not_in_allowlist = True
|
||||
# Chooses to ban recomputation of reductions. This is generally a good idea, as
|
||||
# the result of reductions is generally very small but recomputing reductions in
|
||||
# a fusion can be expensive.
|
||||
ban_recompute_reductions = True
|
||||
# Prevents the partitioner from ever saving views (i.e. always recompute them).
|
||||
# Generally a good idea since views are free to recompute.
|
||||
recompute_views = False
|
||||
# Set this flag to enable considering non-built-in ops, including triton and custom
|
||||
# ops, for recomputation during the knapsack optimization solver.
|
||||
is_non_builtin_to_include = False
|
||||
|
||||
# Rematerialize AC nodes for graphs with forward+loss+backward in one graph.
|
||||
# This optimization minimizes activation checkpoint node lifetimes by computing them
|
||||
# just-in-time. For AC nodes only used in backward, they are deferred to backward region
|
||||
# instead of being computed and saved in forward. This reduces peak memory usage.
|
||||
# Note: This only applies to forward+loss+backward graphs where torch.autograd.grad is allowed
|
||||
# in the graph. Joint graphs (standard AOTAutograd) use the partitioner instead.
|
||||
remat_using_tags_for_fwd_loss_bwd_graph = True
|
||||
|
||||
# By default, the partitioner is purely trying to optimize for runtime (although
|
||||
# it should always use less memory than eager)
|
||||
# This knob controls the partitioner to make that tradeoff for you, choosing the
|
||||
# fastest option that saves less activations than the memory budget.
|
||||
# Specifically, 0.0 corresponds to the activation memory from applying
|
||||
# activation checkpointing to the full compiled region, and 1.0 corresponds to
|
||||
# the activation memory from the default runtime-optimized strategy. So, 0.4
|
||||
# would result in a strategy that saves 40% of the activations compared to the
|
||||
# default strategy.
|
||||
# It solves a 0-1 knapsack to find the minimum recompute necessary to stay below
|
||||
# the activation memory budget.
|
||||
# NOTE: This *cannot* be treated as
|
||||
activation_memory_budget = 1.0
|
||||
|
||||
# This controls how we estimate the runtime when deciding what the cheapest
|
||||
# operators to recompute are. The 3 options are
|
||||
# "flops": Bases it off of the flop count provided by torch.utils.flop_counter
|
||||
# "profile": Benchmarks each operator to come up with a runtime
|
||||
# "testing": Returns 1 for everything
|
||||
activation_memory_budget_runtime_estimator = "flops"
|
||||
|
||||
# This controls the solver used for the 0-1 knapsack. By default we use a
|
||||
# quantized DP solution ("dp"). The other approaches are a "greedy", an "ilp"
|
||||
# (which has a scipy dependency) and "dp_knapsack_sliding_hirschberg", which
|
||||
# used memory-efficient quantized DP solution
|
||||
activation_memory_budget_solver = "dp"
|
||||
|
||||
# This dumps out a SVG visualization of the expected runtime vs. activation
|
||||
# memory tradeoffs for all memory budget values from 0 to 1 in increments of
|
||||
# 0.5. See an example here:
|
||||
# https://github.com/pytorch/pytorch/pull/126320#discussion_r1625104015
|
||||
visualize_memory_budget_pareto = (
|
||||
os.environ.get("PARTITIONER_MEMORY_BUDGET_PARETO", "0") == "1"
|
||||
)
|
||||
|
||||
# This controls the directory in which to dump the SVG plot with the pareto
|
||||
# frontier of the activation checkpointing memory-vs-runtime tradeoffs.
|
||||
memory_budget_pareto_dir = os.environ.get("PARTITIONER_MEMORY_BUDGET_PARETO_DIR")
|
||||
|
||||
# Sets all of the ban_recompute heuristics to False except ban_recompute_reductions
|
||||
# Generally, this will probably result in some memory improvement, but at the
|
||||
# cost of some performance
|
||||
aggressive_recomputation = False
|
||||
|
||||
# activation offloading enablement (testing purpose)
|
||||
enable_activation_offloading = False
|
||||
|
||||
# activation offloading with separate CUDA stream
|
||||
activation_offload_separate_stream = False
|
||||
|
||||
# activation offloading wait sinking when using separate stream (fwd graph)
|
||||
activation_offload_sink_wait = False
|
||||
|
||||
# activation reloading with prefetching when using separate streams (bwd graph)
|
||||
activation_reload_prefetch = False
|
||||
|
||||
# CPU ↔ GPU bandwidth in GB/s, used to estimate transfer times for prefetch
|
||||
# scheduling. This is hardware-specific and should be set by the user.
|
||||
activation_offload_cpu_gpu_bw: float = 50.0
|
||||
|
||||
# If FakeTensor.data_ptr() should error.
|
||||
# This option is independent of AOTAutograd and torch.compile, but our policy
|
||||
# is to turn it off during torch.compile.
|
||||
fake_tensor_allow_unsafe_data_ptr_access = True
|
||||
|
||||
# Unlifts effect tokens from the inputs/outputs in the traced graph and instead
|
||||
# inserts make_token/sink_token calls in the graph to create tokens and then
|
||||
# sink them at the end. Note that this means the graph is no longer functional
|
||||
# which may lead to silent errors unless the backend knows how to handle the
|
||||
# tokens.
|
||||
unlift_effect_tokens = False
|
||||
|
||||
# NOTE: [The default layout constraint for custom operators.]
|
||||
# This must be the name of one of the layout constraint tags
|
||||
# (that is, one of {"needs_fixed_stride_order", "flexible_layout"}),
|
||||
# If the custom op does not have a layout constraint tag already
|
||||
# then we assume the following applies.
|
||||
#
|
||||
# This config is respected by Inductor and we recommend other backends also
|
||||
# respect it.
|
||||
# This config is in torch._functorch and not torch._inductor because it affects
|
||||
# ProxyTensor tracing.
|
||||
custom_op_default_layout_constraint: Literal[
|
||||
"needs_exact_strides", "needs_fixed_stride_order", "flexible_layout"
|
||||
] = "needs_exact_strides"
|
||||
|
||||
|
||||
# Run aot eager decomp partition with CrossRefFakeMode
|
||||
# options = False, "all", "custom_ops"
|
||||
fake_tensor_crossref = False
|
||||
|
||||
# This mode specifies that we should also keep track of the real
|
||||
# tensor along with the fake tensor, and do real compute. While
|
||||
# seemingly this eliminates the whole point of fake tensors, there are
|
||||
# two obvious use cases for it:
|
||||
#
|
||||
# 1. When users call item()/other data dependent operations,
|
||||
# if we propagate_real_tensors we are able to determine what
|
||||
# the true value is and keep going.
|
||||
#
|
||||
# 2. It can be useful for testing, when you want to see if the fake
|
||||
# and real tensors agree with each other. (Note that there are
|
||||
# currently known inaccuracies in how we clone real tensors, that
|
||||
# would have to be tightened up for this to be useful in this
|
||||
# case.)
|
||||
#
|
||||
# Note that fake tensors are typically understood to be cheap to store
|
||||
# indefinitely, so we tend to hold on to them longer than we would
|
||||
# hold onto the real tensors. So we also support you explicitly
|
||||
# deallocating the real tensor associated with a fake tensor, at which
|
||||
# point we will stop propagating real tensors.
|
||||
#
|
||||
# One more thing: when you provide a real tensor to fakeify, we will
|
||||
# clone it, so that we can safely perform mutations on it if necessary.
|
||||
# This will increase live memory usage. This could potentially be
|
||||
# optimized by using COW. We also currently do not faithfully
|
||||
# maintain autograd metadata on the real tensor; this is fine because
|
||||
# AOTAutograd will only use the fake tensor to determine leafness/etc
|
||||
# of tensors in question.
|
||||
fake_tensor_propagate_real_tensors = False
|
||||
|
||||
# AOTDispatcher traces out a backward graph at the time of the forward pass.
|
||||
# This flag controls whether or not that backward graph gets autocast behavior
|
||||
# applied to it.
|
||||
#
|
||||
# The options are either:
|
||||
# - "same_as_forward". We assume that the backward of the torch.compile'ed region
|
||||
# will be run under the same autocast context manager that the region was run
|
||||
# under. This is equivalent to running the following code in eager:
|
||||
#
|
||||
# with torch.amp.autocast(...):
|
||||
# y = region(x)
|
||||
# ...
|
||||
# z.backward()
|
||||
#
|
||||
# - "off". We assume that the backward of the torch.compile'd region will
|
||||
# not be run under any autocast context managers.
|
||||
# This is equivalent to running the following code in eager:
|
||||
#
|
||||
# with torch.amp.autocast(...):
|
||||
# y = region(x)
|
||||
# ...
|
||||
# z.backward()
|
||||
#
|
||||
# - or a list of kwargs dicts that represent an autocast context manager to turn
|
||||
# on during the backward pass.
|
||||
#
|
||||
# e.g. [{"device_type": "cuda"}] is equivalent to running the following code in eager:
|
||||
#
|
||||
# y = region(x)
|
||||
# ...
|
||||
# with torch.amp.autocast(device="cuda"):
|
||||
# z.backward()
|
||||
backward_pass_autocast = "same_as_forward"
|
||||
|
||||
# This controls whether we collect donated buffers. This flag must be set
|
||||
# False if a user wants to retain_graph=True for backward.
|
||||
donated_buffer = not is_fbcode()
|
||||
|
||||
# Controls the default graph output format used by draw_graph
|
||||
# Supported formats are defined here https://graphviz.org/docs/outputs/
|
||||
torch_compile_graph_format = os.environ.get("TORCH_COMPILE_GRAPH_FORMAT", "svg")
|
||||
|
||||
# Valid only if fake_tensor_propagate_real_tensors = True; if a fake-real
|
||||
# kernel mismatch is detected, bypasses by making a fake kernel from the
|
||||
# real tensor outputs.
|
||||
generate_fake_kernels_from_real_mismatches = False
|
||||
|
||||
# When there are device mismatches in FakeTensor device propagation,
|
||||
# prefer a specific device type over others. This is particularly useful
|
||||
# in full compiled mode where intermediate tensors with device mismatches
|
||||
# represent only logical differences during compilation - these intermediate
|
||||
# tensors will never physically materialize in the binary execution, so the
|
||||
# device mismatch is not a real runtime concern. Enabling this allows the
|
||||
# compiler to proceed with compilation by choosing the preferred device type
|
||||
# for consistency. For example, set to "mtia" to prefer MTIA devices over
|
||||
# CPU, or "cuda" to prefer CUDA devices over CPU.
|
||||
fake_tensor_prefer_device_type: str | None = None
|
||||
|
||||
# CUDAGraph safe run_with_rng functionalization.
|
||||
# TODO: turn on by default
|
||||
graphsafe_rng_functionalization = True
|
||||
|
||||
# Whether or not to eagerly compile the backward
|
||||
# used by AOT compile and other settings
|
||||
# TODO: once AOT compile calls aot autograd directly instead of
|
||||
# through compile_fx, we can remove this
|
||||
force_non_lazy_backward_lowering = False
|
||||
|
||||
# only for testing, used to turn functionalization off in AOTDispatcher
|
||||
_test_disable_functionalization = True
|
||||
|
||||
# Error on BypassAOTAutogradCache instead of just a warning
|
||||
# Used for tests
|
||||
strict_autograd_cache = False
|
||||
|
||||
# Note [Recomputing collectives in the partitioner]
|
||||
# The purpose of this config is as follows:
|
||||
# - We have many passes in the compiler (min-cut partitioning, DCE, etc)
|
||||
# which can reorder or delete duplicate nodes in the graph
|
||||
# - If any of these passes reorder/delete/duplicate a collective
|
||||
# in a setting where the compiler is being run independently on multiple
|
||||
# ranks, we run the risk that the compiler will make a different decision on
|
||||
# different ranks, resulting in a NCCL hang when using torch.compile
|
||||
# To handle this, we will (by default) ensure that collectives are not modified
|
||||
# by the compiler.
|
||||
#
|
||||
# A few examples:
|
||||
# - don't dead-code-eliminate collectives
|
||||
# (in case they are dead on rank i but not rank j)
|
||||
# - don't recompute collectives in partitioning
|
||||
# (in case we recompute on rank i but not rank j)
|
||||
#
|
||||
# Today this flag **must** be set to false, but eventually
|
||||
# we want the option to set it to true.
|
||||
# In order to potentially optimize collectives, we'll need the compiler
|
||||
# to broadcast information across ranks at compile time to ensure
|
||||
# that any decisions on collectives are made consistently.
|
||||
unsafe_allow_optimization_of_collectives = False
|
||||
|
||||
# See Note [AOTAutograd Tangent Subclassness for mutated inputs]
|
||||
# TODO(ivankobzarev): Remove this config, being able to deduce it compile time.
|
||||
disable_guess_zero_tangent_for_mutated_input_subclass = False
|
||||
|
||||
# See Note [Tangents memory format]
|
||||
# By default tangents strideness is guessed to be contiguous,
|
||||
# At runtime non contiguous tangents will be coerced to be contiguous.
|
||||
# This config changes this guess for tangents strides to be the same as outputs.
|
||||
# TODO(ivankobzarev): Remove this config once extra memory usage is investigated.
|
||||
guess_tangent_strides_as_outputs = not is_fbcode()
|
||||
|
||||
|
||||
# This is a temporary config to ensure all ranks take the same decision in the partitioner
|
||||
# it will ultimately be removed once we share size_hints across ranks through compiler collectives
|
||||
_sync_decision_cross_ranks = False
|
||||
|
||||
# By default apply inlined saved_tensors_hooks only for "donated" buffers.
|
||||
# "donated" buffers are invisible to the user, they are intermediates of the forward graph.
|
||||
# Applying saved tensors hooks for memory optimizations only for intermediates
|
||||
# guarantees that original saved tensors could be deallocated.
|
||||
# This config enables saved_tensors_hooks are applied for **all** saved tensors,
|
||||
# that could include inputs, parameters, outputs.
|
||||
# "donated" - applied only to saved intermediates of the graph
|
||||
# "no_static" - applied to all saved but not "static"
|
||||
# (this includes parameters and user marked as static)
|
||||
# "all" - no filtering, everything saved for backward.
|
||||
saved_tensors_hooks_filtering_mode = "donated"
|
||||
|
||||
|
||||
# This callback is invoked on the joint graph before partitioning
|
||||
joint_custom_pass: Callable = None # type: ignore[assignment]
|
||||
|
||||
force_autograd_cache = False
|
||||
|
||||
# Note [Selective Decomposition]
|
||||
# This config allows selective decomposition of certain operators in the graph.
|
||||
# When True, it does NOT decompose any nodes, except those nodes that users explicitly
|
||||
# annotated with regional inductor compile. Please read torch.fx.passes.regional_inductor
|
||||
# on to explicitly annotate. This is currently only used by inductor lite mode.
|
||||
selective_decompose: bool = False
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.utils._config_typing import * # noqa: F401, F403
|
||||
|
||||
|
||||
# adds patch, save_config, invalid config checks, etc
|
||||
install_config_module(sys.modules[__name__])
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
The APIs in this file are exposed as `functorch.*`. They are thin wrappers
|
||||
around the torch.func.* APIs that have deprecation warnings -- we're trying
|
||||
to move people to the torch.func.* equivalents.
|
||||
|
||||
NB: We don't use *args, **kwargs in the signatures because that changes the
|
||||
documentation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import textwrap
|
||||
import warnings
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch._functorch.apis as apis
|
||||
import torch._functorch.eager_transforms as _impl
|
||||
import torch._functorch.make_functional as _nn_impl
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch._functorch.eager_transforms import argnums_t
|
||||
from torch._functorch.vmap import in_dims_t, out_dims_t
|
||||
|
||||
|
||||
def get_warning(
|
||||
api: str, new_api: str | None = None, replace_newlines: bool = False
|
||||
) -> str:
|
||||
if new_api is None:
|
||||
new_api = f"torch.func.{api}"
|
||||
warning = (
|
||||
f"We've integrated functorch into PyTorch. As the final step of the \n"
|
||||
f"integration, `functorch.{api}` is deprecated as of PyTorch \n"
|
||||
f"2.0 and will be deleted in a future version of PyTorch >= 2.3. \n"
|
||||
f"Please use `{new_api}` instead; see the PyTorch 2.0 release notes \n"
|
||||
f"and/or the `torch.func` migration guide for more details \n"
|
||||
f"https://pytorch.org/docs/main/func.migrating.html"
|
||||
)
|
||||
if replace_newlines:
|
||||
warning = warning.replace("\n", "")
|
||||
return warning
|
||||
|
||||
|
||||
def warn_deprecated(api: str, new_api: str | None = None) -> None:
|
||||
warning = get_warning(api, new_api, replace_newlines=True)
|
||||
warnings.warn(warning, FutureWarning, stacklevel=3)
|
||||
|
||||
|
||||
def setup_docs(
|
||||
functorch_api: Callable[..., Any],
|
||||
torch_func_api: Callable[..., Any] | None = None,
|
||||
new_api_name: str | None = None,
|
||||
) -> None:
|
||||
api_name = functorch_api.__name__
|
||||
if torch_func_api is None:
|
||||
torch_func_api = getattr(_impl, api_name)
|
||||
# See https://docs.python.org/3/using/cmdline.html#cmdoption-OO
|
||||
if torch_func_api.__doc__ is None:
|
||||
return
|
||||
|
||||
warning = get_warning(api_name, new_api_name)
|
||||
warning_note = "\n.. warning::\n\n" + textwrap.indent(warning, " ")
|
||||
warning_note = textwrap.indent(warning_note, " ")
|
||||
functorch_api.__doc__ = torch_func_api.__doc__ + warning_note
|
||||
|
||||
|
||||
def vmap(
|
||||
func: Callable[..., Any],
|
||||
in_dims: in_dims_t = 0,
|
||||
out_dims: out_dims_t = 0,
|
||||
randomness: str = "error",
|
||||
*,
|
||||
chunk_size: int | None = None,
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("vmap", "torch.vmap")
|
||||
return apis.vmap(func, in_dims, out_dims, randomness, chunk_size=chunk_size)
|
||||
|
||||
|
||||
def grad(
|
||||
func: Callable[..., Any], argnums: argnums_t = 0, has_aux: bool = False
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("grad")
|
||||
return apis.grad(func, argnums, has_aux)
|
||||
|
||||
|
||||
def grad_and_value(
|
||||
func: Callable[..., Any], argnums: argnums_t = 0, has_aux: bool = False
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("grad_and_value")
|
||||
return apis.grad_and_value(func, argnums, has_aux)
|
||||
|
||||
|
||||
def vjp(func: Callable[..., Any], *primals: Any, has_aux: bool = False) -> Any:
|
||||
warn_deprecated("vjp")
|
||||
return _impl.vjp(func, *primals, has_aux=has_aux)
|
||||
|
||||
|
||||
def jvp(
|
||||
func: Callable[..., Any],
|
||||
primals: Any,
|
||||
tangents: Any,
|
||||
*,
|
||||
strict: bool = False,
|
||||
has_aux: bool = False,
|
||||
) -> Any:
|
||||
warn_deprecated("jvp")
|
||||
return _impl.jvp(func, primals, tangents, strict=strict, has_aux=has_aux)
|
||||
|
||||
|
||||
def jacrev(
|
||||
func: Callable[..., Any],
|
||||
argnums: int | tuple[int, ...] = 0,
|
||||
*,
|
||||
has_aux: bool = False,
|
||||
chunk_size: int | None = None,
|
||||
_preallocate_and_copy: bool = False,
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("jacrev")
|
||||
return _impl.jacrev(
|
||||
func,
|
||||
argnums,
|
||||
has_aux=has_aux,
|
||||
chunk_size=chunk_size,
|
||||
_preallocate_and_copy=_preallocate_and_copy,
|
||||
)
|
||||
|
||||
|
||||
def jacfwd(
|
||||
func: Callable[..., Any],
|
||||
argnums: argnums_t = 0,
|
||||
has_aux: bool = False,
|
||||
*,
|
||||
randomness: str = "error",
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("jacfwd")
|
||||
return _impl.jacfwd(func, argnums, has_aux, randomness=randomness)
|
||||
|
||||
|
||||
def hessian(func: Callable[..., Any], argnums: int = 0) -> Callable[..., Any]:
|
||||
warn_deprecated("hessian")
|
||||
return _impl.hessian(func, argnums=argnums)
|
||||
|
||||
|
||||
def functionalize(
|
||||
func: Callable[..., Any], *, remove: str = "mutations"
|
||||
) -> Callable[..., Any]:
|
||||
warn_deprecated("functionalize")
|
||||
return _impl.functionalize(func, remove=remove)
|
||||
|
||||
|
||||
def make_functional(model: nn.Module, disable_autograd_tracking: bool = False) -> Any:
|
||||
warn_deprecated("make_functional", "torch.func.functional_call")
|
||||
return _nn_impl.make_functional(model, disable_autograd_tracking)
|
||||
|
||||
|
||||
def make_functional_with_buffers(
|
||||
model: nn.Module, disable_autograd_tracking: bool = False
|
||||
) -> Any:
|
||||
warn_deprecated("make_functional_with_buffers", "torch.func.functional_call")
|
||||
return _nn_impl.make_functional_with_buffers(model, disable_autograd_tracking)
|
||||
|
||||
|
||||
def combine_state_for_ensemble(models: list[nn.Module]) -> Any:
|
||||
warn_deprecated("combine_state_for_ensemble", "torch.func.stack_module_state")
|
||||
return _nn_impl.combine_state_for_ensemble(models)
|
||||
|
||||
|
||||
setup_docs(vmap, apis.vmap, "torch.vmap")
|
||||
setup_docs(grad, apis.grad)
|
||||
setup_docs(grad_and_value, apis.grad_and_value)
|
||||
setup_docs(vjp)
|
||||
setup_docs(jvp)
|
||||
setup_docs(jacrev)
|
||||
setup_docs(jacfwd)
|
||||
setup_docs(hessian)
|
||||
setup_docs(functionalize)
|
||||
setup_docs(make_functional, _nn_impl.make_functional, "torch.func.functional_call")
|
||||
setup_docs(
|
||||
make_functional_with_buffers, _nn_impl.make_functional, "torch.func.functional_call"
|
||||
)
|
||||
setup_docs(
|
||||
combine_state_for_ensemble,
|
||||
_nn_impl.combine_state_for_ensemble,
|
||||
"torch.func.stack_module_state",
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,264 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torch._functorch.utils import exposed_in
|
||||
|
||||
|
||||
@exposed_in("torch.func")
|
||||
def functional_call(
|
||||
module: torch.nn.Module,
|
||||
parameter_and_buffer_dicts: dict[str, Tensor] | Sequence[dict[str, Tensor]],
|
||||
args: Any = None,
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
*,
|
||||
tie_weights: bool = True,
|
||||
strict: bool = False,
|
||||
) -> Any:
|
||||
r"""Performs a functional call on the module by replacing the module parameters
|
||||
and buffers with the provided ones.
|
||||
|
||||
.. note:: If the module has active parametrizations, passing a value in the
|
||||
:attr:`parameter_and_buffer_dicts` argument with the name set to the regular parameter
|
||||
name will completely disable the parametrization.
|
||||
If you want to apply the parametrization function to the value passed
|
||||
please set the key as ``{submodule_name}.parametrizations.{parameter_name}.original``.
|
||||
|
||||
.. note:: If the module performs in-place operations on parameters/buffers, these will be reflected
|
||||
in the ``parameter_and_buffer_dicts`` input.
|
||||
|
||||
|
||||
Example::
|
||||
|
||||
>>> a = {'foo': torch.zeros(())}
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> mod = Foo() # does self.foo = self.foo + 1
|
||||
>>> print(mod.foo) # tensor(0.)
|
||||
>>> functional_call(mod, a, torch.ones(()))
|
||||
>>> print(mod.foo) # tensor(0.)
|
||||
>>> print(a['foo']) # tensor(1.)
|
||||
|
||||
.. note:: If the module has tied weights, whether or not functional_call respects the tying is determined by the
|
||||
tie_weights flag.
|
||||
|
||||
Example::
|
||||
|
||||
>>> a = {'foo': torch.zeros(())}
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> mod = Foo() # has both self.foo and self.foo_tied which are tied. Returns x + self.foo + self.foo_tied
|
||||
>>> print(mod.foo) # tensor(1.)
|
||||
>>> mod(torch.zeros(())) # tensor(2.)
|
||||
>>> functional_call(mod, a, torch.zeros(())) # tensor(0.) since it will change self.foo_tied too
|
||||
>>> functional_call(mod, a, torch.zeros(()), tie_weights=False) # tensor(1.)--self.foo_tied is not updated
|
||||
>>> new_a = {'foo': torch.zeros(()), 'foo_tied': torch.zeros(())}
|
||||
>>> functional_call(mod, new_a, torch.zeros()) # tensor(0.)
|
||||
|
||||
An example of passing multiple dictionaries
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
a = (
|
||||
{"weight": torch.ones(1, 1)},
|
||||
{"buffer": torch.zeros(1)},
|
||||
) # two separate dictionaries
|
||||
mod = nn.Bar(1, 1) # return self.weight @ x + self.buffer
|
||||
print(mod.weight) # tensor(...)
|
||||
print(mod.buffer) # tensor(...)
|
||||
x = torch.randn((1, 1))
|
||||
print(x)
|
||||
functional_call(mod, a, x) # same as x
|
||||
print(mod.weight) # same as before functional_call
|
||||
|
||||
|
||||
And here is an example of applying the grad transform over the parameters
|
||||
of a model.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.func import functional_call, grad
|
||||
|
||||
x = torch.randn(4, 3)
|
||||
t = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
|
||||
|
||||
def compute_loss(params, x, t):
|
||||
y = functional_call(model, params, x)
|
||||
return nn.functional.mse_loss(y, t)
|
||||
|
||||
|
||||
grad_weights = grad(compute_loss)(dict(model.named_parameters()), x, t)
|
||||
|
||||
.. note:: If the user does not need grad tracking outside of grad transforms, they can detach all of the
|
||||
parameters for better performance and memory usage
|
||||
|
||||
Example::
|
||||
|
||||
>>> detached_params = {k: v.detach() for k, v in model.named_parameters()}
|
||||
>>> grad_weights = grad(compute_loss)(detached_params, x, t)
|
||||
>>> grad_weights.grad_fn # None--it's not tracking gradients outside of grad
|
||||
|
||||
This means that the user cannot call ``grad_weight.backward()``. However, if they don't need autograd tracking
|
||||
outside of the transforms, this will result in less memory usage and faster speeds.
|
||||
|
||||
Args:
|
||||
module (torch.nn.Module): the module to call
|
||||
parameters_and_buffer_dicts (Dict[str, Tensor] or tuple of Dict[str, Tensor]): the parameters that will be used in
|
||||
the module call. If given a tuple of dictionaries, they must have distinct keys so that all dictionaries can
|
||||
be used together
|
||||
args (Any or tuple): arguments to be passed to the module call. If not a tuple, considered a single argument.
|
||||
kwargs (dict): keyword arguments to be passed to the module call
|
||||
tie_weights (bool, optional): If True, then parameters and buffers tied in the original model will be treated as
|
||||
tied in the reparameterized version. Therefore, if True and different values are passed for the tied
|
||||
parameters and buffers, it will error. If False, it will not respect the originally tied parameters and
|
||||
buffers unless the values passed for both weights are the same. Default: True.
|
||||
strict (bool, optional): If True, then the parameters and buffers passed in must match the parameters and
|
||||
buffers in the original module. Therefore, if True and there are any missing or unexpected keys, it will
|
||||
error. Default: False.
|
||||
|
||||
Returns:
|
||||
Any: the result of calling ``module``.
|
||||
"""
|
||||
if isinstance(parameter_and_buffer_dicts, dict):
|
||||
parameters_and_buffers = parameter_and_buffer_dicts
|
||||
elif isinstance(parameter_and_buffer_dicts, Sequence):
|
||||
if not all(isinstance(d, dict) for d in parameter_and_buffer_dicts):
|
||||
raise ValueError(
|
||||
"Expected all elements of parameter_and_buffer_dicts to be dictionaries"
|
||||
)
|
||||
all_keys = [k for d in parameter_and_buffer_dicts for k in d]
|
||||
all_keys_counter: dict[str, int] = {}
|
||||
for k in all_keys:
|
||||
v = all_keys_counter.get(k, 0)
|
||||
all_keys_counter[k] = v + 1
|
||||
repeated_keys = [key for key, n in all_keys_counter.items() if n > 1]
|
||||
if len(repeated_keys) > 0:
|
||||
raise ValueError(
|
||||
f"{repeated_keys} appeared in multiple dictionaries; behavior of functional call is ambiguous"
|
||||
)
|
||||
parameters_and_buffers = {
|
||||
k: v for d in parameter_and_buffer_dicts for k, v in d.items()
|
||||
}
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Expected parameter_and_buffer_dicts to be a dict, or a list/tuple of dicts, "
|
||||
f"but got {type(parameter_and_buffer_dicts)}"
|
||||
)
|
||||
|
||||
return nn.utils.stateless._functional_call(
|
||||
module,
|
||||
parameters_and_buffers,
|
||||
args,
|
||||
kwargs,
|
||||
tie_weights=tie_weights,
|
||||
strict=strict,
|
||||
)
|
||||
|
||||
|
||||
@exposed_in("torch.func")
|
||||
def stack_module_state(
|
||||
models: Sequence[nn.Module] | nn.ModuleList,
|
||||
) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""stack_module_state(models) -> params, buffers
|
||||
|
||||
Prepares a list of torch.nn.Modules for ensembling with :func:`vmap`.
|
||||
|
||||
Given a list of ``M`` ``nn.Modules`` of the same class, returns two dictionaries
|
||||
that stack all of their parameters and buffers together, indexed by name.
|
||||
The stacked parameters are optimizable (i.e. they are new leaf nodes in the
|
||||
autograd history that are unrelated to the original parameters and can be
|
||||
passed directly to an optimizer).
|
||||
|
||||
Here's an example of how to ensemble over a very simple model:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
num_models = 5
|
||||
batch_size = 64
|
||||
in_features, out_features = 3, 3
|
||||
models = [torch.nn.Linear(in_features, out_features) for i in range(num_models)]
|
||||
data = torch.randn(batch_size, 3)
|
||||
|
||||
|
||||
def wrapper(params, buffers, data):
|
||||
return torch.func.functional_call(models[0], (params, buffers), data)
|
||||
|
||||
|
||||
params, buffers = stack_module_state(models)
|
||||
output = vmap(wrapper, (0, 0, None))(params, buffers, data)
|
||||
|
||||
assert output.shape == (num_models, batch_size, out_features)
|
||||
|
||||
When there's submodules, this follows state dict naming conventions
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Foo(nn.Module):
|
||||
def __init__(self, in_features, out_features):
|
||||
super().__init__()
|
||||
hidden = 4
|
||||
self.l1 = nn.Linear(in_features, hidden)
|
||||
self.l2 = nn.Linear(hidden, out_features)
|
||||
|
||||
def forward(self, x):
|
||||
return self.l2(self.l1(x))
|
||||
|
||||
|
||||
num_models = 5
|
||||
in_features, out_features = 3, 3
|
||||
models = [Foo(in_features, out_features) for i in range(num_models)]
|
||||
params, buffers = stack_module_state(models)
|
||||
print(list(params.keys())) # "l1.weight", "l1.bias", "l2.weight", "l2.bias"
|
||||
|
||||
.. warning::
|
||||
All of the modules being stacked together must be the same (except for
|
||||
the values of their parameters/buffers). For example, they should be in the
|
||||
same mode (training vs eval).
|
||||
"""
|
||||
if len(models) == 0:
|
||||
raise RuntimeError("stack_module_state: Expected at least one model, got 0.")
|
||||
if not (all(m.training for m in models) or all(not m.training for m in models)):
|
||||
raise RuntimeError(
|
||||
"stack_module_state: Expected all models to have the same training/eval mode."
|
||||
)
|
||||
model0_typ = type(models[0])
|
||||
if not all(type(m) is model0_typ for m in models):
|
||||
raise RuntimeError(
|
||||
"stack_module_state: Expected all models to be of the same class."
|
||||
)
|
||||
all_params = [dict(model.named_parameters()) for model in models]
|
||||
params = {
|
||||
k: construct_stacked_leaf(tuple(params[k] for params in all_params), k)
|
||||
for k in all_params[0]
|
||||
}
|
||||
all_buffers = [dict(model.named_buffers()) for model in models]
|
||||
buffers = {
|
||||
k: construct_stacked_leaf(tuple(buffers[k] for buffers in all_buffers), k)
|
||||
for k in all_buffers[0]
|
||||
}
|
||||
|
||||
return params, buffers
|
||||
|
||||
|
||||
def construct_stacked_leaf(
|
||||
tensors: tuple[Tensor, ...] | list[Tensor], name: str
|
||||
) -> Tensor:
|
||||
all_requires_grad = all(t.requires_grad for t in tensors)
|
||||
none_requires_grad = all(not t.requires_grad for t in tensors)
|
||||
if not all_requires_grad and not none_requires_grad:
|
||||
raise RuntimeError(
|
||||
f"Expected {name} from each model to have the same .requires_grad"
|
||||
)
|
||||
result = torch.stack(tensors)
|
||||
if all_requires_grad:
|
||||
result = result.detach().requires_grad_()
|
||||
return result
|
||||
@@ -0,0 +1,562 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from functools import partial, wraps
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch.hub import tqdm
|
||||
from torch.multiprocessing.reductions import StorageWeakRef
|
||||
from torch.utils._content_store import ContentStoreWriter
|
||||
|
||||
from .compile_utils import get_outputs, get_placeholders
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
|
||||
is_tuple = object()
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoadTensorMeta:
|
||||
size: tuple[int, ...]
|
||||
stride: tuple[int, ...]
|
||||
dtype: torch.dtype
|
||||
device: torch.device
|
||||
|
||||
|
||||
class ConcreteProp(torch.fx.Interpreter):
|
||||
def __init__(
|
||||
self,
|
||||
mod: fx.GraphModule,
|
||||
*,
|
||||
writer: ContentStoreWriter | None = None,
|
||||
skip_offload: bool = False,
|
||||
) -> None:
|
||||
super().__init__(mod)
|
||||
self.writer = writer
|
||||
self.skip_offload = skip_offload
|
||||
self.seen_storages: set[StorageWeakRef] = set()
|
||||
self.pbar: Any = None
|
||||
|
||||
def run_node(self, n: fx.Node) -> Any:
|
||||
self.pbar.update(1)
|
||||
r = super().run_node(n)
|
||||
name = n.name
|
||||
|
||||
if isinstance(r, torch.Tensor):
|
||||
if self.writer is None:
|
||||
n.meta["concrete_value"] = r
|
||||
else:
|
||||
if StorageWeakRef(r.untyped_storage()) in self.seen_storages:
|
||||
# Refuse to offload tensors which alias other live
|
||||
# tensors, because this will violate operator contracts
|
||||
n.meta["concrete_value"] = None
|
||||
else:
|
||||
if not self.skip_offload:
|
||||
self.writer.write_tensor(os.path.join("eager", name), r)
|
||||
n.meta["concrete_value"] = LoadTensorMeta(
|
||||
r.size(), r.stride(), r.dtype, r.device
|
||||
)
|
||||
self.seen_storages.add(StorageWeakRef(r.untyped_storage()))
|
||||
else:
|
||||
n.meta["concrete_value"] = is_tuple
|
||||
|
||||
return r
|
||||
|
||||
def propagate(self, *args: Any) -> Any:
|
||||
mod = self.module
|
||||
if not isinstance(mod, fx.GraphModule):
|
||||
raise AssertionError(f"expected fx.GraphModule, got {type(mod)}")
|
||||
with tqdm(
|
||||
desc="Saving intermediates for delta debugging",
|
||||
total=len(mod.graph.nodes),
|
||||
disable=self.writer is None,
|
||||
) as pbar:
|
||||
self.pbar = pbar
|
||||
r = super().run(*args)
|
||||
if not self.skip_offload:
|
||||
pbar.set_description(
|
||||
"Saved! To skip next time, run with --skip-saving-eager-intermediates"
|
||||
)
|
||||
return r
|
||||
|
||||
|
||||
def is_load_tensor_node(node: fx.Node) -> bool:
|
||||
return (
|
||||
node.op == "call_function"
|
||||
and node.target is torch.ops.debugprims.load_tensor.default
|
||||
)
|
||||
|
||||
|
||||
# inplace modifies node/inps
|
||||
def _convert_node_to_placeholder(
|
||||
graph: fx.Graph, node: fx.Node, inps: list[torch.Tensor]
|
||||
) -> bool:
|
||||
if node.op == "output" or node.op == "placeholder":
|
||||
return False
|
||||
|
||||
if is_load_tensor_node(node):
|
||||
return False
|
||||
|
||||
concrete_val = node.meta.get("concrete_value", None)
|
||||
|
||||
if isinstance(concrete_val, torch.Tensor):
|
||||
node.op = "placeholder"
|
||||
node.target = node.name
|
||||
node.args = ()
|
||||
node.kwargs = {}
|
||||
|
||||
inps.append(concrete_val)
|
||||
return True
|
||||
|
||||
elif concrete_val is None:
|
||||
return False
|
||||
|
||||
elif concrete_val is is_tuple:
|
||||
r = False
|
||||
for tuple_user in list(node.users):
|
||||
r = _convert_node_to_placeholder(graph, tuple_user, inps) or r
|
||||
# NB: We must not erase the node at this point, because
|
||||
# we are iterating over the nodes and this would change
|
||||
# the iteration order
|
||||
# graph.erase_node(node)
|
||||
return r
|
||||
|
||||
elif isinstance(concrete_val, LoadTensorMeta):
|
||||
node.op = "call_function"
|
||||
node.target = torch.ops.debugprims.load_tensor.default
|
||||
node.args = (
|
||||
os.path.join("eager", node.name),
|
||||
concrete_val.size,
|
||||
concrete_val.stride,
|
||||
)
|
||||
node.kwargs = {
|
||||
"device": concrete_val.device,
|
||||
"dtype": concrete_val.dtype,
|
||||
}
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def create_minified_hlo_graph(
|
||||
minified_fx_graph: fx.GraphModule, inputs: Sequence[torch.Tensor]
|
||||
) -> None:
|
||||
"""
|
||||
Takes minified FX graph as primary input, and ports it to HLO via StableHLO
|
||||
Provides minified HLO graph as output, and archive them to local directory
|
||||
"""
|
||||
hlo_dir = f"{os.getcwd()}/hlo_files"
|
||||
os.makedirs(hlo_dir, exist_ok=True)
|
||||
|
||||
from torch_xla.stablehlo import save_torch_model_as_stablehlo
|
||||
|
||||
save_torch_model_as_stablehlo(minified_fx_graph, inputs, hlo_dir)
|
||||
|
||||
|
||||
def dump_state(fx_g: fx.GraphModule, inps: Sequence[torch.Tensor]) -> None:
|
||||
print(
|
||||
f"""
|
||||
# Working Repro with {len(fx_g.graph.nodes)} nodes
|
||||
inps = {[(i.shape, i.dtype, i.device.type) for i in inps]}
|
||||
inps = [torch.zeros(())] + [torch.ones(shape, dtype=dtype, device=device) for (shape, dtype, device) in inps]
|
||||
{fx_g.code}
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def is_power_of_two(n: int) -> bool:
|
||||
if n == 0:
|
||||
return False
|
||||
return (n & (n - 1)) == 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReproState:
|
||||
graph: fx.Graph
|
||||
inps: Sequence[torch.Tensor]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
ph_nodes = get_placeholders(self.graph)
|
||||
if len(ph_nodes) != len(self.inps):
|
||||
raise AssertionError(
|
||||
f"len(ph_nodes)={len(ph_nodes)} != len(self.inps)={len(self.inps)}"
|
||||
)
|
||||
|
||||
|
||||
def minifier(
|
||||
fail_f: fx.GraphModule,
|
||||
inps: Sequence[torch.Tensor],
|
||||
module_fails: Callable[[fx.GraphModule, Sequence[torch.Tensor]], bool],
|
||||
dump_state: Callable[[fx.GraphModule, Sequence[torch.Tensor]], None] = dump_state,
|
||||
*,
|
||||
save_dir: str | None = None,
|
||||
offload_to_disk: bool = False,
|
||||
skip_offload: bool = False,
|
||||
skip_sanity: bool = False,
|
||||
max_granularity: int | None = None,
|
||||
) -> tuple[fx.GraphModule, Sequence[torch.Tensor]]:
|
||||
"""
|
||||
Minimizes a FX graph with given inputs, such that the resulting FX graph still returns True for module_fails.
|
||||
|
||||
Does 2 main strategies:
|
||||
1. Truncates suffix: Removes some suffix from the graph and sets a new output.
|
||||
2. Delta Debugging: Tries replacing half of the graph with inputs. If fails,
|
||||
tries replacing quarter of the graph, etc.
|
||||
|
||||
>>> # xdoctest: +SKIP(failing)
|
||||
>>> failing_function = fx.symbolic_trace(f)
|
||||
>>> minimize(failing_function, [torch.randn(5)], lambda fx_g, inps: fx_g(*inps))
|
||||
|
||||
note: module_fails returns True if it fails.
|
||||
"""
|
||||
|
||||
failing_graph = fail_f.graph
|
||||
cur_size = len(failing_graph.nodes)
|
||||
|
||||
if max_granularity is not None and not is_power_of_two(max_granularity):
|
||||
raise RuntimeError(f"max_granularity {max_granularity} not power of two")
|
||||
|
||||
num_queries = 0
|
||||
|
||||
def deepcopy_fx_graph(fx_graph: fx.Graph) -> fx.Graph:
|
||||
return fx.GraphModule(fail_f, copy.deepcopy(fx_graph)).graph
|
||||
|
||||
def graph_fails(graph: fx.Graph, inps: Sequence[torch.Tensor]) -> bool:
|
||||
nonlocal num_queries
|
||||
graph = copy.deepcopy(graph)
|
||||
num_queries += 1
|
||||
mod = fx.GraphModule(fail_f, graph)
|
||||
mod.graph.lint()
|
||||
return module_fails(mod, inps)
|
||||
|
||||
writer = None
|
||||
if offload_to_disk:
|
||||
if save_dir is None:
|
||||
raise AssertionError("save_dir must not be None when offload_to_disk=True")
|
||||
writer = ContentStoreWriter(save_dir)
|
||||
|
||||
ConcreteProp(fail_f, writer=writer, skip_offload=skip_offload).propagate(*inps)
|
||||
if not skip_sanity and not graph_fails(failing_graph, inps):
|
||||
raise RuntimeError("Input graph did not fail the tester")
|
||||
print(f"Started off with {cur_size} nodes", file=sys.stderr)
|
||||
|
||||
def _register_strategy(
|
||||
strategy: Callable[[fx.Graph, Sequence[torch.Tensor], int], ReproState | None],
|
||||
name: str,
|
||||
) -> Callable[[ReproState, int], ReproState | None]:
|
||||
@wraps(strategy)
|
||||
def new_func(old_state: ReproState, granularity: int = 1) -> ReproState | None:
|
||||
print(file=sys.stderr)
|
||||
print(
|
||||
f"Strategy: {name} (G: {granularity}) "
|
||||
f"({len(old_state.graph.nodes)} nodes, {len(old_state.inps)} inputs)",
|
||||
file=sys.stderr,
|
||||
)
|
||||
new_state = strategy(
|
||||
deepcopy_fx_graph(old_state.graph), list(old_state.inps), granularity
|
||||
)
|
||||
if new_state is not None:
|
||||
new_nodes = len(new_state.graph.nodes)
|
||||
old_nodes = len(old_state.graph.nodes)
|
||||
new_inps = len(new_state.inps)
|
||||
old_inps = len(old_state.inps)
|
||||
new_outs = len(get_outputs(new_state.graph))
|
||||
old_outs = len(get_outputs(old_state.graph))
|
||||
progress_made = False
|
||||
if new_nodes < old_nodes:
|
||||
progress_made = True
|
||||
print(
|
||||
f"SUCCESS: Went from {old_nodes} to {new_nodes} nodes",
|
||||
file=sys.stderr,
|
||||
)
|
||||
if new_inps > old_inps:
|
||||
progress_made = True
|
||||
print(
|
||||
f"SUCCESS: Went from {old_inps} to {new_inps} inputs",
|
||||
file=sys.stderr,
|
||||
)
|
||||
if new_outs < old_outs:
|
||||
progress_made = True
|
||||
print(
|
||||
f"SUCCESS: Went from {old_outs} to {new_outs} outputs",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
if not progress_made:
|
||||
raise RuntimeError("Success raised but no progress made?")
|
||||
|
||||
if not graph_fails(new_state.graph, new_state.inps):
|
||||
print(
|
||||
"WARNING: Something went wrong, not applying this minification",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return None
|
||||
return new_state
|
||||
else:
|
||||
print(f"FAIL: {name}", file=sys.stderr)
|
||||
return None
|
||||
|
||||
return new_func
|
||||
|
||||
def register_strategy(
|
||||
name: str,
|
||||
) -> Callable[
|
||||
[Callable[[fx.Graph, Sequence[torch.Tensor], int], ReproState | None]],
|
||||
Callable[[ReproState, int], ReproState | None],
|
||||
]:
|
||||
return partial(_register_strategy, name=name)
|
||||
|
||||
@register_strategy("Truncate suffix")
|
||||
def remove_suffix(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
tested: set[int] = set()
|
||||
new_graph = fx.Graph()
|
||||
env: dict[fx.Node, fx.Node] = {}
|
||||
for idx, node in enumerate(cur_graph.nodes):
|
||||
new_node = new_graph.node_copy(node, lambda x: env[x])
|
||||
if node.op not in ["placeholder", "output"]:
|
||||
# If idx is divisible by (granularity * 2), it would have been checked already.
|
||||
if (
|
||||
idx % granularity == 0
|
||||
and (idx % (granularity * 2) != 0)
|
||||
and idx not in tested
|
||||
):
|
||||
output_node = new_graph.output((new_node,))
|
||||
if len(new_graph.nodes) < len(cur_graph.nodes) and graph_fails(
|
||||
new_graph, cur_inps
|
||||
):
|
||||
return ReproState(new_graph, cur_inps)
|
||||
else:
|
||||
tested.add(idx)
|
||||
new_graph.erase_node(output_node)
|
||||
env[node] = new_node
|
||||
return None
|
||||
|
||||
@register_strategy("Remove outputs")
|
||||
def remove_outputs(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
granularity = max(1, granularity // 2)
|
||||
output: fx.Node | None = None
|
||||
for idx, node in enumerate(cur_graph.nodes):
|
||||
node.idx = idx # type: ignore[attr-defined]
|
||||
if node.op == "output":
|
||||
output = node
|
||||
break
|
||||
|
||||
if output is None:
|
||||
return None
|
||||
|
||||
if isinstance(output.args[0], fx.Node):
|
||||
return None
|
||||
|
||||
# output.args[0] is a tuple/list of nodes when returning multiple outputs
|
||||
output_args_raw = output.args[0]
|
||||
if not isinstance(output_args_raw, (list, tuple)):
|
||||
raise AssertionError(
|
||||
f"expected output_args_raw to be list or tuple, got {type(output_args_raw)}"
|
||||
)
|
||||
output_args = sorted(
|
||||
output_args_raw,
|
||||
key=lambda x: x.idx if isinstance(x, fx.Node) else int(1e9), # type: ignore[attr-defined]
|
||||
)
|
||||
if len(output_args) == 1:
|
||||
return None
|
||||
|
||||
for idx in range(0, len(output_args), granularity):
|
||||
output.args = (output_args[:idx] + output_args[idx + granularity :],)
|
||||
if graph_fails(cur_graph, cur_inps):
|
||||
return ReproState(cur_graph, cur_inps)
|
||||
return None
|
||||
|
||||
def remove_unused_inputs_unchecked(cur_state: ReproState) -> ReproState | None:
|
||||
cur_graph = cur_state.graph
|
||||
cur_inps = cur_state.inps
|
||||
ph_nodes = list(get_placeholders(cur_graph))
|
||||
if len(ph_nodes) != len(cur_inps):
|
||||
raise AssertionError(
|
||||
f"len(ph_nodes)={len(ph_nodes)} != len(cur_inps)={len(cur_inps)}"
|
||||
)
|
||||
|
||||
new_inps: list[torch.Tensor] = []
|
||||
for idx in range(len(ph_nodes)):
|
||||
if len(ph_nodes[idx].users) == 0:
|
||||
cur_graph.erase_node(ph_nodes[idx])
|
||||
else:
|
||||
new_inps.append(cur_inps[idx])
|
||||
if len(new_inps) < len(cur_inps):
|
||||
return ReproState(cur_graph, new_inps)
|
||||
return None
|
||||
|
||||
def remove_unused_inputs_checked(cur_state: ReproState) -> ReproState | None:
|
||||
new_state = remove_unused_inputs_unchecked(cur_state)
|
||||
if new_state is not None and graph_fails(new_state.graph, new_state.inps):
|
||||
return new_state
|
||||
return None
|
||||
|
||||
def _remove_unused_wrapper(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
return remove_unused_inputs_checked(ReproState(cur_graph, cur_inps))
|
||||
|
||||
remove_unused_inputs = register_strategy("Remove unused inputs")(
|
||||
_remove_unused_wrapper
|
||||
)
|
||||
|
||||
@register_strategy("Eliminate dead code")
|
||||
def eliminate_dead_code(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
if cur_graph.eliminate_dead_code() and graph_fails(cur_graph, cur_inps):
|
||||
return ReproState(cur_graph, cur_inps)
|
||||
return None
|
||||
|
||||
def _consolidate_placeholders(
|
||||
cur_graph: fx.Graph, inps: list[torch.Tensor]
|
||||
) -> fx.Graph:
|
||||
new_graph = fx.Graph()
|
||||
env = {}
|
||||
seen_non_placeholder = False
|
||||
|
||||
# Move all placeholders to the front; also, if any load_tensor
|
||||
# is at the front, convert it into an input (because it can be live
|
||||
# all the time)
|
||||
for node in cur_graph.nodes:
|
||||
if node.op == "placeholder":
|
||||
new_node = new_graph.node_copy(node, lambda x: env[x])
|
||||
env[node] = new_node
|
||||
elif not seen_non_placeholder and is_load_tensor_node(node):
|
||||
new_node = new_graph.placeholder(node.name)
|
||||
env[node] = new_node
|
||||
inps.append(
|
||||
torch.ops.debugprims.load_tensor.default(*node.args, **node.kwargs)
|
||||
)
|
||||
else:
|
||||
seen_non_placeholder = True
|
||||
|
||||
# Move everyone else
|
||||
for node in cur_graph.nodes:
|
||||
if node not in env:
|
||||
new_node = new_graph.node_copy(node, lambda x: env[x])
|
||||
env[node] = new_node
|
||||
return new_graph
|
||||
|
||||
@register_strategy("Delta Debugging")
|
||||
def delta_debugging(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
num_nodes = len(cur_graph.nodes)
|
||||
for start_range in range(0, num_nodes, granularity):
|
||||
is_removing = False
|
||||
new_graph = deepcopy_fx_graph(cur_graph)
|
||||
new_inps = list(cur_inps[:])
|
||||
end_range = min(num_nodes, start_range + granularity)
|
||||
for idx in range(start_range, end_range):
|
||||
new_node = list(new_graph.nodes)[idx]
|
||||
if _convert_node_to_placeholder(new_graph, new_node, new_inps):
|
||||
is_removing = True
|
||||
if not is_removing:
|
||||
continue
|
||||
new_graph.eliminate_dead_code()
|
||||
new_graph = _consolidate_placeholders(new_graph, new_inps)
|
||||
new_state = remove_unused_inputs_unchecked(ReproState(new_graph, new_inps))
|
||||
if new_state is None:
|
||||
new_state = ReproState(new_graph, new_inps)
|
||||
if graph_fails(new_state.graph, new_state.inps):
|
||||
return ReproState(new_state.graph, new_state.inps)
|
||||
|
||||
return None
|
||||
|
||||
@register_strategy("Consolidate Inputs")
|
||||
def consolidate_inputs(
|
||||
cur_graph: fx.Graph, cur_inps: Sequence[torch.Tensor], granularity: int
|
||||
) -> ReproState | None:
|
||||
old_len = len(cur_inps)
|
||||
new_inps = list(cur_inps[:])
|
||||
cur_graph = _consolidate_placeholders(cur_graph, new_inps)
|
||||
if len(cur_inps) > old_len and graph_fails(cur_graph, new_inps):
|
||||
return ReproState(cur_graph, new_inps)
|
||||
return None
|
||||
|
||||
failing_state = ReproState(failing_graph, inps)
|
||||
|
||||
def try_granularity(
|
||||
failing_state: ReproState, granularity: int, use_non_granular: bool
|
||||
) -> ReproState | None:
|
||||
print(f"Trying granularity {granularity}", file=sys.stderr)
|
||||
|
||||
strategies = []
|
||||
num_nodes = len(failing_state.graph.nodes)
|
||||
num_outputs = len(get_outputs(failing_state.graph))
|
||||
if num_outputs > num_nodes // 2:
|
||||
strategies += [remove_outputs]
|
||||
|
||||
if use_non_granular:
|
||||
strategies += [
|
||||
eliminate_dead_code,
|
||||
remove_unused_inputs,
|
||||
consolidate_inputs,
|
||||
]
|
||||
|
||||
strategies += [remove_suffix, delta_debugging]
|
||||
|
||||
for strategy in strategies:
|
||||
new_state = strategy(failing_state, granularity)
|
||||
if new_state is not None:
|
||||
return new_state
|
||||
return None
|
||||
|
||||
while True:
|
||||
dump_state(fx.GraphModule(fail_f, failing_state.graph), failing_state.inps)
|
||||
granularity = int(2 ** (math.floor(math.log2(len(failing_state.graph.nodes)))))
|
||||
if max_granularity is not None:
|
||||
granularity = min(max_granularity, granularity)
|
||||
new_state = try_granularity(failing_state, granularity, use_non_granular=True)
|
||||
if new_state is not None:
|
||||
failing_state = new_state
|
||||
continue
|
||||
|
||||
granularity //= 2
|
||||
has_progress = False
|
||||
while granularity >= 1:
|
||||
new_state = try_granularity(
|
||||
failing_state, granularity, use_non_granular=False
|
||||
)
|
||||
if new_state is not None:
|
||||
failing_state = new_state
|
||||
has_progress = True
|
||||
break
|
||||
granularity //= 2
|
||||
if has_progress:
|
||||
continue
|
||||
|
||||
new_state = remove_outputs(failing_state, 1)
|
||||
if new_state is not None:
|
||||
failing_state = new_state
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
if not graph_fails(failing_state.graph, failing_state.inps):
|
||||
raise RuntimeError("Uh oh, something went wrong :( Final graph is not failing")
|
||||
|
||||
print(f"Made {num_queries} queries", file=sys.stderr)
|
||||
failing_fx = fx.GraphModule(fail_f, failing_state.graph)
|
||||
|
||||
# If XLA debugging environment is enabled, create minified HLO graph as well
|
||||
if "XLA_HLO_DEBUG" in os.environ:
|
||||
create_minified_hlo_graph(failing_fx, failing_state.inps)
|
||||
|
||||
dump_state(failing_fx, failing_state.inps)
|
||||
print("Wrote minimal repro out to repro.py", file=sys.stderr)
|
||||
return failing_fx, failing_state.inps
|
||||
@@ -0,0 +1,669 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from typing import Any, NoReturn, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torch.nn.utils._named_member_accessor import NamedMemberAccessor
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterable, Sequence
|
||||
|
||||
# Utilities to make nn.Module "functional"
|
||||
# In particular the goal is to be able to provide a function that takes as input
|
||||
# the parameters and evaluate the nn.Module using fixed inputs.
|
||||
|
||||
|
||||
def raise_parameter_tying_error() -> NoReturn:
|
||||
raise RuntimeError(
|
||||
"make_functional(module): we don't yet support models that "
|
||||
"do parameter tying (also sometimes known as weight sharing). "
|
||||
"Please try to rewrite your model by replacing all instances of the "
|
||||
"tied parameter with another and/or comment your support in "
|
||||
"https://github.com/pytorch/functorch/issues/446"
|
||||
)
|
||||
|
||||
|
||||
def create_names_map(
|
||||
named_params: dict[str, Tensor] | Iterable[tuple[str, Tensor]],
|
||||
tied_named_params: dict[str, Tensor] | Iterable[tuple[str, Tensor]],
|
||||
) -> dict[str, list[str]]:
|
||||
"""
|
||||
named_params is a dictionary of tensors: {'A': A, 'B': B}
|
||||
tied_named_params is another dictionary of tensors {'A': A, 'B': B, 'B_tied': B}
|
||||
with potentially tied (or 'duplicated') tensors
|
||||
|
||||
This function creates a mapping from the names in named_params to the
|
||||
names in tied_named_params: {'A': ['A'], 'B': ['B', 'B_tied']}.
|
||||
"""
|
||||
named_params_dict = dict(named_params)
|
||||
tied_named_params_dict = dict(tied_named_params)
|
||||
|
||||
tensors_dict_keys = set(named_params_dict.keys())
|
||||
tied_tensors_dict_keys = set(tied_named_params_dict.keys())
|
||||
if not tensors_dict_keys.issubset(tied_tensors_dict_keys):
|
||||
raise AssertionError(
|
||||
f"tensors_dict_keys {tensors_dict_keys} is not a subset of "
|
||||
f"tied_tensors_dict_keys {tied_tensors_dict_keys}"
|
||||
)
|
||||
|
||||
tensor_to_mapping: dict[Tensor, tuple[str, list[str]]] = {}
|
||||
for key, tensor in named_params_dict.items():
|
||||
tensor_to_mapping[tensor] = (key, [])
|
||||
for key, tensor in tied_named_params_dict.items():
|
||||
if tensor not in tensor_to_mapping:
|
||||
raise AssertionError(
|
||||
f"tensor for key '{key}' not found in tensor_to_mapping"
|
||||
)
|
||||
tensor_to_mapping[tensor][1].append(key)
|
||||
return dict(tensor_to_mapping.values())
|
||||
|
||||
|
||||
def _extract_members(
|
||||
mod: nn.Module,
|
||||
named_members: Callable[..., Iterable[tuple[str, Tensor]]],
|
||||
subclass: Callable[[Tensor], Tensor],
|
||||
) -> tuple[tuple[Tensor, ...], tuple[str, ...], dict[str, list[str]]]:
|
||||
all_named_members = tuple(named_members(remove_duplicate=False))
|
||||
unique_named_members = tuple(named_members(remove_duplicate=True))
|
||||
names_map = create_names_map(unique_named_members, all_named_members)
|
||||
|
||||
# Remove all the members in the model
|
||||
# pyrefly: ignore [implicit-any]
|
||||
memo = {}
|
||||
accessor = NamedMemberAccessor(mod)
|
||||
for name, p in all_named_members:
|
||||
if p not in memo:
|
||||
memo[p] = subclass(torch.empty_like(p, device="meta"))
|
||||
replacement = memo[p]
|
||||
accessor.set_tensor(name, replacement)
|
||||
|
||||
if len(unique_named_members) == 0:
|
||||
names, params = (), ()
|
||||
else:
|
||||
names, params = zip(*unique_named_members) # type: ignore[assignment]
|
||||
return params, names, names_map
|
||||
|
||||
|
||||
def extract_weights(
|
||||
mod: nn.Module,
|
||||
) -> tuple[tuple[Tensor, ...], tuple[str, ...], dict[str, list[str]]]:
|
||||
"""
|
||||
This function removes all the Parameters from the model and
|
||||
return them as a tuple as well as their original attribute names.
|
||||
The weights must be re-loaded with `load_weights` before the model
|
||||
can be used again.
|
||||
Note that this function modifies the model in place and after this
|
||||
call, mod.parameters() will be empty.
|
||||
"""
|
||||
return _extract_members(mod, mod.named_parameters, nn.Parameter)
|
||||
|
||||
|
||||
def extract_buffers(
|
||||
mod: nn.Module,
|
||||
) -> tuple[tuple[Tensor, ...], tuple[str, ...], dict[str, list[str]]]:
|
||||
return _extract_members(mod, mod.named_buffers, lambda x: x)
|
||||
|
||||
|
||||
def load_weights(
|
||||
mod: nn.Module,
|
||||
names: Sequence[str],
|
||||
params: Sequence[Tensor],
|
||||
as_params: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Reload a set of weights so that `mod` can be used again to perform a forward pass.
|
||||
Note that the `params` are regular Tensors (that can have history) and so are left
|
||||
as Tensors. This means that mod.parameters() will still be empty after this call.
|
||||
"""
|
||||
accessor = NamedMemberAccessor(mod)
|
||||
if as_params:
|
||||
params = [nn.Parameter(p) for p in params]
|
||||
accessor.set_tensors(names, params)
|
||||
|
||||
|
||||
def _swap_state(
|
||||
mod: nn.Module, names_map: dict[str, list[str]], elems: Iterable[Tensor]
|
||||
) -> list[Tensor]:
|
||||
result: list[Tensor] = []
|
||||
accessor = NamedMemberAccessor(mod)
|
||||
for (_, attr_names), elem in zip(names_map.items(), elems):
|
||||
for i, attr_name in enumerate(attr_names):
|
||||
if i == 0:
|
||||
result.append(accessor.swap_tensor(attr_name, elem))
|
||||
else:
|
||||
accessor.set_tensor(attr_name, elem)
|
||||
return result
|
||||
|
||||
|
||||
def load_buffers(
|
||||
mod: nn.Module,
|
||||
names: Sequence[str],
|
||||
buffers: Sequence[Tensor],
|
||||
as_params: bool = False,
|
||||
) -> None:
|
||||
accessor = NamedMemberAccessor(mod)
|
||||
accessor.set_tensors(names, buffers)
|
||||
|
||||
|
||||
def load_state(
|
||||
model: nn.Module,
|
||||
weights: Sequence[Tensor],
|
||||
weight_names: Sequence[str],
|
||||
buffers: Sequence[Tensor] = (),
|
||||
buffer_names: Sequence[str] = (),
|
||||
) -> nn.Module:
|
||||
"""load_state(model, weights, weight_names, buffers=(), buffer_names=()) -> model
|
||||
|
||||
load_state takes `weights` and `buffers` and assigns them to the model.
|
||||
This is the inverse operation of `make_functional_deprecated_v1`.
|
||||
"""
|
||||
if len(weight_names) != len(weights):
|
||||
raise AssertionError(
|
||||
f"len(weight_names)={len(weight_names)} != len(weights)={len(weights)}"
|
||||
)
|
||||
load_weights(model, weight_names, weights)
|
||||
if len(buffers) > 0:
|
||||
if len(buffer_names) != len(buffers):
|
||||
raise AssertionError(
|
||||
f"len(buffer_names)={len(buffer_names)} != len(buffers)={len(buffers)}"
|
||||
)
|
||||
load_buffers(model, buffer_names, buffers)
|
||||
return model
|
||||
|
||||
|
||||
def make_functional_deprecated_v1(
|
||||
model: nn.Module,
|
||||
) -> tuple[tuple[Tensor, ...], Callable[..., Any], tuple[str, ...]]:
|
||||
"""make_functional_deprecated_v1(model) -> weights, func, weight_names
|
||||
|
||||
Given an nn.Module, make_functional_deprecated_v1 extracts the state (weights)
|
||||
and returns a functional version of the model, `func`. This makes
|
||||
it so that it is possible use transforms over the parameters of
|
||||
`model`.
|
||||
|
||||
`func` can be invoked as follows:
|
||||
```
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
weights, func, _ = make_functional_deprecated_v1(model)
|
||||
func(weights, (x,))
|
||||
```
|
||||
|
||||
And here is an example of applying the grad transform:
|
||||
```
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
weights, _, func = make_functional_deprecated_v1(model)
|
||||
grad_weights = grad(func)(weights, (x,))
|
||||
```
|
||||
|
||||
To put the state back into a model, use `load_state`.
|
||||
"""
|
||||
buffers = list(model.buffers())
|
||||
if len(buffers) > 0:
|
||||
raise RuntimeError(
|
||||
"make_functional_deprecated_v1(model): `model` has buffers. Please use "
|
||||
"make_functional_with_buffers_deprecated_v1(model) instead."
|
||||
)
|
||||
weights, descriptors, _ = extract_weights(model)
|
||||
|
||||
def fun(weights: tuple[Tensor, ...], data: tuple[Any, ...]) -> Any:
|
||||
mutable_model = copy.deepcopy(model)
|
||||
load_weights(mutable_model, descriptors, weights)
|
||||
return mutable_model(*data)
|
||||
|
||||
return weights, fun, descriptors
|
||||
|
||||
|
||||
def make_functional_with_buffers_deprecated_v1(
|
||||
model: nn.Module,
|
||||
) -> tuple[
|
||||
tuple[Tensor, ...],
|
||||
tuple[Tensor, ...],
|
||||
Callable[..., Any],
|
||||
tuple[str, ...],
|
||||
tuple[str, ...],
|
||||
]:
|
||||
"""make_functional_with_buffers_deprecated_v1(model) -> weights, buffers, func, weight_names, buffer_names
|
||||
|
||||
Given an nn.Module, make_functional_with_buffers_deprecated_v1 extracts the state (weights and buffers)
|
||||
and returns a functional version of the model, `func`.
|
||||
|
||||
`func` can be invoked as follows:
|
||||
```
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
weights, buffers, func, _, _ = make_functional_with_buffers_deprecated_v1(model)
|
||||
func(weights, buffers, (x,))
|
||||
```
|
||||
|
||||
And here is an example of applying the grad transform:
|
||||
```
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
weights, buffers, func, _, _ = make_functional_with_buffers_deprecated_v1(model)
|
||||
func(weights, buffers, (x,))
|
||||
grad_weights = grad(func)(weights, buffers, (x,))
|
||||
```
|
||||
|
||||
To put the state back into a model, use `load_state`.
|
||||
"""
|
||||
weights, weight_descriptors, _ = extract_weights(model)
|
||||
buffers, buf_descriptors, _ = extract_buffers(model)
|
||||
|
||||
def fun(
|
||||
weights: tuple[Tensor, ...],
|
||||
buffers: tuple[Tensor, ...],
|
||||
data: tuple[Any, ...],
|
||||
) -> Any:
|
||||
mutable_model = copy.deepcopy(model)
|
||||
load_weights(mutable_model, weight_descriptors, weights)
|
||||
load_buffers(mutable_model, buf_descriptors, buffers)
|
||||
return mutable_model(*data)
|
||||
|
||||
return weights, buffers, fun, weight_descriptors, buf_descriptors
|
||||
|
||||
|
||||
class FunctionalModuleWithBuffers(nn.Module):
|
||||
"""
|
||||
This is the callable object returned by :func:`make_functional_with_buffers`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stateless_model: nn.Module,
|
||||
param_names: tuple[str, ...],
|
||||
buffer_names: tuple[str, ...],
|
||||
param_names_map: dict[str, list[str]],
|
||||
buffer_names_map: dict[str, list[str]],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.stateless_model = stateless_model
|
||||
self.param_names = param_names
|
||||
self.buffer_names = buffer_names
|
||||
|
||||
self.all_names_map = dict(param_names_map)
|
||||
self.all_names_map.update(buffer_names_map)
|
||||
|
||||
@staticmethod
|
||||
def _create_from(
|
||||
model: nn.Module, disable_autograd_tracking: bool = False
|
||||
) -> tuple[FunctionalModuleWithBuffers, tuple[Tensor, ...], tuple[Tensor, ...]]:
|
||||
# TODO: We don't need to copy the model to create a stateless copy
|
||||
model_copy = copy.deepcopy(model)
|
||||
params, param_names, param_names_map = extract_weights(model_copy)
|
||||
buffers, buffer_names, buffer_names_map = extract_buffers(model_copy)
|
||||
if disable_autograd_tracking:
|
||||
for param in params:
|
||||
param.requires_grad_(False)
|
||||
return (
|
||||
FunctionalModuleWithBuffers(
|
||||
model_copy, param_names, buffer_names, param_names_map, buffer_names_map
|
||||
),
|
||||
params,
|
||||
buffers,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
params: Iterable[Tensor],
|
||||
buffers: Iterable[Tensor],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
# Temporarily load the state back onto self.stateless_model
|
||||
old_state = _swap_state(
|
||||
self.stateless_model,
|
||||
self.all_names_map,
|
||||
tuple(params) + tuple(buffers),
|
||||
)
|
||||
try:
|
||||
return self.stateless_model(*args, **kwargs)
|
||||
finally:
|
||||
# Remove the loaded state on self.stateless_model
|
||||
_swap_state(self.stateless_model, self.all_names_map, old_state)
|
||||
|
||||
|
||||
class FunctionalModule(nn.Module):
|
||||
"""
|
||||
This is the callable object returned by :func:`make_functional`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stateless_model: nn.Module,
|
||||
param_names: tuple[str, ...],
|
||||
names_map: dict[str, list[str]],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.stateless_model = stateless_model
|
||||
self.param_names = param_names
|
||||
self.names_map = names_map
|
||||
|
||||
@staticmethod
|
||||
def _create_from(
|
||||
model: nn.Module, disable_autograd_tracking: bool = False
|
||||
) -> tuple[FunctionalModule, tuple[Tensor, ...]]:
|
||||
# TODO: We don't need to copy the model to create a stateless copy
|
||||
model_copy = copy.deepcopy(model)
|
||||
params, param_names, names_map = extract_weights(model_copy)
|
||||
if disable_autograd_tracking:
|
||||
for param in params:
|
||||
param.requires_grad_(False)
|
||||
return FunctionalModule(model_copy, param_names, names_map), params
|
||||
|
||||
def forward(self, params: Iterable[Tensor], *args: Any, **kwargs: Any) -> Any:
|
||||
# Temporarily load the state back onto self.stateless_model
|
||||
old_state = _swap_state(self.stateless_model, self.names_map, params)
|
||||
try:
|
||||
return self.stateless_model(*args, **kwargs)
|
||||
finally:
|
||||
# Remove the loaded state on self.stateless_model
|
||||
_swap_state(self.stateless_model, self.names_map, old_state)
|
||||
|
||||
|
||||
def make_functional(
|
||||
model: nn.Module, disable_autograd_tracking: bool = False
|
||||
) -> tuple[FunctionalModule, tuple[Tensor, ...]]:
|
||||
"""make_functional(model, disable_autograd_tracking=False) -> func, params
|
||||
|
||||
Given a ``torch.nn.Module``, :func:`make_functional` extracts the state
|
||||
(params) and returns a functional version of the model, ``func``. This
|
||||
makes it so that it is possible use transforms over the parameters of
|
||||
``model``.
|
||||
|
||||
``func`` can be invoked as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from functorch import make_functional
|
||||
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
func, params = make_functional(model)
|
||||
func(params, x)
|
||||
|
||||
And here is an example of applying the grad transform over the parameters
|
||||
of a model.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from functorch import make_functional, grad
|
||||
|
||||
x = torch.randn(4, 3)
|
||||
t = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
func, params = make_functional(model)
|
||||
|
||||
|
||||
def compute_loss(params, x, t):
|
||||
y = func(params, x)
|
||||
return nn.functional.mse_loss(y, t)
|
||||
|
||||
|
||||
grad_weights = grad(compute_loss)(params, x, t)
|
||||
|
||||
If the model has any buffers, please use :func:`make_functional_with_buffers` instead.
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module): Input model.
|
||||
disable_autograd_tracking (bool): Flag to disable gradients tracking for output parameters.
|
||||
The returned params are unrelated to the set of params from the original model. If False (default),
|
||||
the params will have ``requires_grad=True`` on them (aka they will be trackable with regular
|
||||
PyTorch autograd), matching the requires_grad-ness of the params from the original model.
|
||||
Otherwise, the returned params will have ``requires_grad=False``. Default, False.
|
||||
If you plan on using regular PyTorch autograd (e.g., if you want to call ``.backward()`` or
|
||||
``torch.autograd.grad()``, then set ``disable_autograd_tracking=False``.
|
||||
Otherwise, if you're only planning on using functorch's gradient transforms,
|
||||
then please set ``disable_autograd_tracking=True`` to avoid unnecessarily tracking
|
||||
history with PyTorch autograd.
|
||||
|
||||
"""
|
||||
buffers = list(model.buffers())
|
||||
if len(buffers) > 0:
|
||||
raise RuntimeError(
|
||||
"make_functional(model): `model` has buffers. Please use "
|
||||
"make_functional_with_buffers(model) instead."
|
||||
)
|
||||
return FunctionalModule._create_from(
|
||||
model, disable_autograd_tracking=disable_autograd_tracking
|
||||
)
|
||||
|
||||
|
||||
def make_functional_with_buffers(
|
||||
model: nn.Module, disable_autograd_tracking: bool = False
|
||||
) -> tuple[FunctionalModuleWithBuffers, tuple[Tensor, ...], tuple[Tensor, ...]]:
|
||||
"""make_functional_with_buffers(model, disable_autograd_tracking=False) -> func, params, buffers
|
||||
|
||||
Given a ``torch.nn.Module``, make_functional_with_buffers extracts the
|
||||
state (params and buffers) and returns a functional version of the model
|
||||
``func`` that can be invoked like a function.
|
||||
|
||||
``func`` can be invoked as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from functorch import make_functional_with_buffers
|
||||
|
||||
x = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
func, params, buffers = make_functional_with_buffers(model)
|
||||
func(params, buffers, x)
|
||||
|
||||
And here is an example of applying the grad transform over the parameters
|
||||
of a model:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from functorch import make_functional_with_buffers, grad
|
||||
|
||||
x = torch.randn(4, 3)
|
||||
t = torch.randn(4, 3)
|
||||
model = nn.Linear(3, 3)
|
||||
func, params, buffers = make_functional_with_buffers(model)
|
||||
|
||||
|
||||
def compute_loss(params, buffers, x, t):
|
||||
y = func(params, buffers, x)
|
||||
return nn.functional.mse_loss(y, t)
|
||||
|
||||
|
||||
grad_weights = grad(compute_loss)(params, buffers, x, t)
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module): Input model.
|
||||
disable_autograd_tracking (bool): Flag to disable gradients tracking for output parameters.
|
||||
The returned params are unrelated to the set of params from the original model. If False (default),
|
||||
the params will have ``requires_grad=True`` on them (aka they will be trackable with regular
|
||||
PyTorch autograd), matching the requires_grad-ness of the params from the original model.
|
||||
Otherwise, the returned params will have ``requires_grad=False``. Default, False.
|
||||
If you plan on using regular PyTorch autograd (e.g., if you want to call ``.backward()`` or
|
||||
``torch.autograd.grad()``, then set ``disable_autograd_tracking=False``.
|
||||
Otherwise, if you're only planning on using functorch's gradient transforms,
|
||||
then please set ``disable_autograd_tracking=True`` to avoid unnecessarily tracking
|
||||
history with PyTorch autograd.
|
||||
|
||||
"""
|
||||
return FunctionalModuleWithBuffers._create_from(
|
||||
model, disable_autograd_tracking=disable_autograd_tracking
|
||||
)
|
||||
|
||||
|
||||
def transpose_stack(
|
||||
tuple_of_tuple_of_tensors: tuple[tuple[Tensor, ...], ...],
|
||||
) -> tuple[Tensor, ...]:
|
||||
tuple_of_tuple_of_tensors = tuple(zip(*tuple_of_tuple_of_tensors))
|
||||
results = tuple(
|
||||
torch.stack(shards).detach() for shards in tuple_of_tuple_of_tensors
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def combine_state_for_ensemble(
|
||||
models: Sequence[nn.Module],
|
||||
) -> tuple[FunctionalModuleWithBuffers, tuple[Tensor, ...], tuple[Tensor, ...]]:
|
||||
"""combine_state_for_ensemble(models) -> func, params, buffers
|
||||
|
||||
Prepares a list of torch.nn.Modules for ensembling with :func:`vmap`.
|
||||
|
||||
Given a list of ``M`` ``nn.Modules`` of the same class, stacks all of their
|
||||
parameters and buffers together to make ``params`` and ``buffers``.
|
||||
Each parameter and buffer in the result will have an additional dimension
|
||||
of size ``M``.
|
||||
|
||||
:func:`combine_state_for_ensemble` also returns ``func``, a functional
|
||||
version of one of the models in :attr:`models`. One cannot directly run
|
||||
``func(params, buffers, *args, **kwargs)`` directly, you probably want to
|
||||
use ``vmap(func, ...)(params, buffers, *args, **kwargs)``
|
||||
|
||||
Here's an example of how to ensemble over a very simple model:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
num_models = 5
|
||||
batch_size = 64
|
||||
in_features, out_features = 3, 3
|
||||
models = [torch.nn.Linear(in_features, out_features) for i in range(num_models)]
|
||||
data = torch.randn(batch_size, 3)
|
||||
|
||||
fmodel, params, buffers = combine_state_for_ensemble(models)
|
||||
output = vmap(fmodel, (0, 0, None))(params, buffers, data)
|
||||
|
||||
assert output.shape == (num_models, batch_size, out_features)
|
||||
|
||||
.. warning::
|
||||
All of the modules being stacked together must be the same (except for
|
||||
the values of their parameters/buffers). For example, they should be in the
|
||||
same mode (training vs eval).
|
||||
|
||||
This API is subject to change -- we're investigating better ways to
|
||||
create ensembles and would love your feedback how to improve this.
|
||||
"""
|
||||
if len(models) == 0:
|
||||
raise RuntimeError(
|
||||
"combine_state_for_ensemble: Expected at least one model, got 0."
|
||||
)
|
||||
if not (all(m.training for m in models) or all(not m.training for m in models)):
|
||||
raise RuntimeError(
|
||||
"combine_state_for_ensemble: Expected all models to "
|
||||
"have the same training/eval mode."
|
||||
)
|
||||
model0_typ = type(models[0])
|
||||
if not all(type(m) is model0_typ for m in models):
|
||||
raise RuntimeError(
|
||||
"combine_state_for_ensemble: Expected all models to be of the same class."
|
||||
)
|
||||
funcs, params, buffers = zip(
|
||||
*[make_functional_with_buffers(model) for model in models]
|
||||
)
|
||||
params = transpose_stack(params)
|
||||
buffers = transpose_stack(buffers)
|
||||
return funcs[0], params, buffers
|
||||
|
||||
|
||||
def functional_init(
|
||||
model_class: type[nn.Module],
|
||||
ensemble_shape: tuple[()] | tuple[int, ...] = (),
|
||||
device: torch.types.Device = "cpu",
|
||||
) -> Callable[..., tuple[tuple[Tensor, ...], Callable[..., Any], tuple[str, ...]]]:
|
||||
def wrapped(
|
||||
*args: Any, **kwargs: Any
|
||||
) -> tuple[tuple[Tensor, ...], Callable[..., Any], tuple[str, ...]]:
|
||||
if len(ensemble_shape) >= 2:
|
||||
raise ValueError("NYI: ensemble_shape with more than 1 element")
|
||||
if len(ensemble_shape) == 0:
|
||||
model = model_class(*args, **kwargs).to(device)
|
||||
return make_functional_deprecated_v1(model)
|
||||
num_models = ensemble_shape[0] # type: ignore[misc]
|
||||
if num_models <= 0:
|
||||
raise ValueError(f"num_models {num_models} should be > 0")
|
||||
# NB: Not very efficient, more of a POC
|
||||
models = tuple(
|
||||
model_class(*args, **kwargs).to(device) for _ in range(num_models)
|
||||
)
|
||||
_, fn, names = make_functional_deprecated_v1(model_class(*args, **kwargs))
|
||||
weights = tuple(make_functional_deprecated_v1(model)[0] for model in models)
|
||||
weights = tuple(zip(*weights))
|
||||
weights = tuple(torch.stack(shards).detach() for shards in weights)
|
||||
return weights, fn, names
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def functional_init_with_buffers(
|
||||
model_class: type[nn.Module],
|
||||
ensemble_shape: tuple[()] | tuple[int, ...] = (),
|
||||
device: torch.types.Device = "cpu",
|
||||
) -> Callable[
|
||||
...,
|
||||
tuple[
|
||||
tuple[Tensor, ...],
|
||||
tuple[Tensor, ...],
|
||||
Callable[..., Any],
|
||||
tuple[str, ...],
|
||||
tuple[str, ...],
|
||||
]
|
||||
| tuple[tuple[Tensor, ...], Callable[..., Any], tuple[str, ...]],
|
||||
]:
|
||||
def wrapped(
|
||||
*args: Any, **kwargs: Any
|
||||
) -> (
|
||||
tuple[
|
||||
tuple[Tensor, ...],
|
||||
tuple[Tensor, ...],
|
||||
Callable[..., Any],
|
||||
tuple[str, ...],
|
||||
tuple[str, ...],
|
||||
]
|
||||
| tuple[tuple[Tensor, ...], Callable[..., Any], tuple[str, ...]]
|
||||
):
|
||||
if len(ensemble_shape) >= 2:
|
||||
raise ValueError("NYI: ensemble_shape with more than 1 element")
|
||||
if len(ensemble_shape) == 0:
|
||||
model = model_class(*args, **kwargs).to(device)
|
||||
return make_functional_deprecated_v1(model)
|
||||
num_models = ensemble_shape[0] # type: ignore[misc]
|
||||
if num_models <= 0:
|
||||
raise ValueError(f"num_models {num_models} should be > 0")
|
||||
# NB: Not very efficient, more of a POC
|
||||
models = tuple(
|
||||
model_class(*args, **kwargs).to(device) for _ in range(num_models)
|
||||
)
|
||||
(
|
||||
_,
|
||||
_,
|
||||
fn,
|
||||
weight_names,
|
||||
buffer_names,
|
||||
) = make_functional_with_buffers_deprecated_v1(model_class(*args, **kwargs))
|
||||
weights, buffers = zip(
|
||||
*tuple(
|
||||
make_functional_with_buffers_deprecated_v1(model)[:2]
|
||||
for model in models
|
||||
)
|
||||
)
|
||||
weights = tuple(zip(*weights))
|
||||
weights = tuple(torch.stack(shards).detach() for shards in weights)
|
||||
buffers = tuple(zip(*buffers))
|
||||
buffers = tuple(torch.stack(shards).detach() for shards in buffers)
|
||||
return weights, buffers, fn, weight_names, buffer_names
|
||||
|
||||
return wrapped
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,169 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
|
||||
"""
|
||||
This module contains pre-dispatch wrappers for functorch operations
|
||||
that enable proper tracing in PT2 non-strict export/compile fx graph.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._C._functorch import (
|
||||
_add_batch_dim as _add_batch_dim_impl,
|
||||
_remove_batch_dim as _remove_batch_dim_impl,
|
||||
_vmap_decrement_nesting as _vmap_decrement_nesting_impl,
|
||||
_vmap_increment_nesting as _vmap_increment_nesting_impl,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import threading
|
||||
|
||||
|
||||
def _add_batch_dim(self: torch.Tensor, batch_dim: int, level: int) -> torch.Tensor:
|
||||
"""
|
||||
Thin wrapper around torch._C._add_batch_dim that is used to proxy in
|
||||
PT2 export/compile fx graph
|
||||
"""
|
||||
from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
|
||||
|
||||
mode = _maybe_find_pre_dispatch_tf_mode_for_export()
|
||||
batch_dim = self.ndim + batch_dim if batch_dim < 0 else batch_dim
|
||||
|
||||
if mode:
|
||||
return torch.overrides.handle_torch_function(
|
||||
_add_batch_dim, (self,), self, batch_dim, level
|
||||
)
|
||||
|
||||
res = _add_batch_dim_impl(self, batch_dim, level)
|
||||
return res
|
||||
|
||||
|
||||
def _remove_batch_dim(
|
||||
self: torch.Tensor, level: int, batch_size: int, out_dim: int
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Thin wrapper around torch._C._remove_batch_dim that is used to proxy in
|
||||
PT2 export/compile fx graph
|
||||
"""
|
||||
from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
|
||||
|
||||
mode = _maybe_find_pre_dispatch_tf_mode_for_export()
|
||||
|
||||
if mode:
|
||||
return torch.overrides.handle_torch_function(
|
||||
_remove_batch_dim, (self,), self, level, batch_size, out_dim
|
||||
)
|
||||
|
||||
res = _remove_batch_dim_impl(self, level, batch_size, out_dim)
|
||||
return res
|
||||
|
||||
|
||||
def _vmap_increment_nesting(batch_size: int, randomness: str) -> int:
|
||||
"""
|
||||
Thin wrapper around torch._C._vmap_increment_nesting that is used
|
||||
to proxy in export/compile graph
|
||||
"""
|
||||
from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
|
||||
|
||||
mode = _maybe_find_pre_dispatch_tf_mode_for_export()
|
||||
|
||||
if mode:
|
||||
return torch.overrides.handle_torch_function(
|
||||
_vmap_increment_nesting, (batch_size,), batch_size, randomness
|
||||
)
|
||||
res = _vmap_increment_nesting_impl(batch_size, randomness)
|
||||
return res
|
||||
|
||||
|
||||
def _vmap_decrement_nesting() -> int:
|
||||
"""
|
||||
Thin wrapper around torch._C._vmap_increment_nesting that is used
|
||||
to proxy in export/compile graph
|
||||
"""
|
||||
from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
|
||||
|
||||
mode = _maybe_find_pre_dispatch_tf_mode_for_export()
|
||||
|
||||
if mode:
|
||||
return torch.overrides.handle_torch_function(
|
||||
_vmap_decrement_nesting,
|
||||
(),
|
||||
)
|
||||
return _vmap_decrement_nesting_impl()
|
||||
|
||||
|
||||
# Global variables for lazy_load_decompositions
|
||||
DECOMPOSITIONS_LOADED: bool = False
|
||||
DECOMPOSITIONS_LOCK: threading.Lock | None = None
|
||||
VMAP_DECOMPOSITIONS_LIB: torch.library.Library | None = None
|
||||
|
||||
|
||||
def lazy_load_decompositions() -> None:
|
||||
"""
|
||||
Lazy loading of vmap decompositions with pre-dispatch support.
|
||||
"""
|
||||
from torch._export.utils import _maybe_find_pre_dispatch_tf_mode_for_export
|
||||
|
||||
mode = _maybe_find_pre_dispatch_tf_mode_for_export()
|
||||
|
||||
if mode:
|
||||
return torch.overrides.handle_torch_function(lazy_load_decompositions, ())
|
||||
|
||||
global DECOMPOSITIONS_LOADED, DECOMPOSITIONS_LOCK, VMAP_DECOMPOSITIONS_LIB
|
||||
|
||||
if DECOMPOSITIONS_LOADED:
|
||||
return
|
||||
|
||||
# Initialize lock if needed
|
||||
if DECOMPOSITIONS_LOCK is None:
|
||||
import threading
|
||||
|
||||
DECOMPOSITIONS_LOCK = threading.Lock()
|
||||
|
||||
with DECOMPOSITIONS_LOCK:
|
||||
if DECOMPOSITIONS_LOADED:
|
||||
return
|
||||
|
||||
import os
|
||||
|
||||
if not (os.environ.get("PYTORCH_JIT", "1") == "1" and __debug__):
|
||||
DECOMPOSITIONS_LOADED = True
|
||||
return
|
||||
|
||||
# use an alternate way to register an operator into the decomposition table
|
||||
# _register_jit_decomposition doesn't work for some operators, e.g. addr,
|
||||
# because the Tensor types generated cannot be unioned by torchscript
|
||||
# decomp should be type OpOverload
|
||||
VMAP_DECOMPOSITIONS_LIB = torch.library.Library(
|
||||
"aten", "IMPL", "FuncTorchBatched"
|
||||
)
|
||||
|
||||
from torch._decomp import decomposition_table
|
||||
|
||||
def _register_python_decomposition_vmap(decomp: torch._ops.OpOverload) -> None:
|
||||
if VMAP_DECOMPOSITIONS_LIB is None:
|
||||
raise AssertionError("VMAP_DECOMPOSITIONS_LIB must not be None")
|
||||
if decomp in decomposition_table:
|
||||
VMAP_DECOMPOSITIONS_LIB.impl(decomp, decomposition_table[decomp])
|
||||
else:
|
||||
raise RuntimeError(f"could not find decomposition for {decomp}")
|
||||
|
||||
_register_python_decomposition_vmap(torch.ops.aten.mse_loss_backward.default)
|
||||
_register_python_decomposition_vmap(
|
||||
torch.ops.aten.smooth_l1_loss_backward.default
|
||||
)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.huber_loss_backward.default)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.nll_loss_forward.default)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.nll_loss2d_forward.default)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.nll_loss_backward.default)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.nll_loss2d_backward.default)
|
||||
_register_python_decomposition_vmap(torch.ops.aten.addr.default)
|
||||
|
||||
DECOMPOSITIONS_LOADED = True
|
||||
@@ -0,0 +1,338 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from abc import ABC, abstractmethod
|
||||
from functools import cached_property
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._C._functorch import (
|
||||
CFunctionalizeInterpreterPtr,
|
||||
CGradInterpreterPtr,
|
||||
CInterpreter,
|
||||
CJvpInterpreterPtr,
|
||||
CVmapInterpreterPtr,
|
||||
pop_dynamic_layer_stack,
|
||||
push_dynamic_layer_stack,
|
||||
RandomnessType,
|
||||
TransformType,
|
||||
)
|
||||
from torch.autograd.forward_ad import _set_fwd_grad_enabled
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator
|
||||
|
||||
|
||||
"""
|
||||
This file contains the functorch integration with PyDispatcher.
|
||||
|
||||
PyDispatcher does not understand functorch's DynamicLayerStack dispatching
|
||||
logic because it is entirely implemented in C++ in the fallbacks for two
|
||||
dispatch keys, FuncTorchDynamicLayer{Front, Back}Mode (PyDispatcher is unable
|
||||
to directly reuse C++ boxed fallbacks).
|
||||
|
||||
Instead of trying to hammer PyDispatcher into understanding those fallbacks,
|
||||
we re-implement the logic of peeking the top of the stack for an interpreter,
|
||||
selecting the interpreter to dispatch on, etc, in Python. This leads to a
|
||||
simpler design.
|
||||
|
||||
The main difference between C++ functorch and PyDispatcher's functorch logic
|
||||
is that:
|
||||
- C++ functorch needs to manually tweak dispatch keys to ping-pong between
|
||||
DynamicLayerFrontMode and DynamicLayerBackMode.
|
||||
- PyDispatcher's functorch logic pops an Interpreter from the top of the stack
|
||||
and asks it to execute the rule associated with the Interpreter.
|
||||
|
||||
In C++ we do the ping-pong because e.g. vmap rules are associated with the
|
||||
batched DispatchKey, but in PyDispatcher we are able to avoid this by asking
|
||||
the user to register a batching rule directly to a transform that an
|
||||
interpreter then invokes.
|
||||
"""
|
||||
|
||||
|
||||
# FuncTorchInterpreter is the Python version of Interpreter (recall that
|
||||
# the DynamicLayerStack is a stack of interpreters).
|
||||
# It is a wrapper around the actual C++ Interpreter object.
|
||||
#
|
||||
# Keep the methods in sync with aten/src/ATen/functorch/Interpreter.h
|
||||
class FuncTorchInterpreter(ABC):
|
||||
def __init__(self, cptr: Any) -> None:
|
||||
self._cptr = cptr
|
||||
|
||||
# Process an operation. eg for vmap, this is invoking a batching rule.
|
||||
# Conceptually this is analogous to Interpreter::process in C++
|
||||
@abstractmethod
|
||||
def process(self, op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
pass
|
||||
|
||||
# lower an operation from this Interpreter to the next Interpreter on the stack.
|
||||
# Concretely, this involves temporarily popping the current Interpreter.
|
||||
# Conceptually this is analogous to Interpreter::sendToNextInterpreter in C++
|
||||
def lower(self) -> contextlib.AbstractContextManager[Any]:
|
||||
return temporarily_pop_interpreter_stack()
|
||||
|
||||
def level(self) -> int:
|
||||
return self._cptr.level()
|
||||
|
||||
def key(self) -> TransformType:
|
||||
return self._cptr.key()
|
||||
|
||||
def get_state(self) -> tuple[Any, ...]:
|
||||
raise NotImplementedError
|
||||
|
||||
def check_state(self, state: tuple[Any, ...]) -> bool:
|
||||
return state == self.get_state()
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
state = self.__dict__.copy()
|
||||
state.pop("_cptr", None)
|
||||
return state
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def temporarily_pop_interpreter_stack() -> Generator[None, None, None]:
|
||||
try:
|
||||
saved = pop_dynamic_layer_stack()
|
||||
yield
|
||||
finally:
|
||||
push_dynamic_layer_stack(saved)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def temporarily_clear_interpreter_stack() -> Generator[list[Any], None, None]:
|
||||
stack: list[Any] = []
|
||||
try:
|
||||
while torch._C._functorch.peek_interpreter_stack() is not None:
|
||||
stack.append(pop_dynamic_layer_stack())
|
||||
yield list(stack)
|
||||
finally:
|
||||
while stack:
|
||||
push_dynamic_layer_stack(stack.pop())
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def temporarily_restore_interpreter_stack(
|
||||
stack: list[Any] | None,
|
||||
) -> Generator[None, None, None]:
|
||||
pushed: list[Any] = []
|
||||
if stack is None:
|
||||
return
|
||||
try:
|
||||
for s in reversed(stack):
|
||||
push_dynamic_layer_stack(s)
|
||||
pushed.append(s)
|
||||
yield
|
||||
finally:
|
||||
for _ in reversed(pushed):
|
||||
# TODO: would be nice to assert that the layers are the same, but
|
||||
# Python object identity is not preserved
|
||||
pop_dynamic_layer_stack()
|
||||
|
||||
|
||||
class VmapInterpreter(FuncTorchInterpreter):
|
||||
def __init__(self, cdata: CInterpreter) -> None:
|
||||
if cdata.key() != TransformType.Vmap:
|
||||
raise AssertionError(f"expected TransformType.Vmap, got {cdata.key()}")
|
||||
# NOTE: [Interpreter cdata vs cptr]
|
||||
# cdata is a generic CInterpreter. We wrap it in a CVmapInterpreterPtr
|
||||
# so that we can access methods specific to the vmap interpreter
|
||||
self._cdata = cdata
|
||||
|
||||
@cached_property
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _cptr(self) -> CVmapInterpreterPtr:
|
||||
return CVmapInterpreterPtr(self._cdata)
|
||||
|
||||
def process(self, op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
kernel = op.functorch_table[TransformType.Vmap]
|
||||
return kernel(self, *args, **kwargs)
|
||||
|
||||
def batch_size(self) -> int:
|
||||
return self._cptr.batchSize()
|
||||
|
||||
def randomness(self) -> str:
|
||||
typ = self._cptr.randomness()
|
||||
if typ == RandomnessType.Error:
|
||||
return "error"
|
||||
elif typ == RandomnessType.Same:
|
||||
return "same"
|
||||
elif typ == RandomnessType.Different:
|
||||
return "different"
|
||||
raise RuntimeError(f"Unknown RandomnessType: {typ}")
|
||||
|
||||
def get_state(self) -> tuple[Any, ...]:
|
||||
return (self.key().name, self.level(), self.randomness())
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def nested(
|
||||
*contexts: contextlib.AbstractContextManager[Any],
|
||||
) -> Generator[tuple[contextlib.AbstractContextManager[Any], ...], None, None]:
|
||||
with contextlib.ExitStack() as stack:
|
||||
for ctx in contexts:
|
||||
stack.enter_context(ctx)
|
||||
yield contexts
|
||||
|
||||
|
||||
class GradInterpreter(FuncTorchInterpreter):
|
||||
def __init__(self, cdata: CInterpreter) -> None:
|
||||
if cdata.key() != TransformType.Grad:
|
||||
raise AssertionError(f"expected TransformType.Grad, got {cdata.key()}")
|
||||
# See NOTE: [Interpreter cdata vs cptr]
|
||||
self._cdata = cdata
|
||||
|
||||
@cached_property
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _cptr(self) -> CGradInterpreterPtr:
|
||||
return CGradInterpreterPtr(self._cdata)
|
||||
|
||||
def lift(
|
||||
self, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> tuple[tuple[Any, ...], dict[str, Any]]:
|
||||
args, kwargs = pytree.tree_map_only(
|
||||
torch.Tensor, self._cptr.lift, [args, kwargs]
|
||||
)
|
||||
return args, kwargs
|
||||
|
||||
def process(self, op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
kernel = op.functorch_table[TransformType.Grad]
|
||||
args, kwargs = self.lift(args, kwargs)
|
||||
return kernel(self, *args, **kwargs)
|
||||
|
||||
# GradInterpreter has custom lower because of the no_grad interaction
|
||||
# See NOTE [grad and vjp interaction with no_grad]
|
||||
# This logic is mirrored from C++ GradInterpreterPtr::sendToNextInterpreter
|
||||
def lower(self) -> contextlib.AbstractContextManager[Any]:
|
||||
prev_grad_mode = self.prev_grad_mode()
|
||||
if not prev_grad_mode:
|
||||
return nested(torch.no_grad(), super().lower())
|
||||
return super().lower()
|
||||
|
||||
def prev_grad_mode(self) -> bool:
|
||||
return self._cptr.prevGradMode()
|
||||
|
||||
def get_state(self) -> tuple[Any, ...]:
|
||||
return (self.key().name, self.level(), self.prev_grad_mode())
|
||||
|
||||
|
||||
class JvpInterpreter(FuncTorchInterpreter):
|
||||
def __init__(self, cdata: CInterpreter) -> None:
|
||||
if cdata.key() != TransformType.Jvp:
|
||||
raise AssertionError(f"expected TransformType.Jvp, got {cdata.key()}")
|
||||
# See NOTE: [Interpreter cdata vs cptr]
|
||||
self._cdata = cdata
|
||||
|
||||
@cached_property
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _cptr(self) -> CJvpInterpreterPtr:
|
||||
return CJvpInterpreterPtr(self._cdata)
|
||||
|
||||
def lift(
|
||||
self, args: tuple[Any, ...], kwargs: dict[str, Any]
|
||||
) -> tuple[tuple[Any, ...], dict[str, Any]]:
|
||||
args, kwargs = pytree.tree_map_only(
|
||||
torch.Tensor, self._cptr.lift, [args, kwargs]
|
||||
)
|
||||
return args, kwargs
|
||||
|
||||
def process(self, op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
kernel = op.functorch_table[TransformType.Jvp]
|
||||
args, kwargs = self.lift(args, kwargs)
|
||||
return kernel(self, *args, **kwargs)
|
||||
|
||||
# Jvp has custom lower because of the no_fwd_grad interaction
|
||||
# See NOTE [grad and vjp interaction with no_grad] for related info.
|
||||
# This logic is mirrored from C++ JvpInterpreterPtr::sendToNextInterpreter
|
||||
def lower(self) -> contextlib.AbstractContextManager[Any]:
|
||||
prev_fwd_grad_mode = self.prev_fwd_grad_mode()
|
||||
if not prev_fwd_grad_mode:
|
||||
return nested(_set_fwd_grad_enabled(False), super().lower())
|
||||
return super().lower()
|
||||
|
||||
def prev_fwd_grad_mode(self) -> bool:
|
||||
return self._cptr.prevFwdGradMode()
|
||||
|
||||
def get_state(self) -> tuple[Any, ...]:
|
||||
return (self.key().name, self.level(), self.prev_fwd_grad_mode())
|
||||
|
||||
|
||||
class FunctionalizeInterpreter(FuncTorchInterpreter):
|
||||
def __init__(self, cdata: CInterpreter) -> None:
|
||||
if cdata.key() != TransformType.Functionalize:
|
||||
raise AssertionError(
|
||||
f"expected TransformType.Functionalize, got {cdata.key()}"
|
||||
)
|
||||
self._cdata = cdata
|
||||
|
||||
@cached_property
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _cptr(self) -> CFunctionalizeInterpreterPtr:
|
||||
return CFunctionalizeInterpreterPtr(self._cdata)
|
||||
|
||||
def process(self, op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
kernel = op.functorch_table[TransformType.Functionalize]
|
||||
return kernel(self, *args, **kwargs)
|
||||
|
||||
def functionalize_add_back_views(self) -> bool:
|
||||
return self._cptr.functionalizeAddBackViews()
|
||||
|
||||
def get_state(self) -> tuple[Any, ...]:
|
||||
return (self.key().name, self.level())
|
||||
|
||||
|
||||
def coerce_cinterpreter(cinterpreter: CInterpreter) -> FuncTorchInterpreter:
|
||||
key = cinterpreter.key()
|
||||
if key == TransformType.Grad:
|
||||
return GradInterpreter(cinterpreter)
|
||||
if key == TransformType.Vmap:
|
||||
return VmapInterpreter(cinterpreter)
|
||||
if key == TransformType.Jvp:
|
||||
return JvpInterpreter(cinterpreter)
|
||||
if key == TransformType.Functionalize:
|
||||
return FunctionalizeInterpreter(cinterpreter)
|
||||
raise RuntimeError(f"NYI: PyDispatcher has not implemented support for {key}")
|
||||
|
||||
|
||||
def retrieve_current_functorch_interpreter() -> FuncTorchInterpreter:
|
||||
interpreter = torch._C._functorch.peek_interpreter_stack()
|
||||
if interpreter is None:
|
||||
raise AssertionError("interpreter must not be None")
|
||||
return coerce_cinterpreter(interpreter)
|
||||
|
||||
|
||||
def retrieve_all_functorch_interpreters() -> list[FuncTorchInterpreter]:
|
||||
cis = torch._C._functorch.get_interpreter_stack()
|
||||
if cis is None:
|
||||
return []
|
||||
return [coerce_cinterpreter(ci) for ci in cis]
|
||||
|
||||
|
||||
def compare_functorch_state(states: list[tuple[Any, ...]]) -> bool:
|
||||
# There are four possible cases covered here:
|
||||
# 1. Current stack empty AND stack when generated not empty -> Invalidate
|
||||
# 2. Current stack not empty AND stack when generated empty -> Invalidate
|
||||
# 3. Current stack and generated stack empty -> Valid FX graph
|
||||
# 4. Current stack and generated stack not empty -> Valid if both states match
|
||||
peek = torch._C._functorch.peek_interpreter_stack()
|
||||
if (peek is None and len(states) != 0) or (peek is not None and len(states) == 0):
|
||||
return False
|
||||
|
||||
cis = retrieve_all_functorch_interpreters()
|
||||
return len(cis) == len(states) and all(
|
||||
ci.check_state(state) for ci, state in zip(cis, states)
|
||||
)
|
||||
|
||||
|
||||
def dispatch_functorch(op: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
interpreter = retrieve_current_functorch_interpreter()
|
||||
# In traditional PyTorch operators, DispatchKey::FuncTorchTensorWrapper's
|
||||
# unwrap_dead_tensors fallback handles unwrapping dead tensor wrappers.
|
||||
# PyDispatcher sidesteps the PyTorch dispatcher when dealing with functorch
|
||||
# transforms, so we manually unwrap the dead tensors here.
|
||||
# This logic won't need to exist when we have mode-only functorch.
|
||||
args, kwargs = pytree.tree_map_only(
|
||||
torch.Tensor, torch._C._functorch.unwrap_if_dead, (args, kwargs)
|
||||
)
|
||||
return interpreter.process(op, args, kwargs)
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
__all__ = ["make_fx", "dispatch_trace", "PythonKeyTracer", "pythonkey_decompose"]
|
||||
from torch.fx.experimental.proxy_tensor import (
|
||||
decompose,
|
||||
dispatch_trace,
|
||||
make_fx,
|
||||
PythonKeyTracer,
|
||||
)
|
||||
|
||||
|
||||
pythonkey_decompose = decompose
|
||||
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
|
||||
import warnings
|
||||
|
||||
# TODO: remove this file when the migration of the pytree utility is done
|
||||
from torch.utils._pytree import tree_map_, treespec_pprint
|
||||
|
||||
|
||||
__all__ = ["tree_map_", "treespec_pprint"]
|
||||
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("always")
|
||||
warnings.warn(
|
||||
"`torch._functorch.pytree_hacks` is deprecated and will be removed in a future release. "
|
||||
"Please `use torch.utils._pytree` instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
+631
@@ -0,0 +1,631 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
"""
|
||||
From https://docs.google.com/spreadsheets/d/12R3nCOLskxPYjjiNkdqy4OdQ65eQp_htebXGODsjSeA/edit#gid=0
|
||||
Try to keep this list in sync with that.
|
||||
"""
|
||||
|
||||
import operator
|
||||
|
||||
|
||||
top_torch: list[tuple[str, int]] = [
|
||||
("t", 6837449),
|
||||
("tensor", 585786),
|
||||
("mode", 462182),
|
||||
("cat", 394818),
|
||||
("max", 368038),
|
||||
("zeros", 329495),
|
||||
("load", 327756),
|
||||
("no_grad", 294694),
|
||||
("save", 265130),
|
||||
("from_numpy", 243063),
|
||||
("manual_seed", 165044),
|
||||
("ones", 153696),
|
||||
("randn", 150796),
|
||||
("stack", 133358),
|
||||
("sum", 130772),
|
||||
("arange", 98087),
|
||||
("rand", 94715),
|
||||
("mean", 88546),
|
||||
("exp", 73883),
|
||||
("zeros_like", 72831),
|
||||
("min", 72248),
|
||||
("sigmoid", 66798),
|
||||
("log", 62135),
|
||||
("matmul", 47811),
|
||||
("clamp", 45304),
|
||||
("sqrt", 44911),
|
||||
("abs", 43535),
|
||||
("tanh", 42793),
|
||||
("empty", 40311),
|
||||
("argmax", 38435),
|
||||
("bmm", 33984),
|
||||
("pow", 33571),
|
||||
("norm", 31125),
|
||||
("mm", 30995),
|
||||
("is_tensor", 29546),
|
||||
("ones_like", 29512),
|
||||
("nonzero", 28681),
|
||||
("full", 28373),
|
||||
("unsqueeze", 27911),
|
||||
("where", 26585),
|
||||
("randperm", 26450),
|
||||
("eye", 24342),
|
||||
("mul", 23236),
|
||||
("topk", 22537),
|
||||
("as_tensor", 21967),
|
||||
("sort", 21412),
|
||||
("squeeze", 20863),
|
||||
("randint", 20771),
|
||||
("linspace", 20041),
|
||||
("add", 19201),
|
||||
("transpose", 18663),
|
||||
("split", 18325),
|
||||
("gather", 17904),
|
||||
("set_grad_enabled", 16013),
|
||||
("sin", 15669),
|
||||
("cos", 15562),
|
||||
("div", 15513),
|
||||
("index_select", 14866),
|
||||
("multinomial", 14331),
|
||||
("flatten", 14267),
|
||||
("isnan", 14170),
|
||||
("randn_like", 13096),
|
||||
("eq", 12680),
|
||||
("einsum", 12480),
|
||||
("round", 12367),
|
||||
("floor", 11628),
|
||||
("allclose", 11000),
|
||||
("reshape", 10605),
|
||||
("diag", 10167),
|
||||
("chunk", 9581),
|
||||
("std", 9379),
|
||||
("set_default_tensor_type", 9281),
|
||||
("triu", 8559),
|
||||
("meshgrid", 8292),
|
||||
("set_num_threads", 8126),
|
||||
("unique", 7964),
|
||||
("full_like", 7780),
|
||||
("tril", 7538),
|
||||
("dot", 7275),
|
||||
("sign", 6943),
|
||||
("equal", 6916),
|
||||
("normal", 6750),
|
||||
("cumsum", 6556),
|
||||
("dist", 6058),
|
||||
("isfinite", 6030),
|
||||
("gt", 5935),
|
||||
("set_printoptions", 5888),
|
||||
("range", 5491),
|
||||
("empty_like", 5351),
|
||||
("flip", 5342),
|
||||
("masked_select", 5341),
|
||||
("bernoulli", 5262),
|
||||
("atan", 5253),
|
||||
("var", 5247),
|
||||
("prod", 5200),
|
||||
("erf", 5088),
|
||||
("inverse", 5072),
|
||||
("addmm", 4854),
|
||||
("logsumexp", 4582),
|
||||
("fft", 4436),
|
||||
("lt", 4421),
|
||||
("log2", 4316),
|
||||
("enable_grad", 4238),
|
||||
("rand_like", 4187),
|
||||
("argsort", 3972),
|
||||
("seed", 3932),
|
||||
("mv", 3547),
|
||||
("ger", 3309),
|
||||
("ge", 3248),
|
||||
("atan2", 3210),
|
||||
("ceil", 3202),
|
||||
("ne", 3075),
|
||||
("bincount", 3063),
|
||||
("acos", 3055),
|
||||
("rsqrt", 3031),
|
||||
("svd", 3029),
|
||||
("numel", 3003),
|
||||
("log1p", 2840),
|
||||
("unbind", 2808),
|
||||
("le", 2714),
|
||||
("isinf", 2707),
|
||||
("cross", 2646),
|
||||
("set_default_dtype", 2536),
|
||||
("argmin", 2535),
|
||||
("sparse_coo_tensor", 2489),
|
||||
("log10", 2304),
|
||||
("kthvalue", 2192),
|
||||
("set_rng_state", 2158),
|
||||
("get_rng_state", 1996),
|
||||
("get_default_dtype", 1879),
|
||||
("det", 1868),
|
||||
("qr", 1864),
|
||||
("histc", 1852),
|
||||
("symeig", 1832),
|
||||
("trace", 1801),
|
||||
("median", 1795),
|
||||
("addcmul", 1751),
|
||||
("remainder", 1717),
|
||||
("baddbmm", 1693),
|
||||
("lgamma", 1665),
|
||||
("repeat_interleave", 1598),
|
||||
("fmod", 1576),
|
||||
("reciprocal", 1575),
|
||||
("tan", 1560),
|
||||
("initial_seed", 1532),
|
||||
("take", 1529),
|
||||
("stft", 1487),
|
||||
("get_num_threads", 1477),
|
||||
("real", 1459),
|
||||
("cholesky", 1406),
|
||||
("quantize_per_tensor", 1392),
|
||||
("diag_embed", 1364),
|
||||
("lerp", 1363),
|
||||
("asin", 1345),
|
||||
("eig", 1333),
|
||||
("trunc", 1290),
|
||||
("diagonal", 1287),
|
||||
("cosh", 1279),
|
||||
("rfft", 1269),
|
||||
("cumprod", 1260),
|
||||
("addr", 1211),
|
||||
("roll", 1198),
|
||||
("narrow", 1188),
|
||||
("digamma", 1172),
|
||||
("square", 1163),
|
||||
("sinh", 1131),
|
||||
("logspace", 1084),
|
||||
("broadcast_tensors", 1070),
|
||||
("irfft", 1013),
|
||||
("frac", 997),
|
||||
("hann_window", 994),
|
||||
("solve", 989),
|
||||
("logdet", 977),
|
||||
("expm1", 968),
|
||||
("cdist", 946),
|
||||
("addmv", 903),
|
||||
("randint_like", 888),
|
||||
("tensordot", 888),
|
||||
("ifft", 877),
|
||||
("true_divide", 854),
|
||||
("erfinv", 830),
|
||||
("addcdiv", 819),
|
||||
("addbmm", 813),
|
||||
("renorm", 781),
|
||||
("pinverse", 753),
|
||||
("isclose", 740),
|
||||
("erfc", 729),
|
||||
("is_storage", 725),
|
||||
("triangular_solve", 723),
|
||||
("rot90", 709),
|
||||
("logical_not", 686),
|
||||
("geqrf", 681),
|
||||
("slogdet", 677),
|
||||
("lu", 665),
|
||||
("hamming_window", 659),
|
||||
("orgqr", 651),
|
||||
("ormqr", 622),
|
||||
("is_floating_point", 602),
|
||||
("diagflat", 562),
|
||||
("cholesky_solve", 559),
|
||||
("tril_indices", 552),
|
||||
("chain_matmul", 551),
|
||||
("triu_indices", 548),
|
||||
("angle", 522),
|
||||
("poisson", 505),
|
||||
("matrix_power", 485),
|
||||
("unique_consecutive", 471),
|
||||
("quantize_per_channel", 465),
|
||||
("std_mean", 458),
|
||||
("bartlett_window", 447),
|
||||
("var_mean", 428),
|
||||
("lstsq", 421),
|
||||
("logical_and", 419),
|
||||
("mvlgamma", 411),
|
||||
("blackman_window", 400),
|
||||
("bitwise_not", 395),
|
||||
("cholesky_inverse", 388),
|
||||
("as_strided", 384),
|
||||
("floor_divide", 353),
|
||||
("cartesian_prod", 321),
|
||||
("lu_solve", 317),
|
||||
("set_flush_denormal", 310),
|
||||
("empty_strided", 283),
|
||||
("logical_xor", 282),
|
||||
("polygamma", 282),
|
||||
("logical_or", 280),
|
||||
("set_num_interop_threads", 278),
|
||||
("combinations", 274),
|
||||
("trapz", 270),
|
||||
("matrix_rank", 260),
|
||||
("lu_unpack", 255),
|
||||
("result_type", 244),
|
||||
("conj", 231),
|
||||
("cummax", 230),
|
||||
("lobpcg", 229),
|
||||
("bitwise_xor", 217),
|
||||
("promote_types", 213),
|
||||
("get_num_interop_threads", 211),
|
||||
("cummin", 205),
|
||||
("bitwise_and", 198),
|
||||
("dequantize", 192),
|
||||
("bitwise_or", 191),
|
||||
("imag", 191),
|
||||
("can_cast", 184),
|
||||
("istft", 180),
|
||||
("compiled_with_cxx11_abi", 159),
|
||||
("is_complex", 151),
|
||||
("block_diag", 136),
|
||||
("pca_lowrank", 124),
|
||||
("absolute", 122),
|
||||
("svd_lowrank", 108),
|
||||
("neg", 2),
|
||||
]
|
||||
|
||||
top_nn_functional: list[tuple[str, int]] = [
|
||||
("nn.functional.softmax", 10522),
|
||||
("nn.functional.relu", 8572),
|
||||
("nn.functional.interpolate", 7277),
|
||||
("nn.functional.pad", 5207),
|
||||
("nn.functional.log_softmax", 4699),
|
||||
("nn.functional.normalize", 2338),
|
||||
("nn.functional.cross_entropy", 2083),
|
||||
("nn.functional.grid_sample", 1970),
|
||||
("nn.functional.one_hot", 1967),
|
||||
("nn.functional.mse_loss", 1920),
|
||||
("nn.functional.conv2d", 1593),
|
||||
("nn.functional.dropout", 1516),
|
||||
("nn.functional.softplus", 1385),
|
||||
("nn.functional.sigmoid", 1128),
|
||||
("nn.functional.linear", 1036),
|
||||
("nn.functional.gelu", 930),
|
||||
("nn.functional.avg_pool2d", 899),
|
||||
("nn.functional.max_pool2d", 876),
|
||||
("nn.functional.nll_loss", 863),
|
||||
("nn.functional.embedding", 737),
|
||||
("nn.functional.tanh", 664),
|
||||
("nn.functional.leaky_relu", 640),
|
||||
("nn.functional.adaptive_avg_pool2d", 633),
|
||||
("nn.functional.cosine_similarity", 627),
|
||||
("nn.functional.unfold", 609),
|
||||
("nn.functional.conv1d", 596),
|
||||
("nn.functional.binary_cross_entropy_with_logits", 591),
|
||||
("nn.functional.l1_loss", 571),
|
||||
("nn.functional.binary_cross_entropy", 492),
|
||||
("nn.functional.elu", 416),
|
||||
("nn.functional.batch_norm", 413),
|
||||
("nn.functional.upsample", 413),
|
||||
("nn.functional.fold", 305),
|
||||
("nn.functional.affine_grid", 298),
|
||||
("nn.functional.max_pool1d", 297),
|
||||
("nn.functional.torch", 294),
|
||||
("nn.functional.threshold", 263),
|
||||
("nn.functional.smooth_l1_loss", 262),
|
||||
("nn.functional.pairwise_distance", 253),
|
||||
("nn.functional.logsigmoid", 243),
|
||||
("nn.functional.adaptive_max_pool2d", 235),
|
||||
("nn.functional.relu6", 213),
|
||||
("nn.functional.pixel_shuffle", 209),
|
||||
("nn.functional.avg_pool3d", 203),
|
||||
("nn.functional.bilinear", 203),
|
||||
("nn.functional.conv_transpose2d", 201),
|
||||
("nn.functional.gumbel_softmax", 197),
|
||||
("nn.functional.max_unpool2d", 196),
|
||||
("nn.functional.kl_div", 191),
|
||||
("nn.functional.hardtanh", 189),
|
||||
("nn.functional.ctc_loss", 185),
|
||||
("nn.functional.layer_norm", 178),
|
||||
("nn.functional.conv3d", 172),
|
||||
("nn.functional.max_unpool3d", 167),
|
||||
("nn.functional.hardshrink", 165),
|
||||
("nn.functional.hardswish", 156),
|
||||
("nn.functional.selu", 156),
|
||||
("nn.functional.glu", 155),
|
||||
("nn.functional.assert_int_or_pair", 150),
|
||||
("nn.functional.hardsigmoid", 146),
|
||||
("nn.functional.upsample_bilinear", 146),
|
||||
("nn.functional.max_pool3d", 140),
|
||||
("nn.functional.adaptive_avg_pool3d", 139),
|
||||
("nn.functional.instance_norm", 124),
|
||||
("nn.functional.embedding_bag", 122),
|
||||
("nn.functional.upsample_nearest", 110),
|
||||
("nn.functional.avg_pool1d", 105),
|
||||
("nn.functional.prelu", 102),
|
||||
("nn.functional.celu", 92),
|
||||
("nn.functional.dropout2d", 86),
|
||||
("nn.functional.hinge_embedding_loss", 82),
|
||||
("nn.functional.softsign", 81),
|
||||
("nn.functional.max_unpool1d", 74),
|
||||
("nn.functional.silu", 74),
|
||||
("nn.functional.softshrink", 70),
|
||||
("nn.functional.leaky_relu_", 68),
|
||||
("nn.functional.softmin", 67),
|
||||
("nn.functional.channel_shuffle", 66),
|
||||
("nn.functional.multilabel_margin_loss", 66),
|
||||
("nn.functional.dropout3d", 65),
|
||||
("nn.functional.multi_margin_loss", 65),
|
||||
("nn.functional.lp_pool2d", 64),
|
||||
("nn.functional.conv_transpose1d", 62),
|
||||
("nn.functional.triplet_margin_loss", 62),
|
||||
("nn.functional.tanhshrink", 61),
|
||||
("nn.functional.adaptive_max_pool1d", 59),
|
||||
("nn.functional.cosine_embedding_loss", 58),
|
||||
("nn.functional.multi_head_attention_forward", 58),
|
||||
("nn.functional.max_pool1d_with_indices", 53),
|
||||
("nn.functional.poisson_nll_loss", 53),
|
||||
("nn.functional.margin_ranking_loss", 52),
|
||||
("nn.functional.soft_margin_loss", 52),
|
||||
("nn.functional.adaptive_max_pool3d", 51),
|
||||
("nn.functional.group_norm", 51),
|
||||
("nn.functional.local_response_norm", 51),
|
||||
("nn.functional.multilabel_soft_margin_loss", 51),
|
||||
("nn.functional.relu_", 50),
|
||||
("nn.functional.alpha_dropout", 49),
|
||||
("nn.functional.feature_alpha_dropout", 49),
|
||||
("nn.functional.lp_pool1d", 49),
|
||||
("nn.functional.adaptive_max_pool1d_with_indices", 48),
|
||||
("nn.functional.adaptive_max_pool2d_with_indices", 48),
|
||||
("nn.functional.adaptive_max_pool3d_with_indices", 48),
|
||||
("nn.functional.fractional_max_pool2d", 48),
|
||||
("nn.functional.fractional_max_pool2d_with_indices", 48),
|
||||
("nn.functional.fractional_max_pool3d", 48),
|
||||
("nn.functional.fractional_max_pool3d_with_indices", 48),
|
||||
("nn.functional.max_pool2d_with_indices", 48),
|
||||
("nn.functional.max_pool3d_with_indices", 48),
|
||||
("nn.functional.handle_torch_function", 47),
|
||||
("nn.functional.has_torch_function", 47),
|
||||
("nn.functional.adaptive_avg_pool1d", 43),
|
||||
("nn.functional.pdist", 43),
|
||||
("nn.functional.rrelu_", 37),
|
||||
("nn.functional.elu_", 34),
|
||||
("nn.functional.boolean_dispatch", 33),
|
||||
("nn.functional.hardtanh_", 26),
|
||||
("nn.functional.triplet_margin_with_distance_loss", 23),
|
||||
("nn.functional.selu_", 20),
|
||||
("nn.functional.pixel_unshuffle", 19),
|
||||
("nn.functional.conv_transpose3d", 18),
|
||||
("nn.functional.gaussian_nll_loss", 15),
|
||||
("nn.functional.has_torch_function_unary", 15),
|
||||
("nn.functional.has_torch_function_variadic", 15),
|
||||
("nn.functional.celu_", 13),
|
||||
("nn.functional.huber_loss", 7),
|
||||
("nn.functional.mish", 4),
|
||||
("nn.functional.threshold_", 3),
|
||||
("nn.functional.grad", 2),
|
||||
("nn.functional.conv_tbc", 1),
|
||||
("nn.functional.math", 1),
|
||||
]
|
||||
|
||||
top_nn_module: list[tuple[str, int, str | None]] = [
|
||||
("nn.Module", 927129, None),
|
||||
("nn.Linear", 530688, "nn.functional.linear"),
|
||||
("nn.Sequential", 384968, None),
|
||||
("nn.Conv2d", 383320, "nn.functional.conv2d"),
|
||||
("nn.ReLU", 318877, "nn.functional.relu"),
|
||||
("nn.BatchNorm2d", 233265, "nn.functional.batch_norm"),
|
||||
("nn.Dropout", 179268, "nn.functional.dropout"),
|
||||
("nn.ModuleList", 171225, None),
|
||||
("nn.Parameter", 153291, None),
|
||||
("nn.CrossEntropyLoss", 152696, "nn.functional.cross_entropy"),
|
||||
("nn.MaxPool2d", 138619, "nn.functional.max_pool2d"),
|
||||
("nn.Embedding", 111844, "nn.functional.embedding"),
|
||||
("nn.DataParallel", 104238, None),
|
||||
("nn.MSELoss", 82954, "nn.functional.mse_loss"),
|
||||
("nn.Sigmoid", 75810, "nn.functional.sigmoid"),
|
||||
("nn.LeakyReLU", 65632, "nn.functional.leaky_relu"),
|
||||
("nn.BatchNorm1d", 65374, "nn.functional.batch_norm"),
|
||||
("nn.Softmax", 65114, "nn.functional.softmax"),
|
||||
("nn.Tanh", 59445, "nn.functional.tanh"),
|
||||
("nn.AdaptiveAvgPool2d", 59071, "nn.functional.adaptive_avg_pool2d"),
|
||||
("nn.AvgPool2d", 58377, "nn.functional.avg_pool2d"),
|
||||
("nn.ConvTranspose2d", 57524, "nn.functional.conv_transpose2d"),
|
||||
("nn.LSTM", 57411, None),
|
||||
("nn.Conv1d", 41108, "nn.functional.conv1d"),
|
||||
("nn.LayerNorm", 36089, "nn.functional.layer_norm"),
|
||||
("nn.BCELoss", 34005, "nn.functional.binary_cross_entropy"),
|
||||
("nn.Upsample", 32527, "nn.functional.interpolate"),
|
||||
("nn.BCEWithLogitsLoss", 29944, "nn.functional.binary_cross_entropy_with_logits"),
|
||||
("nn.GRU", 25421, None),
|
||||
("nn.Dropout2d", 23512, "nn.functional.dropout2d"),
|
||||
("nn.LogSoftmax", 22897, "nn.functional.log_softmax"),
|
||||
("nn.L1Loss", 22778, "nn.functional.l1_loss"),
|
||||
("nn.GroupNorm", 22183, "nn.functional.group_norm"),
|
||||
("nn.NLLLoss", 21751, "nn.functional.nll_loss"),
|
||||
("nn.Conv3d", 20874, "nn.functional.conv3d"),
|
||||
("nn.Identity", 17911, None),
|
||||
("nn.InstanceNorm2d", 16426, "nn.functional.instance_norm"),
|
||||
("nn.BatchNorm3d", 16378, "nn.functional.batch_norm"),
|
||||
("nn.PReLU", 13472, "nn.functional.prelu"),
|
||||
("nn.ReLU6", 12622, "nn.functional.relu6"),
|
||||
("nn.ELU", 12508, "nn.functional.elu"),
|
||||
("nn.LSTMCell", 10885, None),
|
||||
("nn.Flatten", 10384, "torch.flatten"),
|
||||
("nn.ModuleDict", 10255, None),
|
||||
("nn.ReflectionPad2d", 9954, "nn.functional.pad"),
|
||||
("nn.MaxPool3d", 9526, "nn.functional.max_pool3d"),
|
||||
("nn.MaxPool1d", 9154, "nn.functional.max_pool1d"),
|
||||
("nn.RNN", 9154, None),
|
||||
("nn.ZeroPad2d", 8847, "nn.functional.pad"),
|
||||
("nn.ParameterList", 7702, None),
|
||||
("nn.SyncBatchNorm", 6814, None),
|
||||
("nn.PixelShuffle", 6571, "nn.functional.pixel_shuffle"),
|
||||
("nn.SmoothL1Loss", 6517, "nn.functional.smooth_l1_loss"),
|
||||
("nn.Hardswish", 6458, "nn.functional.hardswish"),
|
||||
("nn.AdaptiveMaxPool2d", 6071, "nn.functional.adaptive_max_pool2d"),
|
||||
("nn.SELU", 6043, "nn.functional.selu"),
|
||||
("nn.ConvTranspose3d", 6039, "nn.functional.conv_transpose3d"),
|
||||
("nn.GRUCell", 5840, None),
|
||||
("nn.ReplicationPad2d", 5600, "nn.functional.pad"),
|
||||
("nn.KLDivLoss", 5541, "nn.functional.kl_div"),
|
||||
("nn.ConvTranspose1d", 5183, "nn.functional.conv_transpose1d"),
|
||||
("nn.Softplus", 5120, "nn.functional.softplus"),
|
||||
("nn.SiLU", 4895, "nn.functional.silu"),
|
||||
("nn.AvgPool3d", 4523, "nn.functional.avg_pool3d"),
|
||||
("nn.CosineSimilarity", 4058, "nn.functional.cosine_similarity"),
|
||||
("nn.GELU", 3932, "nn.functional.gelu"),
|
||||
("nn.UpsamplingBilinear2d", 3673, "nn.functional.interpolate"),
|
||||
("nn.InstanceNorm1d", 3658, "nn.functional.instance_norm"),
|
||||
("nn.Transformer", 3604, None),
|
||||
("nn.MultiheadAttention", 3435, "nn.functional.multi_head_attention_forward"),
|
||||
("nn.AvgPool1d", 3195, "nn.functional.avg_pool1d"),
|
||||
("nn.Dropout3d", 2964, "nn.functional.dropout3d"),
|
||||
("nn.AdaptiveAvgPool3d", 2915, "nn.functional.adaptive_avg_pool3d"),
|
||||
("nn.InstanceNorm3d", 2893, "nn.functional.instance_norm"),
|
||||
("nn.Hardtanh", 2613, "nn.functional.hardtanh"),
|
||||
("nn.MarginRankingLoss", 2568, "nn.functional.margin_ranking_loss"),
|
||||
("nn.GLU", 2526, "nn.functional.glu"),
|
||||
("nn.AdaptiveAvgPool1d", 2481, "nn.functional.adaptive_avg_pool1d"),
|
||||
("nn.EmbeddingBag", 2344, "nn.functional.embedding_bag"),
|
||||
("nn.TransformerEncoderLayer", 2292, None),
|
||||
("nn.TransformerEncoder", 2091, None),
|
||||
("nn.MaxUnpool2d", 2031, "nn.functional.max_unpool2d"),
|
||||
("nn.UpsamplingNearest2d", 2004, "nn.functional.interpolate"),
|
||||
("nn.ConstantPad1d", 1904, "nn.functional.pad"),
|
||||
("nn.ConstantPad2d", 1791, "nn.functional.pad"),
|
||||
("nn.CTCLoss", 1789, "nn.functional.ctc_loss"),
|
||||
("nn.AdaptiveMaxPool1d", 1713, "nn.functional.adaptive_max_pool1d"),
|
||||
("nn.AdaptiveLogSoftmaxWithLoss", 1665, None),
|
||||
("nn.Bilinear", 1664, "nn.functional.bilinear"),
|
||||
("nn.RNNCell", 1653, None),
|
||||
("nn.MultiLabelSoftMarginLoss", 1624, "nn.functional.multilabel_soft_margin_loss"),
|
||||
("nn.Unfold", 1452, "nn.functional.unfold"),
|
||||
("nn.RReLU", 1431, "nn.functional.rrelu"),
|
||||
("nn.CosineEmbeddingLoss", 1357, "nn.functional.cosine_embedding_loss"),
|
||||
("nn.LocalResponseNorm", 1331, "nn.functional.local_response_norm"),
|
||||
("nn.Softmax2d", 1300, "nn.functional.softmax"),
|
||||
("nn.PairwiseDistance", 1241, "nn.functional.pairwise_distance"),
|
||||
("nn.LogSigmoid", 1235, "nn.functional.logsigmoid"),
|
||||
("nn.TripletMarginLoss", 1230, "nn.functional.triplet_margin_loss"),
|
||||
("nn.RNNBase", 1133, None),
|
||||
("nn.Threshold", 1043, "nn.functional.threshold"),
|
||||
("nn.AdaptiveMaxPool3d", 1025, "nn.functional.adaptive_max_pool3d"),
|
||||
("nn.CELU", 1018, "nn.functional.celu"),
|
||||
("nn.NLLLoss2d", 966, "nn.functional.nll_loss"),
|
||||
("nn.Softsign", 877, "nn.functional.softsign"),
|
||||
("nn.ReplicationPad1d", 862, "nn.functional.pad"),
|
||||
("nn.SoftMarginLoss", 856, "nn.functional.soft_margin_loss"),
|
||||
("nn.ParameterDict", 742, None),
|
||||
("nn.ReflectionPad1d", 731, "nn.functional.pad"),
|
||||
("nn.Softshrink", 713, "nn.functional.softshrink"),
|
||||
("nn.AlphaDropout", 710, "nn.functional.alpha_dropout"),
|
||||
("nn.Tanhshrink", 681, "nn.functional.tanhshrink"),
|
||||
("nn.PoissonNLLLoss", 676, "nn.functional.poisson_nll_loss"),
|
||||
("nn.MaxUnpool3d", 660, "nn.functional.max_unpool3d"),
|
||||
("nn.Fold", 630, "nn.functional.fold"),
|
||||
("nn.MultiMarginLoss", 622, "nn.functional.multi_margin_loss"),
|
||||
("nn.TransformerDecoderLayer", 614, None),
|
||||
("nn.TransformerDecoder", 607, None),
|
||||
("nn.Hardshrink", 592, "nn.functional.hardshrink"),
|
||||
("nn.ConstantPad3d", 582, "nn.functional.pad"),
|
||||
("nn.MultiLabelMarginLoss", 580, "nn.functional.multilabel_margin_loss"),
|
||||
("nn.LPPool2d", 550, "nn.functional.lp_pool2d"),
|
||||
("nn.Softmin", 537, "nn.functional.softmin"),
|
||||
("nn.MaxUnpool1d", 518, "nn.functional.max_unpool1d"),
|
||||
("nn.FractionalMaxPool2d", 484, "nn.functional.fractional_max_pool2d"),
|
||||
("nn.Hardsigmoid", 477, "nn.functional.hardsigmoid"),
|
||||
("nn.ReplicationPad3d", 470, "nn.functional.pad"),
|
||||
("nn.HingeEmbeddingLoss", 442, "nn.functional.hinge_embedding_loss"),
|
||||
("nn.LPPool1d", 386, "nn.functional.lp_pool1d"),
|
||||
("nn.FractionalMaxPool3d", 252, "nn.functional.fractional_max_pool3d"),
|
||||
("nn.Container", 217, None),
|
||||
("nn.Unflatten", 206, "nn.functional.unflatten"),
|
||||
("nn.FeatureAlphaDropout", 136, "nn.functional.feature_alpha_dropout"),
|
||||
(
|
||||
"nn.TripletMarginWithDistanceLoss",
|
||||
107,
|
||||
"nn.functional.triplet_margin_with_distance_loss",
|
||||
),
|
||||
("nn.ChannelShuffle", 90, "nn.functional.channel_shuffle"),
|
||||
("nn.RNNCellBase", 88, None),
|
||||
("nn.LazyLinear", 81, "nn.functional.linear"),
|
||||
("nn.UninitializedParameter", 60, None),
|
||||
("nn.CrossMapLRN2d", 59, None),
|
||||
("nn.GaussianNLLLoss", 55, "nn.functional.gaussian_nll_loss"),
|
||||
("nn.PixelUnshuffle", 45, "nn.functional.pixel_unshuffle"),
|
||||
("nn.Mish", 31, "nn.functional.mish"),
|
||||
("nn.ReflectionPad3d", 22, "nn.functional.pad"),
|
||||
("nn.HuberLoss", 18, "nn.functional.huber_loss"),
|
||||
("nn.LazyConv2d", 15, None),
|
||||
("nn.LazyConv1d", 9, None),
|
||||
("nn.LazyConv3d", 8, None),
|
||||
("nn.LazyConvTranspose1d", 8, None),
|
||||
("nn.LazyConvTranspose2d", 8, None),
|
||||
("nn.LazyConvTranspose3d", 8, None),
|
||||
("nn.LazyBatchNorm1d", 3, None),
|
||||
("nn.LazyBatchNorm2d", 3, None),
|
||||
("nn.LazyBatchNorm3d", 3, None),
|
||||
("nn.UninitializedBuffer", 3, None),
|
||||
]
|
||||
|
||||
# No rankings because these are a little hard to get rankings for
|
||||
method_only_ops: list[str] = [
|
||||
"bfloat16",
|
||||
"bool",
|
||||
"byte",
|
||||
"char",
|
||||
"contiguous",
|
||||
"cpu",
|
||||
"cuda",
|
||||
"detach",
|
||||
"double",
|
||||
"expand",
|
||||
"expand_as",
|
||||
"float",
|
||||
"get_device",
|
||||
"half",
|
||||
"hardshrink",
|
||||
"index_add",
|
||||
"index_copy",
|
||||
"index_fill",
|
||||
"index_put",
|
||||
"int",
|
||||
"is_contiguous",
|
||||
"is_pinned",
|
||||
"is_set_to",
|
||||
"is_shared",
|
||||
"is_signed",
|
||||
"item",
|
||||
"long",
|
||||
"masked_scatter",
|
||||
"masked_fill",
|
||||
"narrow_copy",
|
||||
"numpy",
|
||||
"pin_memory",
|
||||
"repeat",
|
||||
"reshape_as",
|
||||
"select",
|
||||
"short",
|
||||
"storage_offset",
|
||||
"sum_to_size",
|
||||
"to",
|
||||
"to_mkldnn",
|
||||
"tolist",
|
||||
"type",
|
||||
"type_as",
|
||||
"unfold",
|
||||
"view",
|
||||
"view_as",
|
||||
]
|
||||
|
||||
|
||||
def get_nn_functional_top_list() -> list[tuple[str, int]]:
|
||||
top_nn_functional_: dict[str, int] = dict(top_nn_functional)
|
||||
for _, count, functional_name in top_nn_module:
|
||||
if functional_name is None:
|
||||
continue
|
||||
if functional_name == "torch.flatten":
|
||||
continue
|
||||
if functional_name not in top_nn_functional_:
|
||||
top_nn_functional_[functional_name] = count
|
||||
else:
|
||||
top_nn_functional_[functional_name] += count
|
||||
|
||||
top_nn_functional_list = list(top_nn_functional_.items())
|
||||
top_nn_functional_list.sort(key=operator.itemgetter(1), reverse=True)
|
||||
return top_nn_functional_list
|
||||
|
||||
|
||||
usage_count: dict[str, int] = dict(get_nn_functional_top_list())
|
||||
usage_count.update(top_torch)
|
||||
@@ -0,0 +1,40 @@
|
||||
import contextlib
|
||||
from collections.abc import Generator
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._C._functorch import (
|
||||
get_single_level_autograd_function_allowed,
|
||||
set_single_level_autograd_function_allowed,
|
||||
unwrap_if_dead,
|
||||
)
|
||||
from torch.utils._exposed_in import exposed_in
|
||||
|
||||
|
||||
__all__ = [
|
||||
"exposed_in",
|
||||
"argnums_t",
|
||||
"enable_single_level_autograd_function",
|
||||
"unwrap_dead_wrappers",
|
||||
]
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def enable_single_level_autograd_function() -> Generator[None, None, None]:
|
||||
try:
|
||||
prev_state = get_single_level_autograd_function_allowed()
|
||||
set_single_level_autograd_function_allowed(True)
|
||||
yield
|
||||
finally:
|
||||
set_single_level_autograd_function_allowed(prev_state)
|
||||
|
||||
|
||||
def unwrap_dead_wrappers(args: tuple[Any, ...]) -> tuple[Any, ...]:
|
||||
# NB: doesn't use tree_map_only for performance reasons
|
||||
result = tuple(
|
||||
unwrap_if_dead(arg) if isinstance(arg, torch.Tensor) else arg for arg in args
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
argnums_t = int | tuple[int, ...]
|
||||
@@ -0,0 +1,567 @@
|
||||
# Copyright (c) Facebook, Inc. and its 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import itertools
|
||||
from collections.abc import Callable # noqa: TC003
|
||||
from functools import partial
|
||||
from typing import Any, cast, NoReturn, TYPE_CHECKING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._C._functorch import is_batchedtensor
|
||||
from torch._functorch.predispatch import (
|
||||
_add_batch_dim,
|
||||
_remove_batch_dim,
|
||||
_vmap_decrement_nesting,
|
||||
_vmap_increment_nesting,
|
||||
lazy_load_decompositions,
|
||||
)
|
||||
from torch.utils._pytree import (
|
||||
_broadcast_to_and_flatten,
|
||||
tree_flatten,
|
||||
tree_map_,
|
||||
tree_unflatten,
|
||||
TreeSpec,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator, Iterable
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
in_dims_t = int | tuple[Any, ...]
|
||||
out_dims_t = int | tuple[int, ...] | None
|
||||
|
||||
|
||||
def doesnt_support_saved_tensors_hooks(f: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
message = (
|
||||
"torch.func.{grad, vjp, jacrev, hessian} don't yet support saved tensor hooks. "
|
||||
"Please open an issue with your use case."
|
||||
)
|
||||
|
||||
@functools.wraps(f)
|
||||
def fn(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
with torch.autograd.graph.disable_saved_tensors_hooks(message):
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return fn
|
||||
|
||||
|
||||
# Checks that all args-to-be-batched have the same batch dim size
|
||||
def _validate_and_get_batch_size(
|
||||
flat_in_dims: list[int | None], flat_args: list[Any]
|
||||
) -> int:
|
||||
batch_sizes = [
|
||||
arg.size(in_dim)
|
||||
for in_dim, arg in zip(flat_in_dims, flat_args)
|
||||
if in_dim is not None
|
||||
]
|
||||
if len(batch_sizes) == 0:
|
||||
raise ValueError("vmap: Expected at least one Tensor to vmap over")
|
||||
if batch_sizes and any(size != batch_sizes[0] for size in batch_sizes):
|
||||
raise ValueError(
|
||||
f"vmap: Expected all tensors to have the same size in the mapped "
|
||||
f"dimension, got sizes {batch_sizes} for the mapped dimension"
|
||||
)
|
||||
return batch_sizes[0]
|
||||
|
||||
|
||||
def _num_outputs(batched_outputs: Tensor | tuple[Tensor, ...]) -> int:
|
||||
if isinstance(batched_outputs, tuple):
|
||||
return len(batched_outputs)
|
||||
return 1
|
||||
|
||||
|
||||
# If value is a tuple, check it has length `num_elements`.
|
||||
# If value is not a tuple, make a tuple with `value` repeated `num_elements` times
|
||||
|
||||
|
||||
def _as_tuple(
|
||||
value: tuple[_R, ...] | _R,
|
||||
num_elements: int,
|
||||
error_message_lambda: Callable[[], str],
|
||||
) -> tuple[_R, ...]:
|
||||
if not isinstance(value, tuple):
|
||||
return (value,) * num_elements
|
||||
if len(value) != num_elements:
|
||||
raise ValueError(error_message_lambda())
|
||||
return value
|
||||
|
||||
|
||||
def _process_batched_inputs(
|
||||
in_dims: in_dims_t, args: tuple[Any, ...], func: Callable[..., Any]
|
||||
) -> tuple[int, list[int | None], list[Any], TreeSpec]:
|
||||
if not isinstance(in_dims, int) and not isinstance(in_dims, tuple):
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): "
|
||||
f"expected `in_dims` to be int or a (potentially nested) tuple "
|
||||
f"matching the structure of inputs, got: {type(in_dims)}."
|
||||
)
|
||||
if len(args) == 0:
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)})(<inputs>): got no inputs. Maybe you forgot to add "
|
||||
f"inputs, or you are trying to vmap over a function with no inputs. "
|
||||
f"The latter is unsupported."
|
||||
)
|
||||
|
||||
flat_args, args_spec = tree_flatten(args)
|
||||
flat_in_dims = _broadcast_to_and_flatten(in_dims, args_spec)
|
||||
if flat_in_dims is None:
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): "
|
||||
f"in_dims is not compatible with the structure of `inputs`. "
|
||||
f"in_dims has structure {tree_flatten(in_dims)[1]} but inputs "
|
||||
f"has structure {args_spec}."
|
||||
)
|
||||
|
||||
for i, (arg, in_dim) in enumerate(zip(flat_args, flat_in_dims)):
|
||||
if not isinstance(in_dim, int) and in_dim is not None:
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): "
|
||||
f"Got in_dim={in_dim} for an input but in_dim must be either "
|
||||
f"an integer dimension or None."
|
||||
)
|
||||
if isinstance(in_dim, int) and not isinstance(arg, Tensor):
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): "
|
||||
f"Got in_dim={in_dim} for an input but the input is of type "
|
||||
f"{type(arg)}. We cannot vmap over non-Tensor arguments, "
|
||||
f"please use None as the respective in_dim"
|
||||
)
|
||||
if in_dim is not None and (in_dim < -arg.dim() or in_dim >= arg.dim()):
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): "
|
||||
f"Got in_dim={in_dim} for some input, but that input is a Tensor "
|
||||
f"of dimensionality {arg.dim()} so expected in_dim to satisfy "
|
||||
f"-{arg.dim()} <= in_dim < {arg.dim()}."
|
||||
)
|
||||
if in_dim is not None and in_dim < 0:
|
||||
flat_in_dims[i] = in_dim % arg.dim()
|
||||
|
||||
return (
|
||||
_validate_and_get_batch_size(flat_in_dims, flat_args),
|
||||
flat_in_dims,
|
||||
flat_args,
|
||||
args_spec,
|
||||
)
|
||||
|
||||
|
||||
# Creates BatchedTensors for every Tensor in arg that should be batched.
|
||||
# Returns the (potentially) batched arguments and the batch_size.
|
||||
|
||||
|
||||
# TODO: See if we can explain how flat works to the type checker
|
||||
def _create_batched_inputs(
|
||||
flat_in_dims: list[int | None],
|
||||
flat_args: list[Any],
|
||||
vmap_level: int,
|
||||
args_spec: TreeSpec,
|
||||
) -> tuple[Any, ...]:
|
||||
# See NOTE [Ignored _remove_batch_dim, _add_batch_dim]
|
||||
batched_inputs = [
|
||||
arg if in_dim is None else _add_batch_dim(arg, in_dim, vmap_level)
|
||||
for in_dim, arg in zip(flat_in_dims, flat_args)
|
||||
]
|
||||
return tree_unflatten(batched_inputs, args_spec)
|
||||
|
||||
|
||||
def _maybe_remove_batch_dim(
|
||||
name: str,
|
||||
batched_output: Any,
|
||||
vmap_level: int,
|
||||
batch_size: int,
|
||||
out_dim: int | None,
|
||||
) -> torch.Tensor:
|
||||
if out_dim is None:
|
||||
if isinstance(batched_output, torch.Tensor) and is_batchedtensor(
|
||||
batched_output
|
||||
):
|
||||
raise ValueError(
|
||||
f"vmap({name}, ...): `{name}` can not return a "
|
||||
f"BatchedTensor when out_dim is None"
|
||||
)
|
||||
return batched_output
|
||||
|
||||
# out_dim is non None
|
||||
if not isinstance(batched_output, torch.Tensor):
|
||||
raise ValueError(
|
||||
f"vmap({name}, ...): `{name}` must only return "
|
||||
f"Tensors, got type {type(batched_output)}. "
|
||||
"Did you mean to set out_dims= to None for output?"
|
||||
)
|
||||
|
||||
return _remove_batch_dim(batched_output, vmap_level, batch_size, out_dim)
|
||||
|
||||
|
||||
# Undos the batching (and any batch dimensions) associated with the `vmap_level`.
|
||||
def _unwrap_batched(
|
||||
batched_outputs: Tensor | tuple[Tensor, ...],
|
||||
out_dims: out_dims_t,
|
||||
vmap_level: int,
|
||||
batch_size: int,
|
||||
func: Callable[..., Any],
|
||||
) -> tuple[Any, ...]:
|
||||
flat_batched_outputs, output_spec = tree_flatten(batched_outputs)
|
||||
|
||||
def incompatible_error() -> NoReturn:
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, ..., out_dims={out_dims})(<inputs>): "
|
||||
f"out_dims is not compatible with the structure of `outputs`. "
|
||||
f"out_dims has structure {tree_flatten(out_dims)[1]} but outputs "
|
||||
f"has structure {output_spec}."
|
||||
)
|
||||
|
||||
flat_out_dims: list[int | None] = []
|
||||
if isinstance(batched_outputs, torch.Tensor):
|
||||
# Some weird edge case requires us to spell out the following
|
||||
# see test_out_dims_edge_case
|
||||
if isinstance(out_dims, int):
|
||||
flat_out_dims = [out_dims]
|
||||
elif isinstance(out_dims, tuple) and len(out_dims) == 1:
|
||||
flat_out_dims = list(out_dims)
|
||||
elif out_dims is None:
|
||||
flat_out_dims = [out_dims]
|
||||
else:
|
||||
incompatible_error()
|
||||
else:
|
||||
broadcast_result = _broadcast_to_and_flatten(out_dims, output_spec)
|
||||
if broadcast_result is None:
|
||||
incompatible_error()
|
||||
else:
|
||||
flat_out_dims = broadcast_result
|
||||
|
||||
flat_outputs = [
|
||||
_maybe_remove_batch_dim(
|
||||
_get_name(func), batched_output, vmap_level, batch_size, out_dim
|
||||
)
|
||||
for batched_output, out_dim in zip(flat_batched_outputs, flat_out_dims)
|
||||
]
|
||||
return tree_unflatten(flat_outputs, output_spec)
|
||||
|
||||
|
||||
def _check_int_or_none(x: Any, func: Callable[..., Any], out_dims: out_dims_t) -> None:
|
||||
if isinstance(x, int):
|
||||
return
|
||||
if x is None:
|
||||
return
|
||||
raise ValueError(
|
||||
f"vmap({_get_name(func)}, ..., out_dims={out_dims}): `out_dims` must be "
|
||||
f"an int, None or a python collection of ints representing where in the outputs the "
|
||||
f"vmapped dimension should appear."
|
||||
)
|
||||
|
||||
|
||||
def _check_out_dims_is_int_or_int_pytree(
|
||||
out_dims: out_dims_t, func: Callable[..., Any]
|
||||
) -> None:
|
||||
if isinstance(out_dims, int):
|
||||
return
|
||||
tree_map_(partial(_check_int_or_none, func=func, out_dims=out_dims), out_dims)
|
||||
|
||||
|
||||
def _get_name(func: Callable[..., Any]) -> str:
|
||||
if hasattr(func, "__name__"):
|
||||
return func.__name__
|
||||
|
||||
if isinstance(func, functools.partial):
|
||||
return f"functools.partial({_get_name(func.func)}, ...)"
|
||||
|
||||
# Not all callables have __name__, in fact, only static functions/methods
|
||||
# do. A callable created via nn.Module, to name one example, doesn't have a
|
||||
# __name__.
|
||||
return repr(func)
|
||||
|
||||
|
||||
def vmap_impl(
|
||||
func: Callable[_P, Tensor | tuple[Tensor, ...]],
|
||||
in_dims: in_dims_t,
|
||||
out_dims: out_dims_t,
|
||||
randomness: str,
|
||||
chunk_size: int | None,
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs,
|
||||
) -> Any:
|
||||
lazy_load_decompositions()
|
||||
_check_out_dims_is_int_or_int_pytree(out_dims, func)
|
||||
batch_size, flat_in_dims, flat_args, args_spec = _process_batched_inputs(
|
||||
in_dims, args, func
|
||||
)
|
||||
|
||||
if chunk_size is not None:
|
||||
chunks_flat_args = _get_chunked_inputs(
|
||||
flat_args, flat_in_dims, batch_size, chunk_size
|
||||
)
|
||||
return _chunked_vmap(
|
||||
func,
|
||||
flat_in_dims,
|
||||
chunks_flat_args,
|
||||
args_spec,
|
||||
out_dims,
|
||||
randomness,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# If chunk_size is not specified.
|
||||
return _flat_vmap(
|
||||
func,
|
||||
batch_size,
|
||||
flat_in_dims,
|
||||
flat_args,
|
||||
args_spec,
|
||||
out_dims,
|
||||
randomness,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def get_chunk_sizes(total_elems: int, chunk_size: int) -> list[int]:
|
||||
n_chunks = total_elems // chunk_size
|
||||
chunk_sizes = [chunk_size] * n_chunks
|
||||
# remainder chunk
|
||||
remainder = total_elems % chunk_size
|
||||
if remainder != 0:
|
||||
chunk_sizes.append(remainder)
|
||||
return chunk_sizes
|
||||
|
||||
|
||||
def _get_chunked_inputs(
|
||||
flat_args: list[Any],
|
||||
flat_in_dims: list[int | None],
|
||||
batch_size: int,
|
||||
chunk_size: int | None,
|
||||
) -> Iterable[tuple[Any, ...]]:
|
||||
split_idxs = (batch_size,)
|
||||
if chunk_size is not None:
|
||||
chunk_sizes = get_chunk_sizes(batch_size, chunk_size)
|
||||
split_idxs = tuple(itertools.accumulate(chunk_sizes))
|
||||
|
||||
flat_args_chunks = tuple(
|
||||
(
|
||||
t.tensor_split(split_idxs, dim=in_dim)
|
||||
if in_dim is not None
|
||||
else [
|
||||
t,
|
||||
]
|
||||
* len(split_idxs)
|
||||
)
|
||||
for t, in_dim in zip(flat_args, flat_in_dims)
|
||||
)
|
||||
|
||||
# transpose chunk dim and flatten structure
|
||||
# chunks_flat_args is a list of flatten args
|
||||
chunks_flat_args = zip(*flat_args_chunks)
|
||||
return chunks_flat_args
|
||||
|
||||
|
||||
def _flatten_chunks_output(
|
||||
chunks_output_: list[Any],
|
||||
) -> tuple[list[tuple[Any, ...]], TreeSpec]:
|
||||
# chunks_output is a list of chunked outputs
|
||||
# flatten chunked outputs:
|
||||
flat_chunks_output: list[list[Any]] = []
|
||||
arg_spec: TreeSpec | None = None
|
||||
for output in chunks_output_:
|
||||
flat_output, arg_specs = tree_flatten(output)
|
||||
flat_chunks_output.append(flat_output)
|
||||
if arg_spec is None:
|
||||
arg_spec = arg_specs
|
||||
|
||||
# transpose chunk dim and flatten structure
|
||||
# flat_output_chunks is flat list of chunks
|
||||
flat_output_chunks = list(zip(*flat_chunks_output))
|
||||
if arg_spec is None:
|
||||
raise AssertionError("arg_spec must not be None")
|
||||
return flat_output_chunks, arg_spec
|
||||
|
||||
|
||||
def _concat_chunked_outputs(
|
||||
out_dims: out_dims_t,
|
||||
arg_spec: TreeSpec,
|
||||
flat_output_chunks: list[tuple[Any, ...] | None],
|
||||
) -> list[Tensor]:
|
||||
# concat chunks on out_dim
|
||||
flat_out_dims = _broadcast_to_and_flatten(out_dims, arg_spec)
|
||||
if flat_out_dims is None:
|
||||
raise AssertionError("flat_out_dims must not be None")
|
||||
if len(flat_out_dims) != len(flat_output_chunks):
|
||||
raise AssertionError(
|
||||
f"len(flat_out_dims)={len(flat_out_dims)} != len(flat_output_chunks)={len(flat_output_chunks)}"
|
||||
)
|
||||
flat_output: list[Tensor] = []
|
||||
for idx, out_dim in enumerate(flat_out_dims):
|
||||
chunk = flat_output_chunks[idx]
|
||||
if chunk is None:
|
||||
raise AssertionError(f"chunk at index {idx} must not be None")
|
||||
flat_output.append(torch.cat(chunk, dim=out_dim))
|
||||
# release tensors
|
||||
flat_output_chunks[idx] = None
|
||||
|
||||
return flat_output
|
||||
|
||||
|
||||
# Applies vmap on chunked_input and returns concatenated output over the chunks.
|
||||
def _chunked_vmap(
|
||||
func: Callable[_P, Tensor | tuple[Tensor, ...]],
|
||||
flat_in_dims: list[int | None],
|
||||
chunks_flat_args: Iterable[tuple[Any, ...]],
|
||||
args_spec: TreeSpec,
|
||||
out_dims: out_dims_t,
|
||||
randomness: str,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
chunks_output: list[Any] = []
|
||||
rs = torch.get_rng_state() if randomness == "same" else None
|
||||
for flat_args_tuple in chunks_flat_args:
|
||||
flat_args = list(flat_args_tuple)
|
||||
batch_size = _validate_and_get_batch_size(flat_in_dims, flat_args)
|
||||
|
||||
# The way we compute split the input in `_get_chunked_inputs`,
|
||||
# we may get a tensor with `0` batch-size. We skip any computation
|
||||
# in that case.
|
||||
# Eg.
|
||||
# >>> chunk_size = 1
|
||||
# >>> batch_size = 6
|
||||
# >>> t = torch.zeros(batch_size, 1)
|
||||
# >>> t.tensor_split([1, 2, 3, 4, 5, 6])
|
||||
# (tensor([[0.]]), tensor([[0.]]), tensor([[0.]]), tensor([[0.]]),
|
||||
# tensor([[0.]]), tensor([[0.]]), tensor([], size=(0, 1)))
|
||||
if batch_size == 0:
|
||||
continue
|
||||
|
||||
if rs is not None:
|
||||
torch.set_rng_state(rs)
|
||||
chunks_output.append(
|
||||
_flat_vmap(
|
||||
func,
|
||||
batch_size,
|
||||
flat_in_dims,
|
||||
flat_args,
|
||||
args_spec,
|
||||
out_dims,
|
||||
randomness,
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
flat_output_chunks, arg_spec = _flatten_chunks_output(chunks_output)
|
||||
|
||||
# chunked output tensors are held by both `flat_output_chunks` and `chunks_output`.
|
||||
# eagerly remove the reference from `chunks_output`.
|
||||
del chunks_output
|
||||
|
||||
# concat chunks on out_dim
|
||||
# Note: We use cast since flat_output_chunks is modified in _concat_chunked_outputs
|
||||
# to set elements to None after processing
|
||||
flat_output = _concat_chunked_outputs(
|
||||
out_dims, arg_spec, cast(list[tuple[Any, ...] | None], flat_output_chunks)
|
||||
)
|
||||
|
||||
# finally unflatten the output
|
||||
return tree_unflatten(flat_output, arg_spec)
|
||||
|
||||
|
||||
# Vmap refactored helper functions:
|
||||
def _check_randomness_arg(randomness: str) -> None:
|
||||
if randomness not in ["error", "different", "same"]:
|
||||
raise RuntimeError(
|
||||
f"Only allowed values for randomness are 'error', 'different', or 'same'. Got {randomness}"
|
||||
)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def vmap_increment_nesting(
|
||||
batch_size: int, randomness: str
|
||||
) -> Generator[int, None, None]:
|
||||
try:
|
||||
vmap_level = _vmap_increment_nesting(batch_size, randomness)
|
||||
yield vmap_level
|
||||
finally:
|
||||
_vmap_decrement_nesting()
|
||||
|
||||
|
||||
def _flat_vmap(
|
||||
func: Callable[..., Tensor | tuple[Tensor, ...]],
|
||||
batch_size: int,
|
||||
flat_in_dims: list[int | None],
|
||||
flat_args: list[Any],
|
||||
args_spec: TreeSpec,
|
||||
out_dims: out_dims_t,
|
||||
randomness: str,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
with vmap_increment_nesting(batch_size, randomness) as vmap_level:
|
||||
batched_inputs = _create_batched_inputs(
|
||||
flat_in_dims, flat_args, vmap_level, args_spec
|
||||
)
|
||||
batched_outputs = func(*batched_inputs, **kwargs)
|
||||
return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
|
||||
|
||||
|
||||
# `restore_vmap` is a private helper function. It is vmap but has the following
|
||||
# differences:
|
||||
# - instead of returning outputs, it returns an (outputs, out_dims) tuple.
|
||||
# out_dims is a pytree of same shape as outputs and contains Optional[int]
|
||||
# specifying where the vmapped dimension, if it exists, is in the corresponding output.
|
||||
# - does no validation on in_dims or inputs (vmap expects at least one Tensor to be vmapped).
|
||||
# restore_vmap allows for no inputs to have the vmap dimension
|
||||
# - does no validation on outputs (vmap expects only Tensor outputs)
|
||||
# restore_vmap allows for return of arbitrary outputs (not just Tensors)
|
||||
#
|
||||
# The TL;DR is that restore_vmap is more general than vmap and has a slightly
|
||||
# different API. The relaxations are so that we can "pause" vmap in the middle
|
||||
# of its execution and then "restore" it later (this is what we do in
|
||||
# the generate_vmap_rule=True implementation of autograd.Function).
|
||||
#
|
||||
# restore_vmap can be technically used in the implementation of vmap, but doing
|
||||
# that refactor is a bit technically challenging because:
|
||||
# - vmap couples the tensor-wrapping code with error checking
|
||||
# - vmap's tensor unwrapping code is in C++; we would need to rewrite part of it
|
||||
# in python because it overlaps with unwrap_batched
|
||||
def restore_vmap(
|
||||
func: Callable[..., _R], in_dims: in_dims_t, batch_size: int, randomness: str
|
||||
) -> Callable[..., tuple[Any, Any]]:
|
||||
def inner(*args: Any, **kwargs: Any) -> tuple[Any, Any]:
|
||||
with vmap_increment_nesting(batch_size, randomness) as vmap_level:
|
||||
batched_inputs = wrap_batched(args, in_dims, vmap_level)
|
||||
batched_outputs = func(*batched_inputs, **kwargs)
|
||||
return unwrap_batched(batched_outputs, vmap_level)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def wrap_batched(
|
||||
args: tuple[Any, ...], bdims: in_dims_t, level: int
|
||||
) -> tuple[Any, ...]:
|
||||
flat_args, spec = tree_flatten(args)
|
||||
flat_bdims = _broadcast_to_and_flatten(bdims, spec)
|
||||
if flat_bdims is None:
|
||||
raise AssertionError("flat_bdims must not be None")
|
||||
result = _create_batched_inputs(flat_bdims, flat_args, level, spec)
|
||||
return result
|
||||
|
||||
|
||||
def unwrap_batched(args: Any, level: int) -> tuple[Any, Any]:
|
||||
flat_args, spec = tree_flatten(args)
|
||||
if len(flat_args) == 0:
|
||||
return args, ()
|
||||
result = [
|
||||
(
|
||||
torch._C._functorch._unwrap_batched(arg, level)
|
||||
if isinstance(arg, torch.Tensor)
|
||||
else (arg, None)
|
||||
)
|
||||
for arg in flat_args
|
||||
]
|
||||
output, bdims = zip(*result)
|
||||
return tree_unflatten(output, spec), tree_unflatten(bdims, spec)
|
||||
Reference in New Issue
Block a user