import operator
from collections import deque
from typing import Dict, List, Set, NamedTuple, Tuple, Deque
import torch
from torch.fx.passes.graph_manipulation import get_size_of_all_nodes
from torch.fx.experimental.partitioner_utils import (
Partition,
Device,
PartitionerConfig,
get_partition_to_latency_mapping,
get_latency_of_partitioned_graph,
NodeLatency,
get_extra_size_of,
PartitionMode,
)
from torch.fx.graph_module import GraphModule
from torch.fx.node import Node, map_arg
from torch.fx.passes.split_module import split_module
class DAGNode:
"""DAGNode class maintains useful information for a partition (submodule),
and its input submodules and output submodules.
"""
def __init__(
self,
submodule_node: Node,
input_nodes: List[Node],
output_nodes: List[Node],
logical_device_ids: List[int],
size_bytes: int,
) -> None:
self.submodule_node: Node = submodule_node
self.input_nodes: List[Node] = input_nodes
self.output_nodes: List[Node] = output_nodes
self.logical_device_ids: List[int] = logical_device_ids
self.size_bytes = size_bytes
def __str__(self) -> str:
return str(self.submodule_node)
class DAG:
"""DAG class contains all the DAG nodes"""
def __init__(self) -> None:
self.nodes: List[DAGNode] = []
def create_node(
self,
submodule_node: Node,
input_nodes: List[Node],
output_nodes: List[Node],
logical_devices: List[int],
size_bytes: int,
) -> None:
node = DAGNode(
submodule_node, input_nodes, output_nodes, logical_devices, size_bytes
)
self.nodes.append(node)
class PartitionResult(NamedTuple):
"""NameTuple used for returning DAG and a new fx module"""
dag: DAG
module_with_submodules: GraphModule
"""Followings are some helper functions for partition manipulation"""
def reset_partition_device(partitions):
for partition in partitions:
partition.logical_device_ids = []
def combine_two_partitions(
partition_0: Partition, partition_1: Partition, partitions: List[Partition]
) -> None:
"""Given a list of partitions and its two partitions,
combine these two partitions into a new one appending to the partitions
and remove the previous two partitions from the list of partitions
"""
partition = Partition(len(partitions))
partition.nodes = partition_0.nodes.union(partition_1.nodes)
partition.recalculate_mem_size()
partitions.append(partition)
partitions.remove(partition_0)
partitions.remove(partition_1)
reorganize_partitions(partitions)
return
def set_parents_and_children(partitions: List[Partition]) -> None:
"""Given a list of partitions, mark parents and children for each partition"""
# Go through all nodes in a partition.
# If a node's user is in other partition,
# then the other partition is this partition's children.
# This partition is the other partition's parent
for partition in partitions:
partition.children = set()
partition.parents = set()
for partition in partitions:
for node in partition.nodes:
# For each node in the current partition, find its users
users = node.users
for n in users:
# Find which the partition the user node belongs to.
# Note that if the node itself is also belongs to that partition,
# that partition is not the child of the current partition
for p in partitions:
if p != partition and n in p.nodes and node not in p.nodes:
partition.children.add(p)
p.parents.add(partition)
return
def reorganize_partitions(partitions: List[Partition]) -> None:
"""Given a list of partitions, reorganize partition id,
its parents and its children for each partition
"""
# Rearrange partition ids
for i, partition in enumerate(partitions):
partition.partition_id = i
set_parents_and_children(partitions)
return
def get_bfs_level_partition(partitions: List[Partition]) -> None:
"""Given a list of partitions,
mark the bfs level for each partition
"""
current_level: Set[Partition] = set()
visited: Set[Partition] = set()
for partition in partitions:
# If a partition has no parent, it should be in root level
if len(partition.parents) == 0:
current_level.add(partition)
next_level: Set[Partition] = set()
level = 0
# bfs
while current_level:
partition = current_level.pop()
partition.bfs_level = level
visited.add(partition)
children = partition.children
for child in children:
if child not in next_level:
next_level.add(child)
if not current_level:
current_level = next_level.copy()
next_level = set()
level += 1
return
def get_node_to_partition_mapping(partitions: List[Partition]) -> Dict[Node, int]:
"""Given a list of partitions,return node to partition mapping"""
node_to_partition: Dict[Node, int] = {}
for partition in partitions:
for node in partition.nodes:
node_to_partition[node] = partition.partition_id
return node_to_partition
def get_logical_id_to_device(devices: List[Device]) -> Dict[int, Device]:
"""Get a mapping from device logical ID to Device object."""
logical_id_to_device: Dict[int, Device] = {}
for d in devices:
logical_id_to_device[d.logical_id] = d
return logical_id_to_device
def get_device_partition_stats(
partitions: List[Partition], devices: List[Device]
) -> Tuple[Dict[Device, List[Partition]], Dict[Device, int], List[Partition]]:
"""Given a list of partitions and a list of devices, returns:
1. A mapping from device to partitions on it;
2. A mapping from device to its remaining memory size;
3. A list of partitions that do not have a device.
"""
# logical id to device
logical_id_to_device = get_logical_id_to_device(devices)
# Track partitions on device
device_to_partitions: Dict[Device, List[Partition]] = {}
# Track device's left mem size
device_to_left_mem_bytes: Dict[Device, int] = {}
for d in devices:
device_to_partitions[d] = []
device_to_left_mem_bytes[d] = d.available_mem_bytes
# Deal with the partitions that already have a device
# and also collect all partitions without a device (no_device_partitions)
no_device_partitions = []
for partition in partitions:
if partition.logical_device_ids != []:
for logical_id in partition.logical_device_ids:
device = logical_id_to_device[logical_id]
device_to_partitions[device].append(partition)
device_to_left_mem_bytes[device] -= partition.used_mem_bytes
else:
no_device_partitions.append(partition)
return (
device_to_partitions,
device_to_left_mem_bytes,
no_device_partitions,
)
def get_device_to_partitions_mapping(
partitions: List[Partition], devices: List[Device]
):
"""Given a list of partitions and a list of devices,
map each partition into a device.
"""
def calculate_extra_mem_bytes_needed_for(
partition: Partition, partitions: List[Partition]
):
all_nodes: Set[Node] = set()
for p in partitions:
all_nodes = all_nodes.union(p.nodes)
if len(all_nodes) == 0:
return partition.used_mem_bytes
all_nodes = all_nodes.union(partition.nodes)
extra_size_needed = 0
for node in partition.nodes:
extra_size_needed += get_extra_size_of(node, all_nodes)
return extra_size_needed
def find_device_for(partition: Partition):
"""Given a partition, find a logical device for the partition
The algorithm is to put the partition on the device
that has just enough mem left for that partition.
device_to_left_mem_bytes is a dictionary between device and its left mem size
sorted by its left mem size
"""
for d in device_to_left_mem_bytes:
extra_size_needed = calculate_extra_mem_bytes_needed_for(
partition, device_to_partitions[d]
)
if extra_size_needed < device_to_left_mem_bytes[d]:
device_to_partitions[d].append(partition)
partition.logical_device_ids.append(d.logical_id)
device_to_left_mem_bytes[d] -= extra_size_needed
return True
return False
(
device_to_partitions,
device_to_left_mem_bytes,
no_device_partitions,
) = get_device_partition_stats(partitions, devices)
# Find devices for all the partitions without a device
found_device = True
for partition in no_device_partitions:
device_to_left_mem_bytes = dict(sorted(device_to_left_mem_bytes.items(), key=lambda item: item[1]))
found_device = find_device_for(partition)
if not found_device:
break
return found_device
def check_dependency(partition):
"""Given a partition,check if there is a circular dependency on
this partition using bfs
"""
visited: Set[Partition] = {partition}
queue: Deque[Partition] = deque([partition])
while queue:
p = queue.popleft()
for child in p.children:
if child == partition:
return True
else:
if child not in visited:
visited.add(child)
queue.append(child)
return False
class Partitioner:
"""A fx module may not fit into one device.
Partitioner class helps partition one fx module into submodules (partitions),
so that the submodules can be executed crossing different accelerators.
The main function of this class is self.partition_graph.
It partitions the fx module based on the scheme specified in partition_config
A DAG structure is returned
along with a new fx module with submodule nodes.
"""
def __init__(self) -> None:
self.partitions: List[Partition] = []
self.node_to_partition: Dict[Node, int] = {}
self.devices: List[Device] = []
def partition_graph(
self,
fx_module: GraphModule,
torch_module: torch.nn.Module,
partitioner_config: PartitionerConfig,
) -> PartitionResult:
"""Given the fx module, torch module and partitioner_config,
find the partitions, do the partitions,
and then return a DAG and a new fx module with submodule nodes (partitions)
"""
self.graph_module = fx_module
self.torch_module = torch_module
self.devices = partitioner_config.devices
if len(self.devices) == 0:
raise RuntimeError("No devices")
# Tag the size in bytes to all nodes in the graph_module.
get_size_of_all_nodes(self.graph_module)
# Check if there are op nodes in the fx module
nodes = self.graph_module.graph.nodes
if all(node.op in {"placeholder", "get_attr", "output"} for node in nodes):
raise RuntimeError("No Partition since no operations in the module")
# Calculate total size of the fx module
total_size_of_graph = 0
for node in nodes:
if node.op == "output":
break
total_size_of_graph += node.size_bytes.total_size
# Find the device with the max mem size
device_with_max_mem = max(self.devices, key=lambda d: d.available_mem_bytes)
# AOT based partition
if partitioner_config.mode == PartitionMode.aot_based:
self.aot_based_partition(
partitioner_config.node_to_partition_mapping,
partitioner_config.partition_to_logical_device_mapping,
)
# Single partition if the whole module can be fit into one device
elif total_size_of_graph <= device_with_max_mem.available_mem_bytes:
self.find_single_partition(
total_size_of_graph, logical_device_id=device_with_max_mem.logical_id
)
elif total_size_of_graph > sum([d.available_mem_bytes for d in self.devices]):
raise RuntimeError("Devices have no enough memory for the module")
else:
# Sparse nn based partition
if partitioner_config.mode == PartitionMode.sparse_nn:
available_mem_bytes = self.devices[0].available_mem_bytes
if not all(
device.available_mem_bytes == available_mem_bytes
for device in self.devices
):
raise RuntimeError("All devices must have same memory size!")
# sparse_nn_partition only support same memory size
# TODO: add different size support for sparse_nn_partition
self.sparse_nn_partition(available_mem_bytes)
# Cost aware partition
elif partitioner_config.mode == PartitionMode.cost_aware:
self.cost_aware_partition(
partitioner_config.transfer_rate_bytes_per_sec,
partitioner_config.node_to_latency_mapping,
)
# KL based partition
elif partitioner_config.mode == PartitionMode.kl_based:
self.kl_based_partition(
partitioner_config.transfer_rate_bytes_per_sec,
partitioner_config.node_to_latency_mapping,
)
else:
self.size_based_partition()
# Saturate host if possible.
if partitioner_config.saturate_host:
self.saturate_host()
# Partition the graph module based on the partition assignment.
module_with_submodules = self.do_partition()
# The DAG contains DAGNodes with info of each partition's input nodes, output nodes
# and how partitions are connected.
dag = self.dump_dag(module_with_submodules)
ret = PartitionResult(dag, module_with_submodules)
return ret
def find_single_partition(
self, total_size_of_graph, logical_device_id: int = 0
) -> None:
"""Fit the whole fx module into one device"""
partition_0 = self.create_partition()
for node in self.graph_module.graph.nodes:
if node.op == "output":
# Skip the output node, but there can
# be nodes after the output in certain cases.
continue
partition_0.nodes.add(node)
partition_0.used_mem_bytes = total_size_of_graph
partition_0.logical_device_ids = [logical_device_id]
# Get the node to partition mapping
self.node_to_partition = get_node_to_partition_mapping(self.partitions)
return
def size_based_partition(self) -> None:
"""This method is to partition the fx module based on memory size.
It uses greedy approach. The result may not be the best.
The basic idea is:
Step 1:
Find a device which has enough memory to fit the current node, create a empty partition
with the size of that device.
Then keep adding the following nodes into the partition until the partition is full.
Step 2:
Repeat Step 1 until no device left
Step 3:
If some nodes are left, create a partition for each left node (single node partition).
and then try to map those partitions into logical devices with enough mem left.
"""
def find_device_based_on_size(node) -> Device:
"""Given a node, this function is to find a logical device
that could fit the node.
"""
mem_size_needed = get_extra_size_of(node, set())
device = Device("", -1, -1)
for d in self.devices:
if (
d not in occupied_devices
and d.available_mem_bytes >= mem_size_needed
):
device = d
break
if device.available_mem_bytes < 0:
raise RuntimeError(str(node) + "is too large to fit any device")
occupied_devices.append(device)
return device
# Track partition and its left mem size
partition_to_left_mem_bytes: Dict[Partition, int] = {}
# Track all the devices that have been used
occupied_devices: List[Device] = []
partition = self.create_partition()
for node in self.graph_module.graph.nodes:
if node.op in {"call_module", "call_method", "call_function"}:
# Check if there are devices left
if len(self.partitions) <= len(self.devices):
total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
# Check if the current partition is the very first partition
if partition.used_mem_bytes == 0:
# Find a device to fit the first node, return available mem size
device = find_device_based_on_size(node)
occupied_devices.append(device)
# Update partition and its left mem size
partition_to_left_mem_bytes[
partition
] = device.available_mem_bytes
# Update available mem for the current partition
partition.logical_device_ids.append(device.logical_id)
else:
# The current partition is not the first partition
# Check if the current node can fit into current partition
if (
partition_to_left_mem_bytes[partition]
< total_size_of_input_nodes
):
# Check if no device is left
if len(self.partitions) == len(self.devices):
# No device is left
# Put the previous partitions into a list (non_single_node_partitions)
non_single_node_partitions = self.partitions[:]
# Create the first single node partition for the current node
self.create_single_node_partition(node)
continue
# Some devices are still left
# Create a new partition with a mem size that is enough for the current node
device = find_device_based_on_size(node)
partition = self.create_partition()
total_size_of_input_nodes = get_extra_size_of(
node, partition.nodes
)
partition_to_left_mem_bytes[
partition
] = device.available_mem_bytes
partition.logical_device_ids.append(device.logical_id)
partition.add_node(node)
partition_to_left_mem_bytes[partition] -= total_size_of_input_nodes
# Create single node partitions if no device is left
else:
self.create_single_node_partition(node)
reorganize_partitions(self.partitions)
# Get the node to partition mapping
self.node_to_partition = get_node_to_partition_mapping(self.partitions)
# Mapping all partitions into device
found_partition_to_device_mapping = get_device_to_partitions_mapping(
self.partitions, self.devices
)
if not found_partition_to_device_mapping:
raise RuntimeError("Cannot Get a Valid Partition to Logical Device Mapping")
return
def saturate_host(self) -> None:
"""Saturate host by assigning replicates to unused devices with enough memory.
It uses a greedy approach to find a next available set of devices to place all split
partitions: For each used device, it searches for an idle device with minimal memory
size that can hold all the partition located on that device; If the search is successful
for all used devices, it then assigns the new devices' logical ID to the corresponding
partition.
"""
(
device_to_partitions,
device_to_left_mem_bytes,
no_device_partitions,
) = get_device_partition_stats(self.partitions, self.devices)
assert (
len(no_device_partitions) == 0
), f"Expect no_device_partitions has 0 device, but get {len(no_device_partitions)}"
# Devices that hold partitions
used_devices = [d for d in self.devices if len(device_to_partitions[d]) > 0]
# Track replicates of the assigned devices
replicated_device_to_used_device: Dict[Device, Device] = {}
while len(used_devices) * 2 + len(replicated_device_to_used_device) <= len(
self.devices
):
# Success flag for this round
success = True
# Devices that have not been assigned
idle_devices = [
d
for d in self.devices
if d not in used_devices and d not in replicated_device_to_used_device
]
# Temporary mapping from replicated device to original device
temp_replicate_mapping = {}
# Find a new device to replicate all partitions on an used device
for used_device in used_devices:
# Idle devices that have enough memory
available_devices = [
d
for d in idle_devices
if d.available_mem_bytes
>= used_device.available_mem_bytes
- device_to_left_mem_bytes[used_device]
]
if len(available_devices) == 0:
success = False
break
new_device = min(available_devices, key=lambda d: d.available_mem_bytes)
idle_devices.remove(new_device)
temp_replicate_mapping[new_device] = used_device
if not success:
break
replicated_device_to_used_device.update(temp_replicate_mapping)
# Update logical device IDs assigned to the partitions
for (
replicate_device,
original_device,
) in replicated_device_to_used_device.items():
logical_id = replicate_device.logical_id
for partition in device_to_partitions[original_device]:
partition.logical_device_ids.append(logical_id)
for p in self.partitions:
print(p.logical_device_ids)
def do_partition(self) -> GraphModule:
"""Return a new fx module with submodule nodes (partitions)."""
module_with_submodules = split_module(
self.graph_module,
self.torch_module,
lambda node: self.node_to_partition[node],
)
return module_with_submodules
def dump_dag(self, module_with_submodules: GraphModule) -> DAG:
"""Return the dag structure and the new fx module with submodules."""
dag = DAG()
for node in module_with_submodules.graph.nodes:
if node.op == "output":
break
if node.op in {"placeholder", "get_attr"}:
continue
if node.target == operator.__getitem__:
continue
input_nodes: Dict[Node, None] = {}
map_arg(node.args, lambda n: input_nodes.setdefault(n))
map_arg(node.kwargs, lambda n: input_nodes.setdefault(n))
# When a node has two or more output nodes,
# it outputs its result to 'getitem' nodes.
# Those 'getitem' nodes are the output node for this node.
# Otherwise, the output node is this node itself.
if len(node.users) > 1:
output_nodes = list(node.users)
else:
output_nodes = [node]
partition_id = int(node.name.rsplit("_", 1)[-1])
device_ids = self.partitions[partition_id].logical_device_ids
size_bytes = self.partitions[partition_id].used_mem_bytes
dag.create_node(
node, list(input_nodes), output_nodes, device_ids, size_bytes
)
return dag
def create_partition(self) -> Partition:
"""Create a partition and append it to self.partitions."""
partition_id = len(self.partitions)
partition = Partition(partition_id)
self.partitions.append(partition)
return partition
def create_single_node_partition(self, node):
"""Create a partition for a single node"""
partition = self.create_partition()
partition.add_node(node)
return
def sparse_nn_partition(self, available_mem_bytes: int) -> None:
"""This method partition a sparse nn module.
It is size based partition but different from size_based_partition,
it only works when all the devices have same memory size (available_mem_bytes).
In the future, devices with different mem sizes will be supported like size_based_partition.
It first traverse all the nodes and do the partitions based on the same memory size.
If the current partition has no enough memory left for a new op node
(call_module, call_method, call_function), a new partition is created.
When crossing the boundary between non-embedding nodes and embedding nodes,
a new partition is created regardlessly.
For example, if the current node is a non-embedding node but the next node is an
embedding node, a new partition is created for the next node.
After the partition, the partitions are combined as much as possible.
The rule is that a non-embedding partition only
combines with another non-embedding one.
So as the embedding partitions.
"""
def combine_partitions_based_on_size(
partitions: List[Partition], available_mem_bytes: int
) -> None:
"""Combining small partitions together to keep as less partitions as possible.
Here is an example of the algorithm to do this:
Assume some partitions, we first sort them based on partition used memory size.
[(partition_4, 1), (partition_3, 1), (partition_2, 2), (partition_1, 7), (partition_0, 9)]
The available memory is 10.
step 1: self.find_partition_to_combine_based_on_size()
First, mark bfs level for each partition
Second, look the smallest partition, partition_4: 10 - 1 = 9
It means any partition has a used memory equal or less than 9 could combine this partition
We go from the largest and selection partition_0.
Check the bfs level for two partitions, if the level difference is less than 2,
it can be combined.
step 2: repeat step 1 until no partitions can be combined
"""
find_combination = True
while find_combination:
# Sort partitions based on memory size
sorted_partitions = sorted(partitions, key=lambda p: p.used_mem_bytes)
# Mark bfs level
get_bfs_level_partition(self.partitions)
find_combination, partitions = find_partition_to_combine_based_on_size(
sorted_partitions, available_mem_bytes, partitions
)
return
def calculate_mem_bytes_needed(p1, p2):
"""Given two partitions, calculate how many mem bytes
are needed if two partitions are combined
"""
nodes = p1.nodes.union(p2.nodes)
mem_bytes_needed = 0
for node in nodes:
mem_bytes_needed += get_extra_size_of(node, nodes)
return mem_bytes_needed
def find_partition_to_combine_based_on_size(
sorted_partitions: List[Partition],
available_mem_bytes: int,
partitions: List[Partition],
) -> Tuple[bool, List[Partition]]:
"""step 1 in combine_partition_based_on_size()"""
find_combination = False
smallest_partition = sorted_partitions.pop(0)
for p in sorted_partitions[::-1]:
if abs(smallest_partition.bfs_level - p.bfs_level) <= 1:
# Calculate how many bytes needed if combined
mem_bytes_needed = calculate_mem_bytes_needed(p, smallest_partition)
if mem_bytes_needed <= available_mem_bytes:
combine_two_partitions(p, smallest_partition, self.partitions)
partitions.remove(smallest_partition)
partitions.remove(p)
partitions.append(self.partitions[-1])
find_combination = True
break
return find_combination, partitions
def reset_partition_in_sparse_nn(partition, new_partition=True):
"""If crossing the boundary between non-embedding nodes and
embedding nodes, create a new partition
"""
if in_embedding_region:
embedding_partitions.append(partition)
else:
non_embedding_partitions.append(partition)
if new_partition:
partition = self.create_partition()
partition.left_mem_bytes = available_mem_bytes
return partition
return None
def is_embedding_node(node: Node) -> bool:
"""Check if a node is an embedding node"""
if node.op == "call_module":
submodule = self.graph_module
for atom in str(node.target).split("."):
if not hasattr(submodule, atom):
raise RuntimeError(
f"Module {submodule} has no attribute {atom}"
)
submodule = getattr(submodule, atom)
if "Embedding" in str(submodule):
return True
return False
# Track embedding partitions and non-embedding partitions separately
embedding_partitions: List[Partition] = []
non_embedding_partitions: List[Partition] = []
# A Flag to check the boundary
in_embedding_region: bool = False
partition = self.create_partition()
for node in self.graph_module.graph.nodes:
if node.op in {"call_module", "call_method", "call_function"}:
# Check if crossing the boundary between embedding nodes and non embedding nodes
if is_embedding_node(node) != in_embedding_region:
# Crossing the boundary
# Check if the current partition is an empty partition
if partition.used_mem_bytes != 0:
# The current partition isn't an empty partition. Create a new one.
partition = reset_partition_in_sparse_nn(partition)
in_embedding_region = not in_embedding_region
total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
if (
total_size_of_input_nodes + partition.used_mem_bytes
> available_mem_bytes
):
partition = reset_partition_in_sparse_nn(partition)
total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
if total_size_of_input_nodes > available_mem_bytes:
raise RuntimeError(
node.target + "is too large to fit into a device"
)
partition.add_node(node)
reset_partition_in_sparse_nn(partition, new_partition=False)
# Set parents and children for partitions
set_parents_and_children(self.partitions)
# Combining non-embedding partitions
combine_partitions_based_on_size(non_embedding_partitions, available_mem_bytes)
# Combining embedding partitions
combine_partitions_based_on_size(embedding_partitions, available_mem_bytes)
total_size_of_non_embedding_partitions = 0
for partition in non_embedding_partitions:
total_size_of_non_embedding_partitions += partition.used_mem_bytes
# Check if devices are enough for all partitions
if len(embedding_partitions) > len(self.devices):
msg = (
"Need "
+ str(len(embedding_partitions))
+ " devices, but only "
+ str(len(self.devices))
+ " provided"
)
raise RuntimeError(msg)
occupied_devices = []
for i, partition in enumerate(embedding_partitions):
# Check if all non-embedding partitions can fit into embedding partition devices
if (
total_size_of_non_embedding_partitions + partition.used_mem_bytes
> available_mem_bytes
):
raise RuntimeError(
"partition_"
+ str(partition.partition_id)
+ "(embedding partition) and non embedding partitions can not fit into one device"
)
else:
# Add logical device to the partition
partition.logical_device_ids = [self.devices[i].logical_id]
occupied_devices.append(self.devices[i].logical_id)
# Add logical devices to the non_embedding_partitions
for partition in non_embedding_partitions:
partition.logical_device_ids = occupied_devices
# Get the node to partition mapping
self.node_to_partition = get_node_to_partition_mapping(self.partitions)
return
def cost_aware_partition(
self,
transfer_rate_bytes_per_sec: float,
node_to_latency_mapping: Dict[Node, NodeLatency],
) -> None:
"""This method is to partition the fx module based on the cost.
The cost is the total latency of running the whole fx module.
In partitioner_utils.py, the cost model is built.
The cost aware partition algorithm is:
#1. At every beginning, each node is a partition.
Then we map all the partitions to the devices
and calculate the cost
#2. Then try to pre-combine any two of the partitions if the two
partitions can be combined.
(the bfs level is less than 2 or two partitions are connected and
can find partition to device mapping)
See if any partition pair could reduce the current cost.
Choose the pair that shows the minimum cost and then combine them
#3. Repeat #2 until the cost cannot be reduced.
"""
def try_combining_partitions(p0_index, p1_index, partitions) -> float:
"""Given two partitions and a list of partitions, combine these two partitions
and see what is the cost of the modified partition list
"""
p0 = partitions[p0_index]
p1 = partitions[p1_index]
"""If two partitions' bfs level are less than 2 or two partitions are connected to each other,
then they can be combined
"""
if (
(abs(p0.bfs_level - p1.bfs_level) <= 1)
or (p0 in p1.parents)
or p0 in (p1.children)
):
combine_two_partitions(p0, p1, partitions)
# Check if a circular dependency exists after combining
if check_dependency(partitions[-1]):
return float("inf")
# Check if the modified partition list can be mapped to devices after combination
reset_partition_device(partitions)
found_deivce = get_device_to_partitions_mapping(
partitions, self.devices
)
if not found_deivce:
return float("inf")
# Calculate the new cost
partition_to_latency_mapping = get_partition_to_latency_mapping(
partitions, node_to_latency_mapping
)
cost = get_latency_of_partitioned_graph(
partitions,
partition_to_latency_mapping,
transfer_rate_bytes_per_sec,
)
return cost
# If two partition can not be combined, the cost is inf
return float("inf")
def search_combination(
transfer_rate_bytes_per_sec, node_to_latency_mapping
) -> bool:
"""Given transfer rate between partitions and each node's latency,
find two partitions to combine so the cost of the partitions can
be reduced.
The algorithm is :
1. Go through all the partition pairs and see
if any pair of partitions can be combined.
2. Calculate the cost after the combination.
3. Select the minimum cost and combine its corresponding partition pair.
"""
partition_to_latency_mapping = get_partition_to_latency_mapping(
self.partitions, node_to_latency_mapping
)
cost = get_latency_of_partitioned_graph(
self.partitions,
partition_to_latency_mapping,
transfer_rate_bytes_per_sec,
)
if len(self.partitions) == 1:
return False
partition_pair: List[int] = []
for i in range(len(self.partitions) - 1):
for j in range(i + 1, len(self.partitions)):
# Try to combine the partition pair
# and see the new cost after combination
new_cost = try_combining_partitions(i, j, self.partitions[:])
if new_cost <= cost:
partition_pair = [i, j]
cost = new_cost
reorganize_partitions(self.partitions)
# If a partition pair is found, combine them
if len(partition_pair) != 0:
p0 = self.partitions[partition_pair[0]]
p1 = self.partitions[partition_pair[1]]
combine_two_partitions(p0, p1, self.partitions)
get_bfs_level_partition(self.partitions)
reset_partition_device(self.partitions)
get_device_to_partitions_mapping(self.partitions, self.devices)
return len(partition_pair) != 0
for node in self.graph_module.graph.nodes:
if node.op not in {"placeholder", "get_attr", "output"}:
self.create_single_node_partition(node)
# Set up parent partitions and children partitions for each partition
set_parents_and_children(self.partitions)
# Get bfs level for each partition
get_bfs_level_partition(self.partitions)
find_combination = True
while find_combination:
# Search for a pair partition to generate the minimum new cost,
# then combine them
find_combination = search_combination(
transfer_rate_bytes_per_sec, node_to_latency_mapping
)
# Make sure all partitions are set up correctly
reorganize_partitions(self.partitions)
# Set up node to partition mapping
self.node_to_partition = get_node_to_partition_mapping(self.partitions)
return
def kl_based_partition(
self,
transfer_rate_bytes_per_sec: float,
node_to_latency_mapping: Dict[Node, NodeLatency],
) -> None:
"""This function is a cost aware partition based
on Kernighan-Lin algorithm.
First, the graph is partitioned using size_based_partition.
Then, each node is swapped with any other node in a different
partition, and at the same time, the cost is estimated after
the swapping.
For example, we have nodes n0, n1, n2, n3 and n4.
Using size_based_partition, n0 and n1 are in Partition p0.
n2, n3 and n4 in Partition p1. The current cost is estimated.
We first tried using n0 to swap with n2 from the other partition.
Then we see that swapping n0 and n2 shows a lower cost
than the current cost and it is the minimum among other pairs like
(n0, None)(This means moving n0 to Partition without swapping other nodes),
(n0, n3) and (n0, n4). We swap n0 and n2 and set the new cost
as the current cost.
Then We repeat this process for all the other nodes until all swapping pairs
are tried.
"""
def swap_nodes(n0, n1, p0, p1):
# Either n0 or n1 could be None
# That means we simply move the node
# to another partition
if n0 is not None:
p0.remove_node(n0)
p1.add_node(n0)
if n1 is not None:
p0.add_node(n1)
p1.remove_node(n1)
def try_swap_nodes(
n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
):
cost = float("inf")
swap_nodes(n0, n1, p0, p1)
# Reorganize partitions after swapping
reorganize_partitions(self.partitions)
# Check if there is a circular dependency after swapping
if (not check_dependency(p0)) and (not check_dependency(p1)):
reset_partition_device(self.partitions)
partition_to_latency_mapping = get_partition_to_latency_mapping(
self.partitions, node_to_latency_mapping
)
# Check if all partitions can be mapped to logical devices after swapping
found_device = get_device_to_partitions_mapping(
self.partitions, self.devices
)
if not found_device:
cost = float("inf")
else:
cost = get_latency_of_partitioned_graph(
self.partitions,
partition_to_latency_mapping,
transfer_rate_bytes_per_sec,
)
# Swap back and reset all partitions back to original
swap_nodes(n1, n0, p0, p1)
reorganize_partitions(self.partitions)
reset_partition_device(self.partitions)
get_device_to_partitions_mapping(self.partitions, self.devices)
return cost
def swap_node_to_partition(
node, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
):
"""This function helps to swap one node from partition p0
with all the nodes in another partition p1
"""
p1_nodes = list(p1.nodes) + [None]
min_cost = float("inf")
node_pair: List[Node] = []
for n1 in p1_nodes:
# Ignore the node if it is not a op node
if n1 is not None and n1.op in {"placeholder", "get_attr"}:
continue
# Try swapping node in p0 with n1 in p1
cost = try_swap_nodes(
node, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
)
if cost < min_cost:
node_pair = [node, n1]
min_cost = cost
return cost, node_pair
# First use size_base_partition
self.size_based_partition()
partition_to_latency_mapping = get_partition_to_latency_mapping(
self.partitions, node_to_latency_mapping
)
# Calculate the cost of the partitions
cost = get_latency_of_partitioned_graph(
self.partitions, partition_to_latency_mapping, transfer_rate_bytes_per_sec
)
# Keep tracking the node pair that shows the better cost
node_pair: List[Node] = []
# Keep tracking the partition pair of node pair
partition_pair: List[Partition] = []
# Collect all the op nodes from the graph
op_nodes = []
for n in self.graph_module.graph.nodes:
if n.op not in {"placeholder", "get_attr", "output"}:
op_nodes.append(n)
for node in op_nodes:
# Find which partition the current node belongs
p0_index = self.node_to_partition[node]
p0 = self.partitions[p0_index]
# Go through all the other partitions to swap
# with other nodes from those partitions
for p1_index, _ in enumerate(self.partitions):
if p0_index != p1_index:
p1 = self.partitions[p1_index]
new_cost, new_node_pair = swap_node_to_partition(
node,
p0,
p1,
node_to_latency_mapping,
transfer_rate_bytes_per_sec,
)
# Update the cost
# Track the swapped node pair and their partitions
if new_cost < cost:
cost = new_cost
node_pair = new_node_pair
partition_pair = [p0, p1]
# Do the swapping after trying all the nodes from a partition
if len(node_pair) != 0:
swap_nodes(
node_pair[0], node_pair[1], partition_pair[0], partition_pair[1]
)
reorganize_partitions(self.partitions)
get_device_to_partitions_mapping(self.partitions, self.devices)
reorganize_partitions(self.partitions)
# Mapping the device to the partition
get_device_to_partitions_mapping(self.partitions, self.devices)
return
def aot_based_partition(
self, node_to_partition_mapping, partition_to_logical_device_mapping
):
"""This function helps to rebuild the partitions given the nodes and its
corresponding partition id
"""
partition_id_to_partition_mapping: Dict[int, Partition] = {}
self.node_to_partition = node_to_partition_mapping
for node in self.node_to_partition:
partition_id = self.node_to_partition[node]
# If the requested partition has not been created, create the partition
if partition_id not in partition_id_to_partition_mapping:
partition = Partition(partition_id)
self.partitions.append(partition)
partition_id_to_partition_mapping[partition_id] = partition
partition.logical_device_ids = partition_to_logical_device_mapping[
partition_id
]
else:
partition = partition_id_to_partition_mapping[
self.node_to_partition[node]
]
# Add the current node into the partition
partition.add_node(node)