For those interested in playing with or doing research using model internals, the TransformerLens [1] project appears to be the leading open-source tooling in this area. It allows for loading dozens of different models, adding hooks, displaying activations in a format compatible with CircuitsVis, and other (mechanistic) interpretability work.

[1] https://github.com/neelnanda-io/TransformerLens

[2] https://github.com/alan-cooney/CircuitsVis

(2020)

I gave a small talk on how to really push using hooks for logging intermediate values (including capturing gradients from torch & fx scripted modules) that may be useful: https://static.sched.com/hosted_files/pytorch2023/40/Interme...

FYI, for anyone interesting in creating and using hooks to better understand what's happening in your model, I created a free lesson covering that:

https://course.fast.ai/Lessons/lesson17.html