| from typing import List |
|
|
| import pandas as pd |
| import tensorrt as trt |
|
|
| from .pipeline_graph import PipelineGraph |
| from .runtime_profiling import RuntimeProfiler |
| from .simplifier import GraphConfig, StageType |
| from .solver import CostGraph, Solver |
| from .tensor_parallel.activation_node import Activation |
| from .tensor_parallel.assertion_node import Assertion |
| from .tensor_parallel.cast_node import Cast |
| from .tensor_parallel.concatenation_node import Concatenation |
| from .tensor_parallel.constant_node import Constant |
| from .tensor_parallel.elementwise_node import ElementWise |
| from .tensor_parallel.fill_node import Fill |
| from .tensor_parallel.gather_node import Gather |
| from .tensor_parallel.identity_node import Identity |
| from .tensor_parallel.input_node import InputNode |
| from .tensor_parallel.matmul_node import MatrixMultiply |
| from .tensor_parallel.node import Node |
| from .tensor_parallel.normalization_node import Normalization |
| from .tensor_parallel.output_node import OuputNode |
| from .tensor_parallel.p2p_node import P2PNode, P2PType |
| from .tensor_parallel.plugin_node import PluginNode |
| from .tensor_parallel.plugin_nodes.gemm_node import GemmPlugin |
| from .tensor_parallel.plugin_nodes.gpt_attention_node import GPTAttentionPlugin |
| from .tensor_parallel.plugin_nodes.identity_node import IdentityPlugin |
| from .tensor_parallel.plugin_nodes.look_up_node import LookupPlugin |
| from .tensor_parallel.plugin_nodes.normalization_node import (LayernormPlugin, |
| RMSnormPlugin) |
| from .tensor_parallel.reduce_node import Reduce |
| from .tensor_parallel.select_node import Select |
| from .tensor_parallel.shape_node import Shape |
| from .tensor_parallel.shuffle_node import Shuffle |
| from .tensor_parallel.slice_node import Slice |
| from .tensor_parallel.softmax_node import SoftMax |
| from .tensor_parallel.unary_node import Unary |
|
|
| LAYER_TYPE_2_NODE_TYPE = { |
| trt.LayerType.ACTIVATION: Activation, |
| trt.LayerType.ASSERTION: Assertion, |
| trt.LayerType.CAST: Cast, |
| trt.LayerType.CONCATENATION: Concatenation, |
| trt.LayerType.CONSTANT: Constant, |
| trt.LayerType.ELEMENTWISE: ElementWise, |
| trt.LayerType.FILL: Fill, |
| trt.LayerType.GATHER: Gather, |
| trt.LayerType.IDENTITY: Identity, |
| trt.LayerType.MATRIX_MULTIPLY: MatrixMultiply, |
| trt.LayerType.NORMALIZATION: Normalization, |
| trt.LayerType.PLUGIN_V2: PluginNode, |
| trt.LayerType.REDUCE: Reduce, |
| trt.LayerType.SELECT: Select, |
| trt.LayerType.SHAPE: Shape, |
| trt.LayerType.SHUFFLE: Shuffle, |
| trt.LayerType.SLICE: Slice, |
| trt.LayerType.SOFTMAX: SoftMax, |
| trt.LayerType.UNARY: Unary, |
| } |
| |
| PLUGIN_LAYER_TYPE_2_NODE_TYPE = { |
| 'GPTAttention': GPTAttentionPlugin, |
| 'Gemm': GemmPlugin, |
| 'Layernorm': LayernormPlugin, |
| 'Rmsnorm': RMSnormPlugin, |
| 'Lookup': LookupPlugin, |
| 'Identity': IdentityPlugin, |
| } |
|
|
|
|
| class NodeGraph: |
|
|
| def __init__(self, graph: PipelineGraph): |
| self._nodes = {} |
|
|
| |
| for input in graph.inputs: |
| self._nodes[input.name] = InputNode(input) |
| for layer in graph.layers: |
| layer.to_base_class() |
| if "p2p_type" in layer.attrs: |
| self._nodes[layer.name] = P2PNode(layer) |
| elif layer.type == trt.LayerType.PLUGIN_V2: |
| layer.to_subclass() |
| plugin_type = layer.as_trt().plugin.plugin_type |
| layer.to_base_class() |
| if plugin_type in PLUGIN_LAYER_TYPE_2_NODE_TYPE: |
| node = PLUGIN_LAYER_TYPE_2_NODE_TYPE[plugin_type](layer) |
| else: |
| node = LAYER_TYPE_2_NODE_TYPE[layer.type](layer) |
| self._nodes[layer.name] = node |
| else: |
| node = LAYER_TYPE_2_NODE_TYPE[layer.type](layer) |
| self._nodes[layer.name] = node |
| for output in graph.outputs: |
| self._nodes[output.name] = OuputNode(output) |
| for node in self.nodes: |
| node.post_init(self) |
| node.node_runtime_profiler = RuntimeProfiler() |
|
|
| def get_node(self, name): |
| return self._nodes[name] |
|
|
| @property |
| def nodes(self) -> List[Node]: |
| return [*self._nodes.values()] |
|
|
| def assign_cost_weights(self, graph_config: GraphConfig): |
| layer_mapping = graph_config.graph_mapping.layer_mapping |
| for layer_name in layer_mapping.values(): |
| node = self.get_node(layer_name) |
| node.sharding_weight += 1 |
| node.resharding_weight += 1 |
| same_spec_layer_mapping = graph_config.graph_mapping.same_spec_layer_mapping |
| for same_spec_layer_name, layer_name in same_spec_layer_mapping.items(): |
| node = self.get_node(layer_name) |
| same_spec_node = self.get_node(same_spec_layer_name) |
| same_spec_node.sharding_weight = node.sharding_weight |
| same_spec_node.resharding_weight = node.resharding_weight |
|
|
| def set_slowest_stage(self, stage_type: StageType, |
| graph_config: GraphConfig): |
| num_micro_batches = graph_config.num_micro_batches |
| block_per_stage = graph_config.num_blocks // graph_config.num_stages |
| block_pipeline_weight = block_per_stage * (num_micro_batches - 1) |
| for node in self.nodes: |
| node.pipeline_weight = 0 |
| node.cost_level = -1 |
| if node.stage_type == StageType.START: |
| if stage_type == StageType.START: |
| node.pipeline_weight = num_micro_batches - 1 |
| node.cost_level = 1 |
| else: |
| node.cost_level = 0 |
| if stage_type == StageType.START and node.in_start_block: |
| node.pipeline_weight = block_pipeline_weight |
| if node.stage_type == StageType.END: |
| if stage_type == StageType.END: |
| node.pipeline_weight = num_micro_batches - 1 |
| node.cost_level = 1 |
| else: |
| node.cost_level = 0 |
| if stage_type == StageType.END and node.in_end_block: |
| node.pipeline_weight = block_pipeline_weight |
| if isinstance(node, P2PNode): |
| if (graph_config.has_cross_host |
| and node.p2p_type == P2PType.CROSS_HOST) or ( |
| not graph_config.has_cross_host |
| and node.p2p_type == P2PType.CROSS_DEVICE): |
| if stage_type == StageType.BLOCK: |
| node.pipeline_weight += num_micro_batches - 1 |
| node.cost_level = 1 |
| else: |
| node.cost_level = 0 |
| elif (graph_config.has_cross_device |
| and node.p2p_type == P2PType.CROSS_DEVICE) or ( |
| not graph_config.has_cross_device |
| and node.p2p_type == P2PType.CROSS_HOST): |
| node.pipeline_weight += num_micro_batches - 1 |
| if stage_type == StageType.BLOCK and node.in_slowest_block: |
| node.pipeline_weight = block_pipeline_weight |
|
|
| def get_cost_graph(self, lmesh): |
| leaf_strategies = [] |
| for node in self.nodes: |
| if node.is_replicated: |
| node.set_strategy(None, lmesh) |
| else: |
| node.collect_strategies(lmesh) |
| for node in self.nodes: |
| strategies_vector = node.update_resharding_cost() |
| if len(strategies_vector) != 0: |
| leaf_strategies.append(strategies_vector) |
| cost_graph = CostGraph(leaf_strategies) |
| return cost_graph |
|
|
| def find_solution(self, cost_graph, memory_budget): |
| solver = Solver(cost_graph, memory_budget=memory_budget) |
| solution = solver.find_solution()[1] |
|
|
| graph_strategy = solution.node_best_strategy |
| for node_name, strategy in graph_strategy.items(): |
| node = self._nodes[node_name] |
| for idx, pre_node in enumerate(node.predecessor_nodes): |
| if pre_node is None: |
| continue |
| if pre_node.node_name not in strategy.best_resharding_cost: |
| continue |
| strategy.best_resharding_cost[ |
| idx] = strategy.best_resharding_cost[pre_node.node_name] |
| strategy.node_names[idx] = pre_node.node_name |
| for key in list(strategy.best_resharding_cost.keys()): |
| if isinstance(key, str): |
| del strategy.best_resharding_cost[key] |
|
|
| return solution |
|
|
| def visualize(self, name='pp_graph'): |
| with open(name + '.dot', 'w') as f: |
| f.write("digraph {\n") |
| ''' |
| f.write(" // Value Nodes\n") |
| for name, tensor in self._tensors.items(): |
| f.write(" \"{}\" [fillcolor = \"green\", label = \"{}\", shape = \"box\", style = \"filled\"];\n".format(name, tensor.shape)) |
| ''' |
| f.write(" // Operation Nodes\n") |
| for name, node in self._nodes.items(): |
| fillcolor = 'white' |
| if 'MATRIX_MULTIPLY' in name: |
| fillcolor = 'green' |
| label = name |
| if len(node.outputs) > 0: |
| label = name + '\\n' + str(node.outputs[0].shape) |
| f.write( |
| " \"{}\" [fillcolor = \"{}\", label = \"{}\", shape = \"box\", style = \"filled\"];\n" |
| .format(name, fillcolor, label)) |
| f.write(" // Edges\n") |
| for name, node in self._nodes.items(): |
| for successor_node in node.successor_nodes: |
| if successor_node: |
| f.write(" \"{}\" ->\"{}\";\n".format( |
| name, successor_node.node_name)) |
| f.write(" }\n") |
|
|
| def visualize_solution(self, |
| solution, |
| fname='pp_graph_solution', |
| ignore_shape_io=True): |
| with open(fname + '.dot', 'w') as f: |
| names, costs, block_ids = [], [], [] |
| f.write("digraph {\n") |
| f.write(" // Operation Nodes\n") |
| for name, node in self._nodes.items(): |
| if ignore_shape_io and node.layer is not None and node.layer.is_shape_io: |
| continue |
| cost = 0.0 |
| fillcolor = 'white' |
| if 'MATRIX_MULTIPLY' in name or 'PLUGIN_V2_Gemm' in name: |
| fillcolor = 'orange' |
| elif '_same_spec' in name: |
| fillcolor = 'gray' |
| elif 'p2p_block' in name: |
| fillcolor = 'blue' |
| elif 'PLUGIN' in name: |
| fillcolor = 'yellow' |
|
|
| shape = 'box' |
| if 'output_node' == node.node_type or 'input_node' == node.node_type: |
| shape = 'ellipse' |
| fillcolor = 'green' |
|
|
| label = name + f'_block{node.building_block_id}_weight{node.sharding_weight}' |
| if len(node.inputs) > 0: |
| for idx, input in enumerate(node.inputs): |
| if not input: |
| continue |
| label = label + f'\\ninput{idx}_' + str( |
| input.shape) + f'_{input.dtype_str_size[0]}_' |
| if node.node_name in solution.node_best_strategy: |
| best_strategy = solution.node_best_strategy[ |
| node.node_name] |
| shard_seq = str( |
| best_strategy.sharding_specs[f'input{idx}']. |
| sharding_sequence) |
| label = label + shard_seq |
| if idx not in best_strategy.best_resharding_cost: |
| continue |
| rcosts = best_strategy.best_resharding_cost[idx][0] |
| comm_action_sequence, resharding_cost = rcosts[ |
| 1], rcosts[2] |
| if len(comm_action_sequence) > 0: |
| label = label + '|' |
| for commspec in comm_action_sequence: |
| comm = [ |
| commspec.comm_pattern, commspec.gather_dim, |
| commspec.shard_dim, |
| commspec.logical_process_axis |
| ] |
| label = label + '->' + str(comm) |
| if resharding_cost > 0: |
| label = label + '_rcost{:.2}'.format( |
| resharding_cost) |
| cost = cost + resharding_cost |
| if len(node.outputs) > 0: |
| best_strategy = None |
| for idx, output in enumerate(node.outputs): |
| label = label + f'\\noutput{idx}_' + str( |
| output.shape) + f'_{output.dtype_str_size[0]}' |
| if node.node_name in solution.node_best_strategy: |
| best_strategy = solution.node_best_strategy[ |
| node.node_name] |
| shard_seq = str( |
| best_strategy.sharding_specs[f'output{idx}']. |
| sharding_sequence) |
| comm = None |
| if f'output{idx}' in best_strategy.communication_actions: |
| commspec = best_strategy.communication_actions[ |
| f'output{idx}'] |
| comm = [ |
| commspec.comm_pattern, commspec.gather_dim, |
| commspec.shard_dim, |
| commspec.logical_process_axis |
| ] |
| label = label + '_' + shard_seq |
| if comm: |
| label = label + f' | {comm}' |
| if best_strategy: |
| cost = cost + best_strategy.sharding_cost + best_strategy.communication_cost |
| label = label + '| scost{:.2}'.format( |
| best_strategy.sharding_cost) |
| if best_strategy.communication_cost > 0: |
| label = label + ' | ccost{:.2}'.format( |
| best_strategy.communication_cost) |
| names.append(name) |
| costs.append(cost) |
| block_ids.append([ |
| node.building_block_id, node.cost_level, |
| node.sharding_weight + node.pipeline_weight, |
| node.same_spec_id |
| ]) |
| f.write( |
| " \"{}\" [fillcolor = \"{}\", label = \"{}\", shape = \"{}\", style = \"filled\"];\n" |
| .format(name, fillcolor, label, shape)) |
| f.write(" // Edges\n") |
| for name, node in self._nodes.items(): |
| if ignore_shape_io and node.layer is not None and node.layer.is_shape_io: |
| continue |
| for successor_node in node.successor_nodes: |
| if successor_node: |
| if ignore_shape_io and successor_node.layer is not None and successor_node.layer.is_shape_io: |
| continue |
| f.write(" \"{}\" ->\"{}\";\n".format( |
| name, successor_node.node_name)) |
| f.write(" }\n") |
| df = pd.DataFrame.from_dict({ |
| 'node': |
| names, |
| 'cost': |
| costs, |
| 'block_id': [block[0] for block in block_ids], |
| 'cost_level': [block[1] for block in block_ids], |
| 'sharding_weight': [block[2] for block in block_ids], |
| 'same_spec_id': [block[3] for block in block_ids] |
| }) |
| df['weight_cost'] = df['sharding_weight'] * df['cost'] |
| df.to_csv(fname + '.csv') |
|
|
