Skip to content

Model trimming

trimbed.model_trim

Trimming a model's embedding table (and output head) down to just the kept tokens.

TOKEN_ID_ATTRIBUTES module-attribute 

TOKEN_ID_ATTRIBUTES = (
    "bos_token_id",
    "eos_token_id",
    "pad_token_id",
    "sep_token_id",
    "cls_token_id",
    "unk_token_id",
    "mask_token_id",
    "decoder_start_token_id",
    "forced_bos_token_id",
    "forced_eos_token_id",
)

Config attributes holding a single token id that must follow the remap.

Which of them a checkpoint sets varies: Qwen3 has only eos_token_id, T5 adds decoder_start_token_id and pad_token_id, BERT sets pad_token_id only.

trim_model 

trim_model(
    model: PreTrainedModel,
    remap: IdRemap,
    config: EmbeddingTrimConfig | None = None,
) -> ModelReport

Shrink a model's vocabulary-sized tensors (embeddings, lm head) down to just the kept tokens.

Rows are gathered before resizing and written back afterwards, because resize_token_embeddings keeps only the first n of the original rows, which is almost never the set we want (we select the right ones non-contiguously).

Parameters:

Name Type Description Default
model PreTrainedModel

A loaded transformers model, e.g. a Qwen3ForCausalLM or a BertForMaskedLM.

 required
remap IdRemap

Mapping from surviving old ids to contiguous new ids. Its new_to_old is used directly as the gather index.

 required
config EmbeddingTrimConfig | None

Embedding-trimming settings. Defaults are used when None.

 None

Returns:

Type Description
ModelReport

Statistics describing the change, e.g. 151,936 embedding rows down to 32,000.

Raises:

Type Description
ValueError

If an id to keep lies outside the existing embedding matrix, which means the remap and the checkpoint disagree about the vocabulary.

View source on GitHub: src/trimbed/model_trim.py lines 69–176