flex-general-2048

flex-general-2048

A pruned and partially distilled variant of allenai/Flex-math-2x7B-1T with a variable-width expert MLP. Expert 1 has been pruned from the full 11,008 intermediate size down to 2048 (19% of original width), then partially recovered via knowledge distillation.

Unlike the math-calibrated variant, this model's pruning was calibrated on general-purpose data — meaning importance scores were computed on a broad data mix rather than math-specific data. 58% of the top-2048 most important neurons differ between the two calibration approaches.

For full details, see the blog post.

How to Use

This repo includes a modeling_pruned_flex_olmo.py file that handles the variable-width expert architecture. Just load with trust_remote_code=True and it works like any other HuggingFace model:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("hbfreed/flex-general-2048", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("allenai/Flex-math-2x7B-1T")

inputs = tokenizer("Hello, world!", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

The tokenizer is the same as the base model's.

How It Was Made

1. Structured pruning: Neuron importance scores were computed on general-purpose data. The least important neurons in Expert 1's gate/up/down projections were removed, reducing intermediate size from 11,008 to 2048.
2. Partial knowledge distillation: The pruned model was partially retrained (~20k steps) using logprobs from the full-sized teacher model. Training was stopped early — the general-calibrated model converged slower and to a higher loss than the math-calibrated variant.

Related Models

License

Apache 2.0 (same as base model)