def get_bytes(node: torch.fx.Node): nonlocal total_output_bytes nonlocal reports if not is_node_output_tensor(node): reports += f"Output {node.name} is not a tensor, this might cause problems during lowering!\n" else: total_output_bytes += get_size_of_node(submod, node)[0]
def split_preview(self, dump_graph: bool = False): reports = "" subgraphs = self.put_nodes_into_subgraphs() acc_subgraphs_num = len([g for g in subgraphs if g.is_acc]) cpu_subgraphs_num = len(subgraphs) - acc_subgraphs_num reports += f"Before removing small acc subgraphs, total {len(subgraphs)} subgraphs are created:" reports += f" {acc_subgraphs_num} acc subgraphs and {cpu_subgraphs_num} cpu subgraphs.\n" subgraphs = self.remove_small_acc_subgraphs(subgraphs) acc_subgraphs_num = len([g for g in subgraphs if g.is_acc]) cpu_subgraphs_num = len(subgraphs) - acc_subgraphs_num reports += f"After removing small acc subgraphs, total {len(subgraphs)} subgraphs are created:" reports += f" {acc_subgraphs_num} acc subgraphs and {cpu_subgraphs_num} cpu subgraphs.\n" for i, subgraph in enumerate(subgraphs): reports += f"_run_on_acc_{i}: " if subgraph.is_acc else f"{self.non_acc_submodule_name}{i}: " reports += f"{len(subgraph.nodes)} node(s)\n" self.tag(subgraphs) split_mod = self.split(remove_tag=True) split_mod.eval() if dump_graph: drawer = FxGraphDrawer( split_mod, "preview", ignore_getattr=True ) dot_graphs = drawer.get_all_dot_graphs() for name, dot_graph in dot_graphs.items(): dot_graph.write_raw(f"{name}.dot") max_qps: float = self.PCIe_BW bottleneck_module = "" for node in split_mod.graph.nodes: if node.op == "call_module" and "acc" in node.target: reports += f"\nProcessing acc submodule {node.target}\n" submod = getattr(split_mod, node.target) def get_submod_inputs(main_mod, submod, example_inputs): sub_inputs = None def get_inputs(self, inputs): nonlocal sub_inputs sub_inputs = inputs handle = submod.register_forward_pre_hook(get_inputs) main_mod(*example_inputs) handle.remove() return sub_inputs submod_inputs = get_submod_inputs( split_mod, submod, self.sample_input ) ShapeProp(submod).propagate(*submod_inputs) total_input_bytes = 0 total_output_bytes = 0 reports += "Checking inputs...\n" for n in submod.graph.nodes: if n.op == "placeholder": if not is_node_output_tensor(n): reports += f"Input {n.name} is not a tensor, this might cause problems during lowering!\n" else: total_input_bytes += get_size_of_node(submod, n)[0] if n.op == "output": output_node = n reports += "Checking outputs...\n" def get_bytes(node: torch.fx.Node): nonlocal total_output_bytes nonlocal reports if not is_node_output_tensor(node): reports += f"Output {node.name} is not a tensor, this might cause problems during lowering!\n" else: total_output_bytes += get_size_of_node(submod, node)[0] map_arg(output_node.args, get_bytes) qps = self.PCIe_BW / max(total_input_bytes, total_output_bytes) reports += f"Total input size in bytes is {total_input_bytes}, total output size in bytes is {total_output_bytes}," reports += f" theoretical max qps (bounds by PCIe bandwidth) for this submodule is {qps}.\n" if qps < max_qps: max_qps = qps bottleneck_module = node.target try: lowered_submod = self._lower_model_to_backend(submod, submod_inputs) except RuntimeError: reports += "Run into an error during lowering!\n" reports += self._find_culprit(submod, submod_inputs) continue try: lowered_submod(*submod_inputs) except RuntimeError: reports += "Run into an error during inference!\n" reports += self._find_culprit(submod, submod_inputs) else: reports += "Lowering and running succeed!\n" reports += f"\nTheoretical max qps (bounds by PCIe bandwidth) for this model is {max_qps}," reports += f" bottleneck is submodule {bottleneck_module}." print(reports) # return the reports for testing purposes return reports