Skip to content

Commit 020b0f6

Browse files
author
Martin Yuan
committed
Print delegation info in export_llama in verbose
1 parent fedb035 commit 020b0f6

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

examples/models/llama/export_llama_lib.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import torch
2424

2525
from executorch.devtools.etrecord import generate_etrecord
26+
from executorch.devtools.backend_debug import get_delegation_info
27+
from tabulate import tabulate
2628
from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass
2729

2830
from executorch.extension.llm.export.builder import DType, LLMEdgeManager
@@ -816,6 +818,13 @@ def _export_llama(args) -> LLMEdgeManager: # noqa: C901
816818

817819
builder = builder.to_executorch(passes=additional_passes)
818820

821+
if args.verbose:
822+
graph_module = builder.edge_manager.exported_program().graph_module
823+
delegation_info = get_delegation_info(graph_module)
824+
print(delegation_info.get_summary())
825+
df = delegation_info.get_operator_delegation_dataframe()
826+
print(tabulate(df, headers="keys", tablefmt="fancy_grid"))
827+
819828
if args.profile_memory:
820829
generate_memory_trace(builder.export_program, "memory_profile.json")
821830

0 commit comments

Comments
 (0)