Hi! Not from the config alone, the config just tells transformers which classes to build. I let it actually build the model, but on PyTorch's meta device, so every module and parameter exists with its real shape and dtype but no memory behind it. That gives me the true nn.Module tree (the same one print(model) would show), and I hash repeated subtrees so 36 identical layers show as one stack with a ×36 badge.
Then I run a fake forward pass (dummy inputs, also on the meta device) with a forward hook on every module. Most ops do shape inference fine without real data, so that records the execution order and every module's input/output shapes. Both things you're asking about are already in the UI:
- Tensor shapes: click any module and the inspector shows its traced input/output (e.g. [1 batch × 7 seq × 4096 hidden], the labels come from matching dim values against config) plus its weight shapes ([151936 vocab × 4096 hidden]). The flow-replay HUD shows the shape transformation at each step.
- Per-layer parameter counts: every node shows its param count and share of the model, there's a treemap of children by params, and a "cost" lens that switches the whole map to compute (MACs), activation memory, or KV cache — with sequence length as a slider.
Params/dtypes are cross-checked against the safetensors headers (fetched via HTTP range requests, no weight download).
I saw a similar page a month back. The design aesthetics are bit different
https://hfviewer.com/
How are you generating the architecture graph from model configs, and do you plan to surface tensor shapes or layer-level parameter counts as well?
Hi! Not from the config alone, the config just tells transformers which classes to build. I let it actually build the model, but on PyTorch's meta device, so every module and parameter exists with its real shape and dtype but no memory behind it. That gives me the true nn.Module tree (the same one print(model) would show), and I hash repeated subtrees so 36 identical layers show as one stack with a ×36 badge.
Then I run a fake forward pass (dummy inputs, also on the meta device) with a forward hook on every module. Most ops do shape inference fine without real data, so that records the execution order and every module's input/output shapes. Both things you're asking about are already in the UI:
- Tensor shapes: click any module and the inspector shows its traced input/output (e.g. [1 batch × 7 seq × 4096 hidden], the labels come from matching dim values against config) plus its weight shapes ([151936 vocab × 4096 hidden]). The flow-replay HUD shows the shape transformation at each step. - Per-layer parameter counts: every node shows its param count and share of the model, there's a treemap of children by params, and a "cost" lens that switches the whole map to compute (MACs), activation memory, or KV cache — with sequence length as a slider.
Params/dtypes are cross-checked against the safetensors headers (fetched via HTTP range requests, no weight download).
Visually inspecting how a particular model is architected is quite interesting. Thank you for coming up with this!
Love it, I thought it only shows the "guts" of a model, turns out it also estimates the cost of serving that model.
extremely cool. I’ve had Claude walk model architectures to debug Loras and fine tunes, but this is delightful
Thanks for good vibes. This is awesome