| """Configuration for GLEAN-pruned OLMoE: variable-width, variable-count experts. |
| """ |
|
|
| from transformers.models.olmoe.configuration_olmoe import OlmoeConfig |
|
|
|
|
| class PrunedOlmoeConfig(OlmoeConfig): |
| """OlmoeConfig plus a per-(layer, expert) width table. |
| |
| ``expert_widths[l]`` lists the SwiGLU intermediate width of each surviving |
| expert in decoder layer ``l``, in expert order. Lists are ragged: layers |
| may keep different numbers of experts (deleted experts simply don't |
| appear — the router in layer ``l`` has ``len(expert_widths[l])`` rows), |
| and each width may differ (multiples of the GEMM block size, 128, for |
| variable-MegaBlocks execution). ``None`` means an unpruned model |
| (uniform ``num_experts`` × ``intermediate_size``). |
| |
| The inherited ``num_experts`` / ``intermediate_size`` keep their ORIGINAL |
| (pre-pruning) values for provenance; the width table is authoritative for |
| the built architecture. |
| """ |
|
|
| model_type = "pruned_olmoe" |
|
|
| def __init__(self, expert_widths: list[list[int]] | None = None, **kwargs): |
| super().__init__(**kwargs) |
| self.expert_widths = expert_widths |
|
|
