Skip to content

Commit 1371374

Browse files
drisspgdrisspg
committed
tweaks
1 parent c4cc1f7 commit 1371374

3 files changed

Lines changed: 164 additions & 17 deletions

File tree

‎examples/memory_viz_transformer.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
"""
2+
Example: 8-layer transformer memory visualization with D3.
3+
4+
Creates an AttentionStack (8 RopeAttention layers), runs a forward + backward
5+
pass on GPU, and captures a memory snapshot rendered as an interactive D3 HTML.
6+
7+
Usage:
8+
python examples/memory_viz_transformer.py
9+
python examples/memory_viz_transformer.py --batch_size 2 --seq_len 1024
10+
"""
11+
12+
import torch
13+
from pathlib import Path
14+
from jsonargparse import CLI
15+
16+
from transformer_nuggets.misc.attention import AttentionStack
17+
from transformer_nuggets.utils.memory_viz import generate_memory_html
18+
19+
20+
def main(
21+
batch_size: int = 2,
22+
seq_len: int = 512,
23+
dim: int = 1024,
24+
num_heads: int = 8,
25+
num_layers: int = 8,
26+
output: str = "data/transformer_memory_viz.html",
27+
) -> None:
28+
output_path = Path(output)
29+
output_path.parent.mkdir(parents=True, exist_ok=True)
30+
31+
torch.cuda.memory._record_memory_history(enabled="all", context="all", stacks="all")
32+
33+
model = AttentionStack(
34+
num_layers=num_layers,
35+
dim=dim,
36+
num_heads=num_heads,
37+
backend="sdpa",
38+
causal=True,
39+
).to("cuda", torch.bfloat16)
40+
41+
x, block_mask = model.get_input(batch_size=batch_size, seq_len=seq_len)
42+
43+
model(x, block_mask=block_mask)
44+
torch.cuda.synchronize()
45+
46+
out = model(x, block_mask=block_mask)
47+
loss = out.sum()
48+
loss.backward()
49+
torch.cuda.synchronize()
50+
51+
snapshot = torch.cuda.memory._snapshot()
52+
torch.cuda.memory._record_memory_history(enabled=None)
53+
54+
html = generate_memory_html(snapshot, title=output_path.stem)
55+
output_path.write_text(html)
56+
print(f"D3 memory visualization saved to: {output_path}")
57+
58+
59+
if __name__ == "__main__":
60+
CLI(main, as_positional=False)

‎transformer_nuggets/utils/memory_viz.py‎

Lines changed: 101 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,15 @@
1515
_BARE_NOISE_PREFIXES = (
1616
"_Py",
1717
"Py_",
18-
"Py",
18+
"PyEval_",
19+
"PyObject_",
20+
"PyRun_",
1921
"pyrun",
2022
"pymain",
2123
"run_mod",
22-
"slot_",
23-
"method_",
24+
"slot_tp_",
2425
"cfunction_",
2526
"vectorcall",
26-
"_call",
2727
"__libc_",
2828
"_start",
2929
)
@@ -325,6 +325,8 @@ def generate_memory_html(
325325
padding: 12px 24px;
326326
border-bottom: 1px solid var(--border);
327327
flex-shrink: 0;
328+
position: relative;
329+
z-index: 60;
328330
}
329331
330332
#header h1 { font-size: 14px; font-weight: 500; font-family: var(--mono); letter-spacing: 0.03em; text-transform: uppercase; flex-shrink: 0; }
@@ -372,6 +374,45 @@ def generate_memory_html(
372374
373375
#help-trigger:hover #help-dropdown { display: block; }
374376
377+
#settings-trigger {
378+
cursor: pointer;
379+
position: relative;
380+
font-size: 14px;
381+
opacity: 0.6;
382+
transition: opacity 0.15s;
383+
user-select: none;
384+
}
385+
#settings-trigger:hover { opacity: 1; }
386+
#settings-dropdown {
387+
display: none;
388+
position: absolute;
389+
top: 100%;
390+
right: 0;
391+
margin-top: 6px;
392+
background: var(--tooltip-bg);
393+
border: 1px solid var(--border);
394+
border-radius: 4px;
395+
padding: 8px 12px;
396+
white-space: nowrap;
397+
z-index: 50;
398+
box-shadow: 0 8px 24px rgba(0, 0, 0, 0.6);
399+
font-family: var(--mono);
400+
font-size: 11px;
401+
color: var(--text-muted);
402+
}
403+
#settings-trigger.open #settings-dropdown { display: block; }
404+
#settings-dropdown label { display: flex; align-items: center; gap: 6px; }
405+
#settings-dropdown select {
406+
background: var(--bg);
407+
color: var(--text);
408+
border: 1px solid var(--border);
409+
border-radius: 3px;
410+
padding: 2px 4px;
411+
font-size: 11px;
412+
font-family: var(--mono);
413+
cursor: pointer;
414+
}
415+
375416
#controls {
376417
display: flex;
377418
gap: 12px;
@@ -873,6 +914,16 @@ def generate_memory_html(
873914
<input type="checkbox" id="dim-persistent-toggle">
874915
Hide never-freed
875916
</label>
917+
<span id="settings-trigger" title="Settings">&#9881;
918+
<div id="settings-dropdown">
919+
<label>Color by
920+
<select id="color-mode">
921+
<option value="stack">stack</option>
922+
<option value="size">size</option>
923+
</select>
924+
</label>
925+
</div>
926+
</span>
876927
</div>
877928
</div>
878929
<div id="main">
@@ -948,8 +999,38 @@ def generate_memory_html(
948999
return PALETTE[stackIdx % PALETTE.length];
9491000
}
9501001
951-
const PERSISTENT_COLOR = '#8A8F98';
952-
const PERSISTENT_ALPHAS = [0.18, 0.24, 0.3];
1002+
const SIZE_PALETTE = [
1003+
'#2E7DB5', '#3E93CC', '#5BA8D9', '#78BBE3',
1004+
'#3ECCC1', '#49C963', '#6DD883', '#8DE49D',
1005+
'#C9CC3E', '#D9DB5B', '#E0C05B', '#CC6B3E',
1006+
'#E08A5B', '#bd93f9', '#a06eed', '#F0A478',
1007+
];
1008+
1009+
const allocSizes = ALLOCS.map(a => a.s);
1010+
const sortedSizes = [...new Set(allocSizes)].sort((a, b) => a - b);
1011+
const sizeToColorIdx = new Map();
1012+
sortedSizes.forEach((s, i) => sizeToColorIdx.set(s, i % SIZE_PALETTE.length));
1013+
1014+
function getSizeColor(allocIdx) {
1015+
return SIZE_PALETTE[sizeToColorIdx.get(ALLOCS[allocIdx].s)];
1016+
}
1017+
1018+
let colorMode = 'stack';
1019+
1020+
function recolorAllocs() {
1021+
let pIdx = 0;
1022+
for (let i = 0; i < ALLOCS.length; i++) {
1023+
const isPersistent = allocPersistent[i];
1024+
allocColors[i] = colorMode === 'size'
1025+
? getSizeColor(i)
1026+
: getColor(ALLOCS[i].si);
1027+
allocAlphas[i] = isPersistent
1028+
? PERSISTENT_ALPHAS[pIdx++ % PERSISTENT_ALPHAS.length]
1029+
: 0.85;
1030+
}
1031+
}
1032+
1033+
const PERSISTENT_ALPHAS = [0.55, 0.62, 0.70];
9531034
9541035
const tooltipEl = document.getElementById('tooltip');
9551036
const detailBody = document.getElementById('detail-body');
@@ -1146,17 +1227,10 @@ def generate_memory_html(
11461227
const allocPersistent = new Uint8Array(ALLOCS.length);
11471228
const allocColors = new Array(ALLOCS.length);
11481229
const allocAlphas = new Float64Array(ALLOCS.length);
1149-
let persistentAlphaIdx = 0;
11501230
for (let i = 0; i < ALLOCS.length; i++) {
1151-
const isPersistent = allocEnds[i] >= META.max_timestep;
1152-
allocPersistent[i] = isPersistent ? 1 : 0;
1153-
allocColors[i] = isPersistent
1154-
? PERSISTENT_COLOR
1155-
: getColor(ALLOCS[i].si);
1156-
allocAlphas[i] = isPersistent
1157-
? PERSISTENT_ALPHAS[persistentAlphaIdx++ % PERSISTENT_ALPHAS.length]
1158-
: 0.85;
1231+
allocPersistent[i] = allocEnds[i] >= META.max_timestep ? 1 : 0;
11591232
}
1233+
recolorAllocs();
11601234
let dimPersistent = false;
11611235
11621236
// Bucket index for O(bucket_size) hit testing instead of O(n)
@@ -1532,6 +1606,12 @@ def generate_memory_html(
15321606
activeKeys.delete(event.key.toLowerCase());
15331607
});
15341608
1609+
const settingsTrigger = document.getElementById('settings-trigger');
1610+
settingsTrigger.addEventListener('click', function(e) {
1611+
if (e.target.closest('#settings-dropdown')) return;
1612+
this.classList.toggle('open');
1613+
});
1614+
15351615
document.getElementById('hwm-toggle').onchange = function() {
15361616
hwmG.style('display', this.checked ? null : 'none');
15371617
};
@@ -1548,6 +1628,12 @@ def generate_memory_html(
15481628
updateChart(currentTransform);
15491629
};
15501630
1631+
document.getElementById('color-mode').onchange = function() {
1632+
colorMode = this.value;
1633+
recolorAllocs();
1634+
drawCanvas();
1635+
};
1636+
15511637
// --- Feature 1: Search & Filter ---
15521638
const searchInput = document.getElementById('search-input');
15531639
const regexToggle = document.getElementById('regex-toggle');

‎transformer_nuggets/utils/merge_traces.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
import gzip
66
import json
7-
import sys
87
from pathlib import Path
98
from typing import Annotated
109

@@ -52,7 +51,9 @@ def merge_traces(input_paths: list[str], output_path: str) -> None:
5251

5352
@app.command()
5453
def main(
55-
traces: Annotated[list[Path], typer.Argument(help="Input trace files, one per rank, in rank order.")],
54+
traces: Annotated[
55+
list[Path], typer.Argument(help="Input trace files, one per rank, in rank order.")
56+
],
5657
output: Annotated[Path, typer.Option("-o", "--output", help="Output path.")] = Path(
5758
"merged_trace.json.gz"
5859
),

0 commit comments

Comments
 (0)