fix: reduce Trie and newmm peak memory; add tcc_pos_array() - #1323
Merged
Conversation
Copilot
AI
changed the title
[WIP] Fix high memory usage in newmm tokenization
fix(newmm): reduce peak memory in safe_mode=False tokenization
Mar 9, 2026
2 tasks
- trie.py: lazy Node.children (None instead of {}); saves ~3.25 MiB
(52 K leaf nodes × 64 B) for the default word-dict trie
- trie.py: remove self.words set; add _word_count int; saves ~2 MiB
of set-object overhead per Trie; rewrite __contains__, __iter__,
__len__, add, remove to use trie traversal / DFS
- trie.py: __iter__ uses recursive DFS with shared mutable prefix to
avoid O(k^2) list copies
- tcc_p.py: add tcc_pos_array() returning bytearray indexed by position;
O(1) array lookup with len(text)+1 bytes vs. ~80 B/position in a set
- newmm.py: use tcc_pos_array() instead of tcc_pos(); replace all four
'in valid_poss' set-membership tests with valid_poss[i] array indexing
- tests: add _word_count accuracy tests (re-add, double-remove, iter);
add tcc_pos_array edge-case tests (empty/None/int input, array length)
Co-authored-by: bact <128572+bact@users.noreply.github.com>
Copilot
AI
changed the title
fix(newmm): reduce peak memory in safe_mode=False tokenization
fix: reduce Trie and newmm tokenizer peak memory usage
Mar 10, 2026
Add explicit `if parent.children is not None:` guard before `del parent.children[ch]` in Trie.remove() so mypy can narrow the Optional[dict] type. The guard is always true at runtime but mypy cannot infer this from the path-traversal logic above. Add CHANGELOG.md entry for the Trie memory and tcc_pos_array changes. Co-authored-by: bact <128572+bact@users.noreply.github.com>
Copilot
AI
changed the title
fix: reduce Trie and newmm tokenizer peak memory usage
fix: reduce Trie and newmm peak memory; add tcc_pos_array()
Mar 10, 2026
Member
|
Peak memory at first run is now 76.7 MB, reduced from 81.9 MB (-5.2MB or 6.3%). |
Consolidate memory optimization details for Trie and TCC lookups.
bact
marked this pull request as ready for review
March 10, 2026 06:14
|
bact
added
enhancement
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.
Trie construction dominates newmm's peak memory since
word_dict_trie()became lazy-loaded via@lru_cache(PR #1186) — the first tokenization call now pays the full trie-build cost. Three redundant allocations inTrieand one in the tokenizer hot path are addressed.What do these changes do
Trie.Node.childrenlazy allocation —Noneby default; dict created on first child insert, reset toNonewhen last child is deleted. Eliminates ~52 K empty dicts × 64 B ≈ 3.25 MiB for the default word-dict trie.self.words: set[str]→_word_count: int— the set duplicated every word string and added ~2 MiB of set-object overhead. A plain counter is incremented/decremented inadd()/remove()onendflag transitions.__contains__traverses the trie;__iter__is a recursive DFS with a shared mutable prefix (avoids O(k²) list copies);__len__returns the counter.tcc_pos_array(text) → bytearray— new function intcc_p.pyalongsidetcc_pos(). Returnsbytearray(len(text)+1)witharr[i] = 1at each valid TCC boundary.newmm._onecutuses this instead ofset[int]; all fourin valid_possset-membership tests becomevalid_poss[i]array-index lookups.union-attrfix —Trie.remove()prune loop adds an explicitif parent.children is not None:guard beforedel parent.children[ch]to narrowOptional[dict]for the type checker (always true at runtime).What was wrong
Trie.Nodeallocated an empty{}dict at construction, even leaf nodes (52 K of them in the default trie).self.wordsstored a full copy of every word string in aset, doubling string-reference overhead and adding ~2 MiB of set bookkeeping.tcc_pos()returned aset[int], requiring ~80 B per position;newmmthen did Pythonin-set tests in the inner loop.Trie.remove()accessedparent.children[ch]without narrowingOptional[dict], causing a mypyunion-attrerror.How this fixes it
Lazy
childrenavoids allocating dicts that are never populated. Replacingwordswith_word_countremoves the per-word string duplication.tcc_pos_arraytrades thesetfor a flatbytearray— O(1) indexed lookup with storage proportional to text length in characters rather than cluster count. Theif parent.children is not None:guard gives mypy the type narrowing it needs.Your checklist for this pull request
💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.