OpenMOSS / Language-Model-SAEs

For OpenMOSS Mechanistic Interpretability Team's Sparse Autoencoder (SAE) research.
21 stars 3 forks source link

[Proposal] Accelerate Inference in TransformerLens #11

Open Hzfinfdu opened 3 weeks ago

Hzfinfdu commented 3 weeks ago

The main bottleneck of SAE training lies in activation gen. It can be annoying when we try to work with larger models.

Try to accelerate TL inference, especially attn forward. What are some possible options? FlashAttn2 or VLLM or something?

Since we usually do not cache Q K V, attn forward can be replaced with some faster alternatives.