Thank you for the wonderful repo and the opensource on llama3, it's clear and easy to read.
I think the common practise of training a SAE written in OpenAI/Anthophic is on the residual streams, or to say the outputs of the MLPs.
The repo is getting the activations through output_hidden_states=True provided in transformers. But I see that hidden_states got with this factor are the inputs of the transformers blocks.
I do not current know whether the residual streams or the hidden_states give better results for analysis. Or I was understanding the code in a wrong way. Can you help?

Thank you very much.

For reference
https://github.com/huggingface/transformers/blob/a5c642fe7a1f25d3bdcd76991443ba6ff7ee34b2/src/transformers/models/llama/modeling_llama.py#L859