https://github.com/milescb/sparseffn
https://github.com/milescb/sparseffn
Last synced: 10 months ago
JSON representation
- Host: GitHub
- URL: https://github.com/milescb/sparseffn
- Owner: milescb
- Created: 2025-03-01T00:23:02.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-03-01T02:07:36.000Z (over 1 year ago)
- Last Synced: 2025-03-01T03:18:45.810Z (over 1 year ago)
- Language: Python
- Size: 4.88 KB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# Sparse FFN
Train a small network to learn the output of feed-forward (FFN) blocks of an LLM. This distilation technique can then be used to predict the sparsity of the FFN blocks in LLMs.
## Setup the environment
```
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```
## Train the model
Training can be accomplished simply by running the `train.py` script:
```
python train.py -d .pt -od -o -bs -lr -e
```
## Use trained model
To test and make plots, simply run:
```
python test.py --data -mod .pt
```
To use trained models, run:
```python
from ffnsparse.query import predict
pred = predict(input_vector, "model_name.pt")
```
where `input_vector` is the embedded input to the FFN block, and `"model_name.pt"` is the path to the trained model.