memory viewer peak memory take into account fragmentation in HLO temp…

…oraries. also add argument size, and hlo temp size etc PiperOrigin-RevId: 551913310
tensorflow · Jul 28, 2023 · 630af93 · 630af93
1 parent b8a419e
commit 630af93
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 14 deletions.
diff --git a/frontend/app/common/interfaces/memory_viewer_preprocess.jsonpb_decls.d.ts.gz b/frontend/app/common/interfaces/memory_viewer_preprocess.jsonpb_decls.d.ts.gz
diff --git a/frontend/app/components/memory_viewer/memory_usage/memory_usage.ts b/frontend/app/components/memory_viewer/memory_usage/memory_usage.ts
@@ -31,7 +31,10 @@ export class MemoryUsage {
   private smallBufferCount: number;
 
   peakHeapSizeBytes: number;
-  unpaddedPeakHeapSizeBytes: number;
+  paddingOverhead: number;
+  hloTempSizeBytes: number;
+  hloTempFragmentation: number;
+  totalArgumentSizeBytes: number;
   peakLogicalBuffers: number[];
   peakHeapSizePosition: number;
   indefiniteMemoryUsageBytes: MemoryUsageBytes;
@@ -68,7 +71,10 @@ export class MemoryUsage {
     this.smallBufferCount = 0;
 
     this.peakHeapSizeBytes = 0;
-    this.unpaddedPeakHeapSizeBytes = 0;
+    this.paddingOverhead = 0;
+    this.hloTempSizeBytes = 0;
+    this.hloTempFragmentation = 0;
+    this.totalArgumentSizeBytes = 0;
     this.peakLogicalBuffers = [];
     this.peakHeapSizePosition = 0;
     this.indefiniteMemoryUsageBytes = {padded: 0, unpadded: 0};
@@ -127,9 +133,21 @@ export class MemoryUsage {
     if (!this.timelineUrl.startsWith('/memory_viewer.json')) {
       this.timelineUrl = '';
     }
-    this.peakHeapSizeBytes = (preprocess.peakHeapMib || 0) * 1024 * 1024;
-    this.unpaddedPeakHeapSizeBytes =
+    this.peakHeapSizeBytes =
+        (preprocess.totalBufferAllocationMib || 0) * 1024 * 1024;
+    this.paddingOverhead = (preprocess.peakHeapMib || 0) * 1024 * 1024 -
         (preprocess.peakUnpaddedHeapMib || 0) * 1024 * 1024;
+    this.totalArgumentSizeBytes =
+        (preprocess.entryComputationParametersMib || 0) * 1024 * 1024;
+    this.hloTempSizeBytes = this.peakHeapSizeBytes -
+        (preprocess.indefiniteBufferAllocationMib || 0) * 1024 * 1024;
+    const framenentationSizeBytes = this.peakHeapSizeBytes -
+        (preprocess.peakUnpaddedHeapMib || 0) * 1024 * 1024;
+    if (this.hloTempSizeBytes) {
+      this.hloTempFragmentation =
+          framenentationSizeBytes / this.hloTempSizeBytes;
+    }
+
     this.peakHeapSizePosition = (preprocess.peakHeapSizePosition || 0);
     this.heapSizes = preprocess.heapSizes || [];
     this.unpaddedHeapSizes = preprocess.unpaddedHeapSizes || [];
@@ -292,7 +310,7 @@ export class MemoryUsage {
     this.peakHeapSizePosition = peakHeapSizePosition;
 
     this.peakHeapSizeBytes = peakHeapSizeBytes;
-    this.unpaddedPeakHeapSizeBytes = unpaddedPeakHeapSizeBytes;
+    this.paddingOverhead = peakHeapSizeBytes - unpaddedPeakHeapSizeBytes;
     this.heapSizes = heapSizes;
     this.unpaddedHeapSizes = unpaddedHeapSizes;
   }

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ng.html b/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ng.html
@@ -5,8 +5,15 @@
     <div>
       <div class="mat-headline">Module name: {{moduleName}}</div>
       <div class="sub-title">Peak memory allocation is {{peakHeapSizeMiB}} MiB</div>
-      <div class="description" [hidden]="!unpaddedPeakHeapSizeMiB">
-        ({{unpaddedPeakHeapSizeMiB}} MiB without padding)
+      <div class="description" [hidden]="!paddingOverhead">
+        {{paddingOverhead}} MiB total padding overhead
+      </div>
+      <div class="description" [hidden]="!totalArgumentSizeBytes">
+        {{totalArgumentSizeBytes}} MiB total argument size
+      </div>
+      <div class="description" [hidden]="!hloTempSizeBytes">
+        {{hloTempSizeBytes}} MiB HLO temporariy variable
+        (fragmentation  {{hloTempFragmentation}}%)
       </div>
       <div class="description">
         Modifying your model's architecture, batch size and data dimensions may

diff --git a/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts b/frontend/app/components/memory_viewer/memory_viewer_main/memory_viewer_main.ts
@@ -30,12 +30,15 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
   /** XLA memory space color */
   @Input() memorySpaceColor: number = 0;
 
-  moduleName: string = '';
+  moduleName = '';
   peakInfo?: BufferAllocationInfo;
   activeInfo?: BufferAllocationInfo;
-  peakHeapSizeMiB: string = '';
-  unpaddedPeakHeapSizeMiB: string = '';
-  timelineUrl: string = '';
+  peakHeapSizeMiB = '';
+  paddingOverhead = '';
+  totalArgumentSizeBytes = '';
+  hloTempSizeBytes = '';
+  hloTempFragmentation = '';
+  timelineUrl = '';
   usage?: MemoryUsage;
   heapSizes: number[] = [];
   maxHeap: HeapObject[] = [];
@@ -149,9 +152,14 @@ export class MemoryViewerMain implements OnDestroy, OnChanges {
 
     this.peakHeapSizeMiB =
         utils.bytesToMiB(this.usage.peakHeapSizeBytes).toFixed(2);
-    this.unpaddedPeakHeapSizeMiB =
-        utils.bytesToMiB(this.usage.unpaddedPeakHeapSizeBytes).toFixed(2);
-
+    this.paddingOverhead =
+        utils.bytesToMiB(this.usage.paddingOverhead).toFixed(2);
+    this.totalArgumentSizeBytes =
+        utils.bytesToMiB(this.usage.totalArgumentSizeBytes).toFixed(2);
+    this.hloTempSizeBytes =
+        utils.bytesToMiB(this.usage.hloTempSizeBytes).toFixed(2);
+    this.hloTempFragmentation =
+        (this.usage.hloTempFragmentation * 100.0).toFixed(2);
     this.heapSizes = this.usage.heapSizes || [];
     this.unpaddedHeapSizes = this.usage.unpaddedHeapSizes || [];
     this.peakInfo = {