|
|
import math |
|
|
import re |
|
|
from dataclasses import dataclass |
|
|
from enum import Enum |
|
|
from typing import Dict, List, Tuple |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
from tensorrt_llm.network import Network |
|
|
|
|
|
from .config import AutoParallelConfig |
|
|
from .device_mesh import PhysicalDeviceMesh |
|
|
from .pipeline_graph import PipelineGraph |
|
|
from .shape_info import ShapeInfo, ShapeType, get_shape_info |
|
|
from .tensor_parallel.p2p_node import P2PType |
|
|
from .utils import get_cache_key, get_sorted_layer_ids, silent_trt_logger |
|
|
|
|
|
|
|
|
class StageType(Enum): |
|
|
START = 0 |
|
|
BLOCK = 1 |
|
|
END = 2 |
|
|
|
|
|
|
|
|
class BuildingBlock: |
|
|
|
|
|
def __init__(self, graph, layer_range) -> None: |
|
|
self.graph = graph |
|
|
self.layer_range = layer_range |
|
|
self.network = graph.as_trt() |
|
|
self.owned_inputs = {} |
|
|
self.is_edges_collected = False |
|
|
self.intra_edges = [] |
|
|
self.src_inter_edges = [] |
|
|
self.dst_inter_edges = [] |
|
|
self.relative_src_inter_edges = [] |
|
|
self.relative_dst_inter_edges = [] |
|
|
self.relative_inter_edges = set() |
|
|
self.edge_hash = None |
|
|
self.outputs = None |
|
|
self.type_id = -1 |
|
|
self.block_id = -1 |
|
|
self.p2p_type = None |
|
|
self.is_superset = False |
|
|
self.is_subset = False |
|
|
self.sorted_layer_ids = [] |
|
|
|
|
|
def collect_edges(self): |
|
|
if self.is_edges_collected: |
|
|
return |
|
|
for layer_index in self.layer_range: |
|
|
trt_layer = self.network.get_layer(layer_index) |
|
|
layer = self.graph.get_layer(trt_layer.name) |
|
|
layer_offset = layer.index - self.layer_range.start |
|
|
for input_index, input in enumerate(layer.inputs): |
|
|
if input is not None: |
|
|
if input.is_graph_input: |
|
|
is_owned = input.graph_input_index in self.owned_inputs |
|
|
if not is_owned and np.all([ |
|
|
layer.index in self.layer_range or np.all([ |
|
|
output.as_trt().is_shape_tensor |
|
|
for output in layer.outputs |
|
|
]) for layer, _ in input.consumers |
|
|
]): |
|
|
self.owned_inputs[input.graph_input_index] = len( |
|
|
self.owned_inputs) |
|
|
is_owned = True |
|
|
if is_owned: |
|
|
self.intra_edges.append( |
|
|
(-1, self.owned_inputs[input.graph_input_index], |
|
|
layer_offset, input_index)) |
|
|
else: |
|
|
self.dst_inter_edges.append( |
|
|
(-1, input.graph_input_index, layer_offset, |
|
|
input_index)) |
|
|
else: |
|
|
src_layer_index = input.producer.index |
|
|
if src_layer_index < self.layer_range.start or src_layer_index >= self.layer_range.stop: |
|
|
self.dst_inter_edges.append( |
|
|
(src_layer_index, input.output_index, |
|
|
layer_offset, input_index)) |
|
|
else: |
|
|
src_layer_offset = src_layer_index - self.layer_range.start |
|
|
self.intra_edges.append( |
|
|
(src_layer_offset, input.output_index, |
|
|
layer_offset, input_index)) |
|
|
for output_index, output in enumerate(layer.outputs): |
|
|
for dst_layer, dst_input_index in output.consumers: |
|
|
dst_layer_index = dst_layer.index |
|
|
if dst_layer_index < self.layer_range.start or dst_layer_index >= self.layer_range.stop: |
|
|
self.src_inter_edges.append( |
|
|
(layer_offset, output_index, dst_layer_index, |
|
|
dst_input_index)) |
|
|
self.edge_hash = tuple(self.intra_edges) |
|
|
self.outputs = sorted( |
|
|
set((edge[0], edge[1]) for edge in self.src_inter_edges)) |
|
|
self.is_edges_collected = True |
|
|
|
|
|
def collect_relative_inter_edges(self, layer_to_block): |
|
|
self.collect_edges() |
|
|
for src_layer_index, src_output_index, dst_layer_index, dst_input_index in self.dst_inter_edges: |
|
|
if src_layer_index in layer_to_block: |
|
|
src_block = layer_to_block[src_layer_index] |
|
|
src_layer_offset = src_layer_index - src_block.layer_range.start |
|
|
dst = (self.type_id, dst_layer_index, dst_input_index) |
|
|
self.relative_dst_inter_edges.append( |
|
|
(src_block.type_id, src_layer_offset, src_output_index, |
|
|
*dst)) |
|
|
else: |
|
|
self.relative_dst_inter_edges.append( |
|
|
(-1, src_layer_index, src_output_index, self.type_id, |
|
|
dst_layer_index, dst_input_index)) |
|
|
self.relative_inter_edges = set(self.relative_dst_inter_edges + |
|
|
self.outputs) |
|
|
|
|
|
def get_input_names(self): |
|
|
self.collect_edges() |
|
|
input_tensor_names = [] |
|
|
for edge in self.dst_inter_edges: |
|
|
layer_index = edge[0] |
|
|
output_index = edge[1] |
|
|
if layer_index == -1: |
|
|
tensor_name = self.network.get_input(output_index).name |
|
|
else: |
|
|
tensor_name = self.network.get_layer(layer_index).get_output( |
|
|
output_index).name |
|
|
input_tensor_names.append(tensor_name) |
|
|
return input_tensor_names |
|
|
|
|
|
def get_input_mapping(self, last_blocks): |
|
|
input_mapping = {} |
|
|
for tensor_name, relative_edge in zip(self.get_input_names(), |
|
|
self.relative_dst_inter_edges): |
|
|
type_id = relative_edge[0] |
|
|
output_index = relative_edge[2] |
|
|
if type_id >= 0: |
|
|
last_block = last_blocks[type_id] |
|
|
layer_offset = relative_edge[1] |
|
|
mapped_layer_index = last_block.layer_range.start + layer_offset |
|
|
mapped_tensor_name = self.network.get_layer( |
|
|
mapped_layer_index).get_output(output_index).name |
|
|
input_mapping[tensor_name] = mapped_tensor_name |
|
|
else: |
|
|
input_mapping[tensor_name] = tensor_name |
|
|
return input_mapping |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class GraphMapping: |
|
|
layer_mapping: Dict[int, int] = None |
|
|
block_mapping: Dict[int, int] = None |
|
|
p2p_types: Dict[int, P2PType] = None |
|
|
p2p_tensors: Dict[int, List[str]] = None |
|
|
block_to_stage: Dict[int, int] = None |
|
|
same_spec_layer_mapping: Dict[str, str] = None |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class GraphConfig: |
|
|
num_micro_batches: int = 1 |
|
|
num_blocks: int = 1 |
|
|
num_stages: int = 1 |
|
|
has_cross_device: bool = False |
|
|
has_cross_host: bool = False |
|
|
graph_mapping: GraphMapping = None |
|
|
phy_mesh: PhysicalDeviceMesh = None |
|
|
stage_phy_meshes: List[PhysicalDeviceMesh] = None |
|
|
|
|
|
|
|
|
class Simplifier: |
|
|
|
|
|
def __init__(self, network: Network, config: AutoParallelConfig): |
|
|
self.config = config |
|
|
self.sharded_io_allowlist = config.sharded_io_allowlist |
|
|
self.same_buffer_io = config.same_buffer_io |
|
|
self.same_spec_io = config.same_spec_io.copy() |
|
|
for key, value in self.same_buffer_io.items(): |
|
|
if key not in self.same_spec_io: |
|
|
self.same_spec_io[key] = value |
|
|
|
|
|
self.llm_network = network |
|
|
self.network = network.trt_network |
|
|
self.module_to_layer_range_map = network._module_call_stack.module_to_layer_range_map |
|
|
self.graph = self.get_graph() |
|
|
self.init_layer_hash() |
|
|
|
|
|
module_tree = self.get_module_tree() |
|
|
building_blocks = self.collect_building_blocks(module_tree) |
|
|
blocks_by_module_hash = self.get_blocks_by_module_hash(building_blocks) |
|
|
self.blocks_by_edge_hash = self.get_blocks_by_edge_hash( |
|
|
blocks_by_module_hash) |
|
|
self.layer_to_block = self.get_layer_to_block() |
|
|
self.blocks = self.get_all_blocks() |
|
|
self.backbone_blocks = self.get_backbone_blocks() |
|
|
self.graph_mapping_for_shape = self.get_graph_mapping_for_shape() |
|
|
self.graph_for_shape = self.create_simplified_graph_for_shape() |
|
|
self.shape_info = None |
|
|
self.num_micro_batches = None |
|
|
|
|
|
def infer_shapes(self, num_micro_batches): |
|
|
if self.num_micro_batches == num_micro_batches: |
|
|
return |
|
|
with silent_trt_logger(): |
|
|
self.shape_info = self.get_full_shape_info(num_micro_batches) |
|
|
self.graph.assign_shapes(self.shape_info) |
|
|
self.num_micro_batches = num_micro_batches |
|
|
|
|
|
def list_all_num_micro_batches(self): |
|
|
opt_batch_size = self.get_opt_batch_size() |
|
|
candidates = [] |
|
|
for num_micro_batches in range(1, self.get_opt_batch_size() + 1): |
|
|
if opt_batch_size % num_micro_batches == 0: |
|
|
candidates.append(num_micro_batches) |
|
|
return candidates |
|
|
|
|
|
def get_graph(self): |
|
|
graph = PipelineGraph.from_trt(self.network) |
|
|
graph._unfilled_weights = self.llm_network._unfilled_weights.copy() |
|
|
graph._io_buffer_mapping |
|
|
for input in graph.inputs: |
|
|
input_name = input.name |
|
|
for pattern, repl in self.same_buffer_io.items(): |
|
|
if re.match(pattern, input_name): |
|
|
output_name = re.sub(pattern, repl, input_name) |
|
|
output = graph.get_output(output_name) |
|
|
if output is not None: |
|
|
graph._io_buffer_mapping[output_name] = input_name |
|
|
return graph |
|
|
|
|
|
def get_opt_batch_size(self): |
|
|
input_tensors = self.llm_network._inputs |
|
|
num_profiles = len(list(input_tensors.values())[0].profiles) |
|
|
opt_batch_sizes = [] |
|
|
for i in range(num_profiles): |
|
|
for input_tensor in input_tensors.values(): |
|
|
shape_profile = input_tensor.profiles[i] |
|
|
opt_shape = shape_profile.opt |
|
|
for j in range(len(input_tensor.shape)): |
|
|
name = input_tensor.trt_tensor.get_dimension_name(j) |
|
|
if name == 'batch_size': |
|
|
opt_batch_sizes.append(opt_shape[j]) |
|
|
return min(opt_batch_sizes) |
|
|
|
|
|
def get_module_hash(self, layer_range): |
|
|
module_hash = () |
|
|
for i in layer_range: |
|
|
assert i < self.network.num_layers, f"layer index {i} in {layer_range} out of range of {self.network.num_layers}" |
|
|
layer_name = self.network.get_layer(i).name |
|
|
layer = self.graph.get_layer(layer_name) |
|
|
module_hash += (layer.attrs["hash"], ) |
|
|
return module_hash |
|
|
|
|
|
def get_network_hash(self) -> str: |
|
|
return str(self.get_module_hash(range(self.network.num_layers))) |
|
|
|
|
|
def collect_building_blocks(self, module_tree): |
|
|
building_blocks = {} |
|
|
queue = [] |
|
|
for tree in module_tree["children"].values(): |
|
|
queue.append(tree) |
|
|
while len(queue) > 0: |
|
|
while len(queue) > 0: |
|
|
tree = queue.pop(0) |
|
|
module_name = tree["name"] |
|
|
if module_name is None: |
|
|
for child in tree["children"].values(): |
|
|
queue.append(child) |
|
|
continue |
|
|
layer_range = self.module_to_layer_range_map[module_name] |
|
|
module_hash = self.get_module_hash(layer_range) |
|
|
if module_hash in building_blocks: |
|
|
building_blocks[module_hash].append(tree) |
|
|
else: |
|
|
building_blocks[module_hash] = [tree] |
|
|
for module_hash in [*building_blocks.keys()]: |
|
|
if len(building_blocks[module_hash]) == 1: |
|
|
tree = building_blocks[module_hash][0] |
|
|
for child in tree["children"].values(): |
|
|
queue.append(child) |
|
|
del building_blocks[module_hash] |
|
|
blocks_by_module_hash = { |
|
|
module_hash: [ |
|
|
BuildingBlock(self.graph, |
|
|
self.module_to_layer_range_map[tree["name"]]) |
|
|
for tree in trees |
|
|
] |
|
|
for module_hash, trees in building_blocks.items() |
|
|
} |
|
|
building_blocks = [] |
|
|
for block_list in blocks_by_module_hash.values(): |
|
|
for block in block_list: |
|
|
building_blocks.append(block) |
|
|
building_blocks = sorted(building_blocks, |
|
|
key=lambda x: x.layer_range.start) |
|
|
if len(building_blocks) >= 2: |
|
|
for block, next_block in zip(building_blocks[:-1], |
|
|
building_blocks[1:]): |
|
|
block.layer_range = range(block.layer_range.start, |
|
|
next_block.layer_range.start) |
|
|
return building_blocks |
|
|
|
|
|
def get_all_blocks(self): |
|
|
building_blocks = [] |
|
|
for block_list in self.blocks_by_edge_hash.values(): |
|
|
for block in block_list: |
|
|
building_blocks.append(block) |
|
|
building_blocks = sorted(building_blocks, |
|
|
key=lambda x: x.layer_range.start) |
|
|
all_blocks = [] |
|
|
current_layer_index = 0 |
|
|
block_id = 0 |
|
|
for block in building_blocks: |
|
|
assert current_layer_index <= block.layer_range.start |
|
|
if current_layer_index < block.layer_range.start: |
|
|
new_block = BuildingBlock( |
|
|
self.graph, |
|
|
range(current_layer_index, block.layer_range.start)) |
|
|
new_block.block_id = block_id |
|
|
block_id += 1 |
|
|
all_blocks.append(new_block) |
|
|
block.block_id = block_id |
|
|
block_id += 1 |
|
|
all_blocks.append(block) |
|
|
current_layer_index = block.layer_range.stop |
|
|
if current_layer_index < self.graph.num_layers: |
|
|
new_block = BuildingBlock( |
|
|
self.graph, range(current_layer_index, self.graph.num_layers)) |
|
|
new_block.block_id = block_id |
|
|
all_blocks.append(new_block) |
|
|
sorted_layer_ids = get_sorted_layer_ids(self.network) |
|
|
for block in all_blocks: |
|
|
block.collect_relative_inter_edges(self.layer_to_block) |
|
|
for layer_id in sorted_layer_ids: |
|
|
if layer_id in block.layer_range: |
|
|
block.sorted_layer_ids.append(layer_id) |
|
|
return all_blocks |
|
|
|
|
|
def get_backbone_blocks(self): |
|
|
sorted_blocks = sorted( |
|
|
self.blocks_by_edge_hash.values(), |
|
|
key=lambda blocks: (len(blocks), len(blocks[0].layer_range)), |
|
|
) |
|
|
if len(sorted_blocks) == 0: |
|
|
return [] |
|
|
else: |
|
|
return sorted_blocks[-1] |
|
|
|
|
|
def get_blocks_by_module_hash(self, blocks): |
|
|
blocks_by_module_hash = {} |
|
|
for block in blocks: |
|
|
module_hash = self.get_module_hash(block.layer_range) |
|
|
if module_hash not in blocks_by_module_hash: |
|
|
blocks_by_module_hash[module_hash] = [] |
|
|
blocks_by_module_hash[module_hash].append(block) |
|
|
for module_hash in [*blocks_by_module_hash.keys()]: |
|
|
if len(blocks_by_module_hash[module_hash]) == 1: |
|
|
del blocks_by_module_hash[module_hash] |
|
|
return blocks_by_module_hash |
|
|
|
|
|
def get_module_tree(self): |
|
|
module_tree = {"children": {}, "name": None} |
|
|
for module_name in self.module_to_layer_range_map.keys(): |
|
|
full_name = module_name.split('.') |
|
|
current_tree = module_tree["children"] |
|
|
for depth, name in enumerate(full_name): |
|
|
if name not in current_tree: |
|
|
current_tree[name] = {"children": {}, "name": None} |
|
|
if depth == len(full_name) - 1: |
|
|
current_tree[name]["name"] = module_name |
|
|
else: |
|
|
current_tree = current_tree[name]["children"] |
|
|
return module_tree |
|
|
|
|
|
def get_blocks_by_edge_hash(self, blocks_by_module_hash): |
|
|
blocks_by_edge_hash = {} |
|
|
for block_list in blocks_by_module_hash.values(): |
|
|
for block in block_list: |
|
|
block.collect_edges() |
|
|
edge_hash = block.edge_hash |
|
|
if edge_hash not in blocks_by_edge_hash: |
|
|
blocks_by_edge_hash[edge_hash] = [] |
|
|
blocks_by_edge_hash[edge_hash].append(block) |
|
|
for edge_hash in [*blocks_by_edge_hash.keys()]: |
|
|
if len(blocks_by_edge_hash[edge_hash]) == 1: |
|
|
del blocks_by_edge_hash[edge_hash] |
|
|
else: |
|
|
block_list = blocks_by_edge_hash[edge_hash] |
|
|
blocks_by_edge_hash[edge_hash] = sorted( |
|
|
block_list, key=lambda x: x.layer_range.start) |
|
|
for type_id, block_list in enumerate(blocks_by_edge_hash.values()): |
|
|
for block in block_list: |
|
|
block.type_id = type_id |
|
|
return blocks_by_edge_hash |
|
|
|
|
|
def get_layer_to_block(self): |
|
|
layer_to_block = {} |
|
|
for block_list in self.blocks_by_edge_hash.values(): |
|
|
for block in block_list: |
|
|
for layer_index in block.layer_range: |
|
|
layer_to_block[layer_index] = block |
|
|
return layer_to_block |
|
|
|
|
|
def clean_blocks(self): |
|
|
for block in self.blocks: |
|
|
block.p2p_type = None |
|
|
block.is_superset = False |
|
|
block.is_subset = False |
|
|
|
|
|
def mark_p2p_type(self, phy_mesh, stage_phy_meshes, |
|
|
graph_config: GraphConfig): |
|
|
if len(self.backbone_blocks) == 0 or len(stage_phy_meshes) == 1: |
|
|
return |
|
|
assert len(self.backbone_blocks) % len(stage_phy_meshes) == 0 |
|
|
block_per_stage = len(self.backbone_blocks) // len(stage_phy_meshes) |
|
|
|
|
|
for block in self.backbone_blocks: |
|
|
block.p2p_type = None |
|
|
for stage_index, stage_phy_mesh in enumerate(stage_phy_meshes[:-1]): |
|
|
next_stage_phy_mesh = stage_phy_meshes[stage_index + 1] |
|
|
last_device_id = stage_phy_mesh.phy_devices_id.flatten()[-1] |
|
|
next_first_device_id = next_stage_phy_mesh.phy_devices_id.flatten( |
|
|
)[0] |
|
|
num_devices_per_host = phy_mesh.num_devices_per_host |
|
|
next_block = self.backbone_blocks[(stage_index + 1) * |
|
|
block_per_stage] |
|
|
if last_device_id // num_devices_per_host != next_first_device_id // num_devices_per_host: |
|
|
next_block.p2p_type = P2PType.CROSS_HOST |
|
|
graph_config.has_cross_host = True |
|
|
else: |
|
|
next_block.p2p_type = P2PType.CROSS_DEVICE |
|
|
graph_config.has_cross_device = True |
|
|
|
|
|
def get_graph_mapping(self): |
|
|
layer_mapping = {} |
|
|
block_mapping = {} |
|
|
p2p_types = {} |
|
|
p2p_tensors = {} |
|
|
for block_list in self.blocks_by_edge_hash.values(): |
|
|
superset_blocks = [] |
|
|
superset_block_index = {} |
|
|
for block in block_list: |
|
|
block_added = False |
|
|
for index, superset_block in enumerate(list(superset_blocks)): |
|
|
if block.p2p_type == superset_block.p2p_type: |
|
|
if block.relative_inter_edges.issubset( |
|
|
superset_block.relative_inter_edges): |
|
|
block.is_subset = True |
|
|
block.is_superset = False |
|
|
superset_block_index[id(block)] = index |
|
|
block_added = True |
|
|
break |
|
|
elif superset_block.relative_inter_edges.issubset( |
|
|
block.relative_inter_edges): |
|
|
superset_block.is_subset = True |
|
|
superset_block.is_superset = False |
|
|
block.is_subset = False |
|
|
block.is_superset = True |
|
|
superset_blocks[index] = block |
|
|
superset_block_index[id(block)] = index |
|
|
block_added = True |
|
|
break |
|
|
if not block_added: |
|
|
block.is_subset = False |
|
|
block.is_superset = True |
|
|
superset_blocks.append(block) |
|
|
superset_block_index[id(block)] = len(superset_blocks) - 1 |
|
|
for block in block_list: |
|
|
assert not (block.is_subset and block.is_superset) |
|
|
if block.is_subset: |
|
|
superset_block = superset_blocks[superset_block_index[id( |
|
|
block)]] |
|
|
block_mapping[block.block_id] = superset_block.block_id |
|
|
owned_inputs = map( |
|
|
lambda x: x[0], |
|
|
sorted(block.owned_inputs.items(), key=lambda x: x[1])) |
|
|
superset_owned_inputs = map( |
|
|
lambda x: x[0], |
|
|
sorted(superset_block.owned_inputs.items(), |
|
|
key=lambda x: x[1])) |
|
|
for from_input_id, to_input_id in zip( |
|
|
owned_inputs, superset_owned_inputs): |
|
|
from_input_name = self.network.get_input( |
|
|
from_input_id).name |
|
|
to_input_name = self.network.get_input(to_input_id).name |
|
|
layer_mapping[from_input_name] = to_input_name |
|
|
for from_layer_id, to_layer_id in zip( |
|
|
block.layer_range, superset_block.layer_range): |
|
|
from_layer = self.network.get_layer(from_layer_id) |
|
|
to_layer = self.network.get_layer(to_layer_id) |
|
|
layer_mapping[from_layer.name] = to_layer.name |
|
|
for i in range(from_layer.num_outputs): |
|
|
from_output = from_layer.get_output(i) |
|
|
if from_output.is_network_output: |
|
|
to_output = to_layer.get_output(i) |
|
|
layer_mapping[from_output.name] = to_output.name |
|
|
if block.p2p_type is not None: |
|
|
p2p_types[block.block_id] = block.p2p_type |
|
|
p2p_tensors[block.block_id] = [ |
|
|
*set(block.get_input_names()) |
|
|
] |
|
|
for from_name, to_name in zip( |
|
|
block.get_input_names(), |
|
|
superset_block.get_input_names()): |
|
|
layer_mapping[ |
|
|
f"p2p_block{block.block_id}_{from_name}"] = f"p2p_block{superset_block.block_id}_{to_name}" |
|
|
stage_id = 0 |
|
|
block_to_stage = {} |
|
|
for block in self.blocks: |
|
|
if block.p2p_type is not None: |
|
|
stage_id += 1 |
|
|
block_to_stage[block.block_id] = stage_id |
|
|
return GraphMapping( |
|
|
layer_mapping, |
|
|
block_mapping, |
|
|
p2p_types, |
|
|
p2p_tensors, |
|
|
block_to_stage, |
|
|
) |
|
|
|
|
|
def create_simplified_graph(self, graph_config: GraphConfig): |
|
|
new_graph = PipelineGraph.create_graph() |
|
|
new_graph._io_buffer_mapping = self.graph._io_buffer_mapping |
|
|
layer_mapping = graph_config.graph_mapping.layer_mapping |
|
|
|
|
|
for i in range(self.network.num_inputs): |
|
|
trt_input = self.network.get_input(i) |
|
|
if trt_input.name not in layer_mapping: |
|
|
new_graph.add_input(trt_input) |
|
|
|
|
|
last_blocks = {} |
|
|
same_spec_mapping = {} |
|
|
same_spec_layer_mapping = {} |
|
|
shape_mapping = {} |
|
|
building_block_id = 0 |
|
|
same_spec_ids = {} |
|
|
same_spec_count = 0 |
|
|
for block in self.blocks: |
|
|
if not block.is_subset: |
|
|
stage_type = None |
|
|
if not block.is_superset: |
|
|
if block.block_id == 0: |
|
|
stage_type = StageType.START |
|
|
elif block.block_id == len(self.blocks) - 1: |
|
|
stage_type = StageType.END |
|
|
input_mapping = block.get_input_mapping(last_blocks) |
|
|
for from_name, to_name in [*input_mapping.items()]: |
|
|
if to_name in same_spec_mapping: |
|
|
input_mapping[from_name] = same_spec_mapping[to_name] |
|
|
if to_name in layer_mapping: |
|
|
input_mapping[from_name] = layer_mapping[to_name] |
|
|
if block.is_superset and block.p2p_type is not None: |
|
|
for from_name, to_name in [*input_mapping.items()]: |
|
|
output_tensor = new_graph.get_tensor(to_name) |
|
|
p2p_layer = new_graph.as_trt().add_identity( |
|
|
output_tensor.as_trt()) |
|
|
p2p_layer.name = f"p2p_block{block.block_id}_{from_name}" |
|
|
p2p_layer.metadata = p2p_layer.name |
|
|
p2p_tensor = p2p_layer.get_output(0) |
|
|
p2p_tensor.name = f"{p2p_layer.name}_output" |
|
|
wrapped_layer = new_graph.register_layer(p2p_layer) |
|
|
wrapped_layer.attrs[ |
|
|
"building_block_id"] = building_block_id |
|
|
wrapped_layer.attrs["p2p_type"] = block.p2p_type |
|
|
input_mapping[from_name] = p2p_tensor.name |
|
|
shape_mapping[p2p_tensor.name] = from_name |
|
|
building_block_id += 1 |
|
|
for i in block.sorted_layer_ids: |
|
|
layer = self.network.get_layer(i) |
|
|
wrapped_layer = new_graph.add_layer( |
|
|
layer, |
|
|
input_mapping=input_mapping, |
|
|
) |
|
|
wrapped_layer.attrs["building_block_id"] = building_block_id |
|
|
wrapped_layer.attrs["stage_type"] = stage_type |
|
|
if block.is_superset: |
|
|
last_blocks[block.type_id] = block |
|
|
|
|
|
if block.type_id in same_spec_ids: |
|
|
same_spec_id = same_spec_ids[block.type_id] |
|
|
update_same_spec_count = False |
|
|
else: |
|
|
same_spec_id = same_spec_count |
|
|
same_spec_ids[block.type_id] = same_spec_id |
|
|
update_same_spec_count = True |
|
|
count = same_spec_id |
|
|
for i, (layer_offset, |
|
|
output_index) in enumerate(block.outputs): |
|
|
layer = self.network.get_layer(block.layer_range.start + |
|
|
layer_offset) |
|
|
tensor_name = layer.get_output(output_index).name |
|
|
output_tensor = new_graph.get_tensor(tensor_name) |
|
|
same_spec_layer = new_graph.as_trt().add_identity( |
|
|
output_tensor.as_trt()) |
|
|
same_spec_layer.name = f"{tensor_name}_same_spec" |
|
|
same_spec_layer.metadata = same_spec_layer.name |
|
|
same_spec_tensor = same_spec_layer.get_output(0) |
|
|
same_spec_tensor.name = f"{same_spec_layer.name}_output" |
|
|
wrapped_layer = new_graph.register_layer( |
|
|
same_spec_layer) |
|
|
wrapped_layer.attrs[ |
|
|
"building_block_id"] = building_block_id |
|
|
wrapped_layer.attrs["same_spec_id"] = count |
|
|
count += 1 |
|
|
same_spec_mapping[tensor_name] = same_spec_tensor.name |
|
|
same_spec_layer_mapping[ |
|
|
same_spec_layer.name] = layer.name |
|
|
shape_mapping[same_spec_tensor.name] = tensor_name |
|
|
for i, graph_input_index in enumerate( |
|
|
block.owned_inputs.keys()): |
|
|
input_name = self.network.get_input( |
|
|
graph_input_index).name |
|
|
input_tensor = new_graph.get_input(input_name) |
|
|
input_tensor.attrs["same_spec_id"] = count |
|
|
count += 1 |
|
|
if update_same_spec_count: |
|
|
same_spec_count = count |
|
|
building_block_id += 1 |
|
|
graph_config.graph_mapping.same_spec_layer_mapping = same_spec_layer_mapping |
|
|
|
|
|
if len(self.backbone_blocks) >= 2: |
|
|
start_block = self.backbone_blocks[0] |
|
|
if start_block.is_subset: |
|
|
start_block = self.blocks[graph_config.graph_mapping. |
|
|
block_mapping[start_block.block_id]] |
|
|
for i in start_block.layer_range: |
|
|
layer_name = self.network.get_layer(i).name |
|
|
layer = new_graph.get_layer(layer_name) |
|
|
layer.attrs["in_start_block"] = True |
|
|
end_block = self.backbone_blocks[-1] |
|
|
if end_block.is_subset: |
|
|
end_block = self.blocks[graph_config.graph_mapping. |
|
|
block_mapping[end_block.block_id]] |
|
|
for i in end_block.layer_range: |
|
|
layer_name = self.network.get_layer(i).name |
|
|
layer = new_graph.get_layer(layer_name) |
|
|
layer.attrs["in_end_block"] = True |
|
|
slowest_p2p_type = None |
|
|
if graph_config.has_cross_host: |
|
|
slowest_p2p_type = P2PType.CROSS_HOST |
|
|
elif graph_config.has_cross_device: |
|
|
slowest_p2p_type = P2PType.CROSS_DEVICE |
|
|
if slowest_p2p_type is not None: |
|
|
for block in self.blocks: |
|
|
if block.is_superset and block.p2p_type == slowest_p2p_type: |
|
|
for i in block.layer_range: |
|
|
layer_name = self.network.get_layer(i).name |
|
|
layer = new_graph.get_layer(layer_name) |
|
|
layer.attrs["in_slowest_block"] = True |
|
|
|
|
|
for i in range(self.network.num_outputs): |
|
|
trt_output = self.network.get_output(i) |
|
|
output = self.graph.get_output(trt_output.name) |
|
|
if output.producer is not None and output.producer.index in self.layer_to_block and self.layer_to_block[ |
|
|
output.producer.index].is_subset: |
|
|
continue |
|
|
if trt_output.is_shape_tensor: |
|
|
new_output = new_graph.add_output_shape(trt_output) |
|
|
else: |
|
|
new_output = new_graph.add_output(trt_output) |
|
|
sharded_io = False |
|
|
for pattern in self.sharded_io_allowlist: |
|
|
if re.match(pattern, new_output.name): |
|
|
sharded_io = True |
|
|
break |
|
|
if not sharded_io: |
|
|
new_output.producer.attrs["is_replicated"] = True |
|
|
|
|
|
for input in new_graph.inputs: |
|
|
input_name = input.name |
|
|
sharded_io = False |
|
|
for pattern in self.sharded_io_allowlist: |
|
|
if re.match(pattern, input_name): |
|
|
sharded_io = True |
|
|
break |
|
|
if not sharded_io: |
|
|
input.attrs["is_replicated"] = True |
|
|
for pattern, repl in self.same_spec_io.items(): |
|
|
if re.match(pattern, input_name): |
|
|
output_name = re.sub(pattern, repl, input_name) |
|
|
output = new_graph.get_output(output_name) |
|
|
if output is not None: |
|
|
if "same_spec_id" in input.attrs: |
|
|
same_spec_id = input.attrs["same_spec_id"] |
|
|
else: |
|
|
same_spec_id = same_spec_count |
|
|
same_spec_count += 1 |
|
|
input.attrs["same_spec_id"] = same_spec_id |
|
|
output.attrs["same_spec_id"] = same_spec_id |
|
|
if math.prod(self.graph.get_input( |
|
|
input_name).shape) < math.prod( |
|
|
self.graph.get_output(output_name).shape): |
|
|
input.attrs["no_memory_footprint"] = True |
|
|
else: |
|
|
output.attrs["no_memory_footprint"] = True |
|
|
|
|
|
return new_graph, shape_mapping |
|
|
|
|
|
def enrich_shape_info(self, shape_mapping): |
|
|
shapes = self.shape_info.shapes.copy() |
|
|
max_shapes = self.shape_info.max_shapes.copy() |
|
|
values = self.shape_info.values.copy() |
|
|
shape_layers = self.shape_info.shape_layers |
|
|
for from_name, to_name in shape_mapping.items(): |
|
|
if to_name in shapes: |
|
|
shapes[from_name] = shapes[to_name] |
|
|
if to_name in max_shapes: |
|
|
max_shapes[from_name] = max_shapes[to_name] |
|
|
if to_name in values: |
|
|
values[from_name] = values[to_name] |
|
|
shape_info = ShapeInfo(shapes, values, shape_layers, max_shapes) |
|
|
return shape_info |
|
|
|
|
|
def simplify_graph( |
|
|
self, phy_mesh: PhysicalDeviceMesh, num_stages: int, |
|
|
num_devices_per_stage: int) -> Tuple[PipelineGraph, GraphConfig]: |
|
|
num_blocks = len(self.backbone_blocks) |
|
|
if num_blocks % num_stages != 0: |
|
|
return None, None |
|
|
graph_config = GraphConfig() |
|
|
graph_config.num_micro_batches = self.num_micro_batches |
|
|
graph_config.num_blocks = num_blocks |
|
|
graph_config.num_stages = num_stages |
|
|
graph_config.phy_mesh = phy_mesh |
|
|
stage_phy_meshes = phy_mesh.split_pipeline_meshes( |
|
|
num_stages, num_devices_per_stage) |
|
|
graph_config.stage_phy_meshes = stage_phy_meshes |
|
|
with silent_trt_logger(): |
|
|
self.clean_blocks() |
|
|
self.mark_p2p_type(phy_mesh, stage_phy_meshes, graph_config) |
|
|
graph_config.graph_mapping = self.get_graph_mapping() |
|
|
new_graph, shape_mapping = self.create_simplified_graph( |
|
|
graph_config) |
|
|
shape_info = self.enrich_shape_info(shape_mapping) |
|
|
new_graph.assign_shapes(shape_info) |
|
|
return new_graph, graph_config |
|
|
|
|
|
def get_graph_mapping_for_shape(self): |
|
|
layer_mapping = {} |
|
|
tensor_mapping = {} |
|
|
for block_list in self.blocks_by_edge_hash.values(): |
|
|
head_block = block_list[0] |
|
|
for block in block_list[1:]: |
|
|
for from_layer_id, to_layer_id in zip(block.layer_range, |
|
|
head_block.layer_range): |
|
|
from_layer = self.network.get_layer(from_layer_id) |
|
|
to_layer = self.network.get_layer(to_layer_id) |
|
|
layer_mapping[from_layer.name] = to_layer.name |
|
|
for i in range(from_layer.num_outputs): |
|
|
tensor_mapping[from_layer.get_output( |
|
|
i).name] = to_layer.get_output(i).name |
|
|
return layer_mapping, tensor_mapping |
|
|
|
|
|
def create_simplified_graph_for_shape(self): |
|
|
new_graph = PipelineGraph.create_graph() |
|
|
|
|
|
for i in range(self.network.num_inputs): |
|
|
trt_input = self.network.get_input(i) |
|
|
new_graph.add_input(trt_input) |
|
|
|
|
|
head_blocks = {} |
|
|
removed_blocks = set() |
|
|
removed_layers = set() |
|
|
for block_list in self.blocks_by_edge_hash.values(): |
|
|
head_block = block_list[0] |
|
|
head_blocks[head_block.type_id] = head_block |
|
|
for block in block_list[1:]: |
|
|
removed_blocks.add(id(block)) |
|
|
for layer_index in block.layer_range: |
|
|
removed_layers.add(layer_index) |
|
|
|
|
|
for block in self.blocks: |
|
|
if not id(block) in removed_blocks: |
|
|
input_mapping = block.get_input_mapping(head_blocks) |
|
|
for i in block.sorted_layer_ids: |
|
|
layer = self.network.get_layer(i) |
|
|
new_graph.add_layer( |
|
|
layer, |
|
|
input_mapping=input_mapping, |
|
|
) |
|
|
|
|
|
for i in range(self.network.num_outputs): |
|
|
trt_output = self.network.get_output(i) |
|
|
output = self.graph.get_output(trt_output.name) |
|
|
if output.producer is not None and output.producer.index in removed_layers: |
|
|
continue |
|
|
if trt_output.is_shape_tensor: |
|
|
new_graph.add_output_shape(trt_output) |
|
|
else: |
|
|
new_graph.add_output(trt_output) |
|
|
|
|
|
return new_graph |
|
|
|
|
|
def get_full_shape_info(self, num_micro_batches): |
|
|
layer_mapping, tensor_mapping = self.graph_mapping_for_shape |
|
|
optimization_profiles = self.llm_network._generate_optimization_profiles( |
|
|
) |
|
|
if len(optimization_profiles) > 0: |
|
|
optimization_profile = optimization_profiles[-1] |
|
|
else: |
|
|
optimization_profile = None |
|
|
shape_info = get_shape_info(self.graph_for_shape.as_trt(), |
|
|
optimization_profile) |
|
|
max_shape_info = get_shape_info(self.graph_for_shape.as_trt(), |
|
|
optimization_profile, |
|
|
shape_type=ShapeType.MAX) |
|
|
shape_info.max_shapes = max_shape_info.shapes |
|
|
for removed_tensor_name, tensor_name in tensor_mapping.items(): |
|
|
shape_info.shapes[removed_tensor_name] = shape_info.shapes[ |
|
|
tensor_name] |
|
|
shape_info.max_shapes[removed_tensor_name] = shape_info.max_shapes[ |
|
|
tensor_name] |
|
|
if tensor_name in shape_info.values: |
|
|
shape_info.values[removed_tensor_name] = shape_info.values[ |
|
|
tensor_name] |
|
|
for removed_layer_name, layer_name in layer_mapping.items(): |
|
|
if layer_name in shape_info.shape_layers: |
|
|
shape_info.shape_layers.add(removed_layer_name) |
|
|
return shape_info |
|
|
|
|
|
def init_layer_hash(self): |
|
|
with silent_trt_logger(): |
|
|
optimization_profiles = self.llm_network._generate_optimization_profiles( |
|
|
) |
|
|
if len(optimization_profiles) > 0: |
|
|
optimization_profile = optimization_profiles[-1] |
|
|
else: |
|
|
optimization_profile = None |
|
|
shape_info = get_shape_info(self.network, optimization_profile) |
|
|
dtypes = {tensor.name: tensor.dtype for tensor in self.graph.tensors} |
|
|
for layer in self.graph.layers: |
|
|
layer_hash = get_cache_key( |
|
|
layer.as_trt(), |
|
|
shape_info.shapes, |
|
|
shape_info.values, |
|
|
dtypes, |
|
|
) |
|
|
layer.attrs["hash"] = layer_hash |
|
|
|
