Timing

optimization.py

@@ -1,6 +1,7 @@
 """
 """
 
+from datetime import datetime
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
@@ -38,17 +39,25 @@ INDUCTOR_CONFIGS = {
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
+    t0 = datetime.now()
+
     @spaces.GPU(duration=1500)
     def compile_transformer():
 
+        print('compile_transformer', -(t0 - (t0 := datetime.now())))
+        
         with capture_component_call(pipeline, 'transformer') as call:
             pipeline(*args, **kwargs)
 
+        print('capture_component_call', -(t0 - (t0 := datetime.now())))
+        
         dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
 
+        print('quantize_', -(t0 - (t0 := datetime.now())))
+        
         hidden_states: torch.Tensor = call.kwargs['hidden_states']
         hidden_states_transposed = hidden_states.transpose(-1, -2).contiguous()
         if hidden_states.shape[-1] > hidden_states.shape[-2]:
@@ -65,6 +74,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             dynamic_shapes=dynamic_shapes,
         )
 
+        print('exported_landscape', -(t0 - (t0 := datetime.now())))
+        
         exported_portrait = torch.export.export(
             mod=pipeline.transformer,
             args=call.args,
@@ -72,10 +83,15 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             dynamic_shapes=dynamic_shapes,
         )
 
-        return (
-            aoti_compile(exported_landscape, INDUCTOR_CONFIGS),
-            aoti_compile(exported_portrait, INDUCTOR_CONFIGS),
-        )
+        print('exported_portrait', -(t0 - (t0 := datetime.now())))
+
+        compiled_landscape = aoti_compile(exported_landscape, INDUCTOR_CONFIGS)
+        print('compiled_landscape', -(t0 - (t0 := datetime.now())))
+
+        compiled_portrait = aoti_compile(exported_portrait, INDUCTOR_CONFIGS)
+        print('compiled_portrait', -(t0 - (t0 := datetime.now())))
+
+        return compiled_landscape, compiled_portrait
 
     compiled_landscape, compiled_portrait = compile_transformer()