Fixes for xFormers profiler

# Changes
* No longer errors when no summary is available
* MemTrace with PT2.0: No longer throws a ValueError when no allocation was made

ghstack-source-id: a907b59fb5af078daad130108727f758afbd0b68
Pull Request resolved: https://github.com/fairinternal/xformers/pull/512

__original_commit__ = fairinternal/xformers@fd34c96c885b547d8f8d093dbb3551626bf55829

xformers/profiler/profiler.py

@@ -137,6 +137,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             return
         assert self.enabled
         snapshot = torch.cuda.memory._snapshot()
+        torch.cuda.memory._record_memory_history(False)
+        # No data was recorded - avoids a `ValueError` in `trace_plot`
+        if all(len(t) == 0 for t in snapshot["device_traces"]):
+            self.main_profiler.summary.append(("MemTrace", "(no allocation recorded)"))
+            return
         # Dump to disk
         filename = os.path.abspath(
             os.path.join(
@@ -151,7 +156,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     snapshot, device=None, plot_segments=False
                 )
             )
-        torch.cuda.memory._record_memory_history(False)
 
     def step(self) -> None:
         pass
@@ -341,10 +345,12 @@ def step(self) -> None:
             self.parents = ["Global"]
             self.update_profilers_on_step()
         if self.done_steps == self.last_step:
-            logger.info("xFormers profiler done - summary:\n%s", self.format_summary())
+            logger.info("xFormers profiler done. %s", self.format_summary())
 
     def format_summary(self) -> str:
+        if len(self.summary) == 0:
+            return ""
         pad_titles = max(len(title) for title, value in self.summary)
-        return "\n".join(
+        return "summary:\n" + "\n".join(
             [f"  {title.ljust(pad_titles)}: {value}" for title, value in self.summary]
         )