Improving memory snapshot (mosaicml#3315)

* add all stack trace * Update memory_snapshot.py * Update memory_snapshot.py --------- Co-authored-by: Charles Tang <[email protected]>
mvpatel2000 · May 28, 2024 · 1c1f36e · 1c1f36e
1 parent 91f2945
commit 1c1f36e
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/composer/callbacks/memory_snapshot.py b/composer/callbacks/memory_snapshot.py
@@ -141,9 +141,8 @@ def start_record_memory_history(self) -> None:
 
         log.info('Starting snapshot record_memory_history')
         torch.cuda.memory._record_memory_history(
-            True,  # type: ignore
-            trace_alloc_max_entries=self.max_entries,
-            trace_alloc_record_context=True,
+            'all',  # type: ignore
+            max_entries=self.max_entries,
         )
 
     def stop_record_memory_history(self) -> None: