Lint and print

Author: jackzhxng

examples/models/llama/export_llama_lib.py

@@ -70,10 +70,10 @@
     replace_sdpa_with_simple_sdpa,
 )
 
-from .source_transformation.vulkan_rope import replace_with_vulkan_rotary_emb
-
 from .source_transformation.torchtune.attention import replace_mha_with_inference_mha
 
+from .source_transformation.vulkan_rope import replace_with_vulkan_rotary_emb
+
 
 IS_FBCODE = True  #  os.environ.get("FBCODE_PLATFORM", False)
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -1019,4 +1019,8 @@ def _get_source_transforms(  # noqa
     if args.vulkan:
         transforms.append(replace_with_vulkan_rotary_emb)
 
+    print(
+        f"Performing the following source transformations: {[transform.__name__ for transform in transforms]}"
+    )
+
     return transforms

examples/models/llama/source_transformation/torchtune/attention.py

@@ -32,6 +32,7 @@ def _replace_mha_with_inference_mha(module: torch.nn.Module) -> None:
         else:
             replace_mha_with_inference_mha(child)
 
+
 def replace_mha_with_inference_mha(module: torch.nn.Module) -> torch.nn.Module:
     """
     Replace TorchTune's MHA with an inference friendly version of MHA that