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).
Specifically, one particular otherwise excellent model I use has an alignment problem (sycophancy) that I've isolated to a specific layer. I can nuke the layer with lora and the behaviour stops - but I'm not sure what else I'm nuking in the process. I'm quite new at this so I'd love any advice. Thank you!
https://hfviewer.com/
The main difference appears to be how the graph gets made. From what I can tell they analyze the config on the server (their page says the first request can take a few minutes). I actually build the model on PyTorch's meta device and run a fake forward pass through it, so I get the real execution order and the tensor shapes at every module. That's what the animated replay and the compute/KV cache numbers are based on.