openmoss / language-model-saes Goto Github PK
View Code? Open in Web Editor NEWFor OpenMOSS Mechanistic Interpretability Team's Sparse Autoencoder (SAE) research.
For OpenMOSS Mechanistic Interpretability Team's Sparse Autoencoder (SAE) research.
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.
I have noticed 2 problems in current frontend + backend service.
AttributeError: 'PreTrainedTokenizerFast' object has no attribute 'byte_decoder'.
Simple truncating might not work. It may fail to catch longer range dependency.
I propose to make some changes in the frontend. What about a preview (shorter local context) which can be expanded to full context? @dest1n1s
The current activation store implementation has some drawbacks. Maybe we need to add some new features for streaming activation store and make some optimizations. Below I list some details.
IMO pretraining data should be packed and SFT data should be sorted by length and batched with post paddings. Activations in the residual stream of s should be ignored in SAE training. I believe this is better fitted to real-world distribution.
We need to add into the configuration to configure this.
It is much easier for people (who may be new to mechanistic interpretability) to get started with detailed tutorial and documentation. Currently this project lacks documentation and comments in many modules. We should raise the documentation coverage, to ensure detailed explanation of every part of our library.
Furthermore, we should consider building a static documentation site with tools like MkDocs. This helps people get an overview of the usage of the library without actually downloading it.
Automatic testing is fundamental to keep a collaborative developed project from endless bugs corrupting modules that originally work. As for a deep learning library, always running the whole training or analyzing process from the outermost can consume lots of time and computational resources. Minor bugs may also not be triggered in a fixed training setting. Thus, it's necessary to test at different levels to ensure proper functioning as much as possible.
I propose adding the following 4 categories of testing:
Continuous Integration (CI) with GitHub workflows should also be added to run testing on every push/PR. PRs should not be merged unless all tests are passed.
We aim to build a really useful infrastructure for SAE research.
Maybe we need to open source our SAEs with a huggingface-type interface. This may require some sort of cloud service for storage or stuff like that? @dest1n1s
This has lower priority since this feature is only of use after we trained considerably good SAEs on larger language models.
We can publish this library on PyPI so that people can use this package simply using pip install lm-saes
! However, before this we should first get this library well-tested and well-documented:
Besides, even if we publish this library, people may still need to clone this repository for using the visualization tools. Perhaps we can publish docker images of the backend & frontend of visualization.
What about doing all stuff with relative path? e.g. saving model path in lm_config.json.
I will revisit this later
TransformerLens already supports stop at a given layer. We can utilize this feature to remove unused calculation. This can be applied at any cases that uses run_with_cache.
Same as Early Stop, what if we do not load these unused weights at all? This may save GPU memory when training early and middle layers.
A natural approach to faster SAE training is data parallel. Maybe we can just simply use DDP to make 8 copies of the TL model to yield activation and synchronize SAE gradients. This may help accelerate activation gen, which is the speed bottleneck for larger LMs.
This may not work on larger size models, say 70B models. Maybe the ultimate solution is a producer-consumer design pattern. Let's leave this for later.
A declarative, efficient, and flexible JavaScript library for building user interfaces.
๐ Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
An Open Source Machine Learning Framework for Everyone
The Web framework for perfectionists with deadlines.
A PHP framework for web artisans
Bring data to life with SVG, Canvas and HTML. ๐๐๐
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
Some thing interesting about web. New door for the world.
A server is a program made to process requests and deliver data to clients.
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
Some thing interesting about visualization, use data art
Some thing interesting about game, make everyone happy.
We are working to build community through open source technology. NB: members must have two-factor auth.
Open source projects and samples from Microsoft.
Google โค๏ธ Open Source for everyone.
Alibaba Open Source for everyone
Data-Driven Documents codes.
China tencent open source team.