print graph break in llama (#2469)

Summary:
bypass-github-pytorch-ci-checks

Pull Request resolved: #2469

print the graph break when exporting llama

Reviewed By: mikekgfb

Differential Revision: D54969119

fbshipit-source-id: f94f0cb22258ff4adc5bd23a15f604523f1fd4c9

examples/models/llama2/builder.py

@@ -20,6 +20,8 @@
 )
 from executorch.exir import EdgeProgramManager
 from executorch.exir.backend.partitioner import Partitioner
+
+from executorch.exir.backend.utils import print_delegated_graph
 from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
 
 from executorch.exir.passes import MemoryPlanningPass
@@ -34,7 +36,6 @@
 from ...portable.utils import export_to_edge, save_pte_program
 from ..model_factory import EagerModelFactory
 
-
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT)
 
@@ -306,6 +307,11 @@ def to_backend(
             assert self.edge_manager is not None
             self.edge_manager = self.edge_manager.to_backend(partitioner)
             if self.verbose:
+                logging.info(
+                    print_delegated_graph(
+                        self.edge_manager.exported_program().graph_module
+                    )
+                )
                 logging.info(f"Applied partitioners: {partitioner}")
         else:
             logging.warning("Invalid partitioner, skipping...")