Repository navigation
perf(pipeline): move the model into the pipeline instead of cloning it - #2299
Closed
ArthurZucker wants to merge 1 commit into
Closed
ArthurZucker wants to merge 1 commit into
ArthurZucker wants to merge 1 commit into
Conversation
`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.
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2279 (
poc/target-encode).What
PipelineTokenizer::try_from(&Tokenizer)cloned the wholeModelWrapperbefore consuming it. While the compact structures are being built from the copy, the vocabulary and merge table are resident twice.A
Tokenizeris normally constructed only to parsetokenizer.jsonand 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 intoPipelineTokenizer::build(&Tokenizer, ModelWrapper). The borrowing impl is unchanged for callers that still need their tokenizer afterwards.Numbers
Peak RSS contributed by the conversion step (
getrusagehigh-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.What this does not fix
The dominant term is
Tokenizer::from_fileitself: it peaks at 198 MB parsing llama-3's 16 MB config into aTokenizerthat 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
tokenizers0.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_itasserts both conversions agree token for token on gpt2, llama-2 and albert. Fulltk-encodesuite passes (391 tests).