Skip to content

perf(pipeline): move the model into the pipeline instead of cloning it - #2299

Closed
ArthurZucker wants to merge 1 commit into
poc/target-encodefrom
perf/pipeline-move-model
Closed

ArthurZucker wants to merge 1 commit into
poc/target-encodefrom
perf/pipeline-move-model

Conversation

@ArthurZucker

Copy link
Copy Markdown
Collaborator

Stacked on #2279 (poc/target-encode).

What

PipelineTokenizer::try_from(&Tokenizer) cloned the whole ModelWrapper before consuming it. While the compact structures are being built from the copy, the vocabulary and merge table are resident twice.

A Tokenizer is normally constructed only to parse tokenizer.json and dropped as soon as the pipeline exists, so that copy is pure waste — and it lands on the peak, which is the number that has to fit in a container.

This adds TryFrom<Tokenizer>, which moves the model across, and factors the shared conversion into PipelineTokenizer::build(&Tokenizer, ModelWrapper). The borrowing impl is unchanged for callers that still need their tokenizer afterwards.

// before — one full copy of the vocab and merges, thrown away immediately
let pipe = PipelineTokenizer::try_from(&Tokenizer::from_file(path)?)?;

// after — the model is moved
let pipe = PipelineTokenizer::try_from(Tokenizer::from_file(path)?)?;

Numbers

Peak RSS contributed by the conversion step (getrusage high-water sampled before and after). The live heap of the finished pipeline is identical either way — this changes only what is resident while it is built.

model conversion Δpeak, before after total peak
gpt2 +11.8 MB +1.7 MB 39.8 → 26.9 MB
llama-2 +11.6 MB +4.1 MB 37.4 → 29.9 MB
llama-3 +53.2 MB +25.9 MB 251.5 → 224.3 MB
mistral-nemo +37.9 MB +11.7 MB 143.6 → 117.5 MB
deepseek-v4 +45.8 MB +23.1 MB 110.9 → 88.4 MB

What this does not fix

The dominant term is Tokenizer::from_file itself: it peaks at 198 MB parsing llama-3's 16 MB config into a Tokenizer that then holds 37 MB — a 5× overshoot inside the legacy loader, before the pipeline is involved at all. Separate problem, not addressed here.

For context on why this matters: measuring loaded-tokenizer memory by RSS delta is misleading in exactly this situation. RSS never falls, so a loader that builds an intermediate and discards it gets billed for memory it no longer holds. Measured by live heap, the pipeline holds less than every other engine tested (llama-3: 12 MB vs 29 MB for tokenizers 0.23.1, 48 MB for kitoken, 84 MB for gigatoken) — the MPHF does what it claims; the peak was the loader's fault.

Test

A moved-from field left silently empty would still build a working pipeline that just emits wrong ids, so moving_the_tokenizer_converts_identically_to_borrowing_it asserts both conversions agree token for token on gpt2, llama-2 and albert. Full tk-encode suite passes (391 tests).

`PipelineTokenizer::try_from(&Tokenizer)` cloned the whole `ModelWrapper`
before consuming it, so while the compact structures were being built the
vocabulary and merge table were resident twice. A `Tokenizer` is normally
constructed only to parse `tokenizer.json` and dropped immediately after, so
that copy is pure waste -- and it lands on the peak, which is the figure that
has to fit in a container.

Adds `TryFrom<Tokenizer>` which moves the model across, and factors the shared
conversion into `PipelineTokenizer::build`, taking the model already owned.
The borrowing impl is unchanged for callers that still need their tokenizer.

Peak RSS added by the conversion step (getrusage high-water, measured before
and after; live heap of the finished pipeline is identical either way):

    gpt2          +11.8 MB -> +1.7 MB    total 39.8 -> 26.9 MB
    llama-2       +11.6 MB -> +4.1 MB    total 37.4 -> 29.9 MB
    llama-3       +53.2 MB -> +25.9 MB   total 251.5 -> 224.3 MB
    mistral-nemo  +37.9 MB -> +11.7 MB   total 143.6 -> 117.5 MB
    deepseek-v4   +45.8 MB -> +23.1 MB   total 110.9 -> 88.4 MB

Note what this does *not* fix: the dominant term is `Tokenizer::from_file`
itself, which peaks at 198 MB parsing llama-3's 16 MB config into a structure
that then holds 37 MB. That is a separate problem in the legacy loader.

A moved-from field left silently empty would still build a working pipeline
that emits wrong ids, so the new test asserts both conversions agree token for
token on gpt2, llama-2 and albert.
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants