llama-1B-5B / README.md
Local-Axiom-AI's picture
Update README.md
8ee2b58 verified
|
Raw
History Blame Contribute Delete
10.2 kB
---
license: apache-2.0
language:
- en
pipeline_tag: text-generation
base_model: allenai/OLMoE-1B-7B-0125-Instruct
tags:
- gguf
- olmoe
- moe
- expert-pruning
- model-compression
- research
---
# OLMoE-1B-7B-0125-Instruct — 30% Expert Pruned
This is an experimental version of
[`allenai/OLMoE-1B-7B-0125-Instruct`](https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct)
with approximately **30% of its routed experts physically removed**.
The original model has 64 experts in each of its 16 mixture-of-experts layers.
This version removes 19 experts independently from every layer, leaving 45
experts per layer. The original top-8 routing is preserved.
The purpose of this release is to explore a simple question:
> How much of a mixture-of-experts model can be removed when experts are ranked
> by usage, contribution, and functional redundancy rather than parameter
> magnitude alone?
This is a research artifact. Removing nearly one-third of the expert pool is
an aggressive intervention and may remove capabilities that appear only in
particular domains or prompt distributions.
In testing against the unpruned base model, this 30%-cut version showed an
approximately **6% average performance decrease** across the same MMLU,
ARC-Challenge, HellaSwag, and GSM8K benchmark screen used during the pruning
experiments. This provides an initial measure of the quality tradeoff: about
30% of the expert pool was removed for an observed average benchmark loss of
about 6%.
## What was removed
| Property | Original OLMoE | 30% expert-pruned model |
|---|---:|---:|
| MoE layers | 16 | 16 |
| Experts per layer | 64 | 45 |
| Total expert modules | 1,024 | 720 |
| Expert modules removed | — | 304 |
| Expert pool removed | — | 29.69% |
| Experts active per token | 8 | 8 |
| Context length | 4,096 | 4,096 |
Nineteen experts were removed from each layer. Expert roles are layer-specific,
so candidates were selected independently in every layer rather than removing
the same expert indices throughout the network.
For each removed expert, both the expert MLP and its corresponding router row
were deleted. The router therefore produces 45 logits instead of 64, and its
softmax renormalizes probability over the experts that remain.
Attention layers, embeddings, normalization layers, and retained expert
tensors are not reduced by the 30% expert cut. This is why the final file is
smaller by slightly less than 30%.
## The pruning idea
Mixture-of-experts models contain a large pool of expert MLPs, but only a small
subset is selected for each token. Some experts are used frequently across
many prompts, while others are activated rarely or behave similarly to other
experts in the same layer.
The pruning hypothesis is that part of this expert pool is redundant. If the
least important experts can be identified, their parameters and router entries
can be removed without changing the rest of the network.
The difficult part is deciding what "least important" means. Routing frequency
alone is not sufficient: an expert may be rarely used globally while remaining
important for code, mathematics, multilingual text, or structured output.
## How experts were ranked
The pruning process used multiple signals for every expert in every layer:
| Signal | Purpose |
|---|---|
| Routed-token frequency | Measures how often an expert is selected |
| Selected router probability | Measures how strongly the router prefers it |
| Gate-weighted output magnitude | Estimates the expert's contribution after routing |
| Raw output magnitude | Measures the scale of the expert's response |
| Functional similarity | Identifies experts with behavior already represented by others |
The signals were combined into a layer-local importance score. Low usage was
therefore not enough by itself to make an expert a removal candidate. An expert
also needed to have low measured contribution or substantial functional
overlap with experts that would remain.
Calibration examples covered general knowledge, science, commonsense,
mathematics, code, multilingual text, and structured formats. Including
several domains reduces the risk of classifying a narrow specialist as unused
only because it is quiet on ordinary English prompts.
The process for each MoE layer was:
1. Collect routing and expert-output statistics on the calibration mix.
2. Compare expert behavior on shared hidden states.
3. Protect strong contributors observed in each calibration domain.
4. Rank the remaining candidates by multi-signal importance.
5. Remove the lowest-ranked 19 experts and their router rows.
6. Preserve the original top-8 routing over the 45 retained experts.
## Deletion rather than weight averaging
This model uses structural expert deletion. Removed experts are not replaced
with averaged MLP weights.
Two expert MLPs may produce similar outputs while using different internal
hidden-unit permutations. Directly averaging their parameters can therefore
damage both functions. Functional similarity is used as evidence of
redundancy, not as justification for naive parameter averaging.
Because the model is physically compacted, the removed expert tensors no
longer occupy storage or resident memory. This differs from router masking,
where all expert parameters remain in the checkpoint even if some experts can
never be selected.
## What the 30% cut changes
The expert pool is reduced from 64 to 45 choices in every layer. Each token is
still routed to 8 experts, so a larger fraction of the remaining pool is active
for each token:
- Original: 8 of 64 experts, or 12.5% of the layer's expert pool.
- Pruned: 8 of 45 experts, or 17.8% of the layer's expert pool.
This release primarily targets model size and resident memory. It does not
promise a proportional reduction in inference FLOPs because the number of
active experts per token remains unchanged.
Pruning can also change routing decisions even when every retained tensor is
copied exactly. Removing router outputs changes the softmax distribution and
may change which experts enter the top-8 selection, particularly when router
scores are close.
## Released artifact
| File | Experts/layer | Quantization | Size |
|---|---:|---:|---:|
| `OLMoE-1B-7B-0125-Instruct-30pct-Q4_K_M.gguf` | 45 | Q4_K_M | 2.84 GiB |
The local upstream Q4_K_M conversion is 3.92 GiB. The 30%-pruned Q4_K_M file is
approximately **27.54% smaller**. The remaining difference from 30% comes from
the model parameters that are not part of the expert MLP pool.
The GGUF metadata reports 16 OLMoE blocks, 45 experts per block, 8 active
experts per token, and a 4,096-token context length.
## Running with llama.cpp
```bash
llama-cli \
-m OLMoE-1B-7B-0125-Instruct-30pct-Q4_K_M.gguf \
-cnv -ngl 99 -c 4096
```
To run an OpenAI-compatible local server:
```bash
llama-server \
-m OLMoE-1B-7B-0125-Instruct-30pct-Q4_K_M.gguf \
-ngl 99 -c 4096 --port 8080
```
Reduce `-ngl` if the model does not fit in available GPU memory. The artifact
contains the upstream chat template in its GGUF tokenizer metadata.
## Testing status
This exact 30%-pruned Q4_K_M artifact was successfully loaded and used for
generation with `llama.cpp` on an NVIDIA RTX 3090.
The pruned model was also compared with the unpruned base model using the same
task benchmark suite used elsewhere in the pruning experiment:
- MMLU;
- ARC-Challenge;
- HellaSwag; and
- GSM8K.
Across that suite, the 30%-cut model had an approximately **6% average drop**
relative to the base model. This is an aggregate summary, not a claim that
every individual task decreased by exactly 6%.
The detailed result file is no longer available, so per-task scores, sample
counts, and statistical uncertainty cannot be reproduced from this release.
The aggregate result is included for transparency but should be treated as a
small benchmark screen rather than a publication-grade evaluation.
## Limitations
- Removing 19 of 64 experts per layer is aggressive and may reduce quality.
- Calibration can miss specialists outside its sampled domains.
- Rare-language, code, structured-output, or other narrow capabilities may be
affected even when general chat behavior appears reasonable.
- Top-8 routing remains unchanged, so active-compute savings are modest.
- Q4_K_M quantization introduces an additional source of quality loss beyond
expert pruning.
- The retained aggregate benchmark result shows an approximately 6% average
decrease, but the deleted detailed results prevent per-task analysis or
independent reproduction of that figure from this repository.
- The model inherits the upstream model's biases and safety limitations.
This model should be evaluated on representative data before practical use. It
is best treated as an experiment in structural MoE compression rather than a
drop-in replacement for the upstream checkpoint.
## Upstream model and license
This work is derived from
[`allenai/OLMoE-1B-7B-0125-Instruct`](https://huggingface.co/allenai/OLMoE-1B-7B-0125-Instruct)
and retains its Apache 2.0 license. See the upstream model card for training
data, post-training, safety information, and citations. Architecture details
are available in the [OLMoE paper](https://arxiv.org/abs/2409.02060).
## Conclusion
Removing approximately 30% of the model's expert pool while observing only an
approximately 6% average drop across the tested benchmarks suggests that a
meaningful amount of expert capacity may be redundant for the evaluated task
distribution. Although the benchmark screen is too limited to establish broad
quality parity, the difference between the scale of the structural reduction
and the measured performance loss warrants further research into MoE pruning.
More broadly, these results support research into less wasteful MoE designs:
models that preserve useful specialization while requiring fewer stored
parameters, less resident memory, and less underused expert capacity. Larger
and more diverse evaluations, improved specialist detection, iterative
pruning, and post-pruning recovery or distillation could help determine how
much redundancy can be removed without disproportionately harming rare-domain
capabilities.