remove legacy tokenizer system - #1989
Merged
Merged
Conversation
Signed-off-by: dimapihtar <dpihtar@gmail.com>
dimapihtar
marked this pull request as draft
January 20, 2026 11:48
Signed-off-by: dimapihtar <dpihtar@gmail.com>
dimapihtar
marked this pull request as ready for review
February 9, 2026 16:49
Contributor
Author
|
/ok to test d98abbc |
Contributor
📝 WalkthroughWalkthroughThis PR consolidates the tokenizer infrastructure by removing legacy tokenizer implementations (BERT, GPT2 BPE, multimodal) and their factory utilities, while introducing a unified Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@src/megatron/bridge/training/tokenizers/tokenizer.py`:
- Around line 47-57: The TikTokenizer branch sets kwargs["vocab_size"] twice
which risks overwriting a valid value with None; in the block where
config.tokenizer_type == "TikTokenizer" remove the duplicate unconditional
assignment of kwargs["vocab_size"] and only set kwargs["vocab_size"] when
config.vocab_size is not None (i.e., keep the existing guarded assignment if
config.vocab_size is present and do not reassign later). Target the tokenizer.py
TikTokenizer branch and the kwargs dict modifications (keys "vocab_size",
"num_special_tokens", "special_tokens", "pattern", "chat_template").
- Around line 15-27: The function build_tokenizer currently clobbers
caller-supplied kwargs by doing kwargs = {}, so change it to preserve incoming
kwargs and merge or update it with config-derived values instead of resetting;
locate the build_tokenizer function and remove the kwargs = {} assignment, then
populate tokenizer_library and any config-derived options by updating kwargs
(e.g., kwargs.update({...}) or creating a new dict merged = {**kwargs,
**config_values}) before passing them into MegatronTokenizer initialization to
ensure caller kwargs are honored.
In `@tests/unit_tests/training/test_tokenizer.py`:
- Around line 125-168: Add the `@pytest.mark.unit` marker to
test_hf_tokenizer_as_local_path_object and remove the network call to
AutoTokenizer.from_pretrained; instead create a tiny local tokenizer with the
tokenizers library, save it to tmp_path, and use that path in the existing
TokenizerConfig so build_tokenizer loads from disk. Specifically, replace the
AutoTokenizer.from_pretrained("bert-base-uncased") step with constructing a
small Tokenizer (e.g., a WordLevel or BPE model), train/populate it with a tiny
vocabulary or trainer, call tokenizer.save_pretrained or the tokenizers
equivalent to write tokenizer files into local_model_path, then proceed with the
existing TokenizerConfig(tokenizer_model=local_model_path) and assertions; keep
assertions comparing build_tokenizer output to the locally saved tokenizer and
verifying files exist.
Comment on lines
+15
to
+27
| def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer: | ||
| """Initialize tokenizer from megatron.core.tokenizers based on the provided configuration. | ||
|
|
||
| Args: | ||
| tokenizer_instance: Tokenizer instance with a `tokenize` method | ||
| default: Fallback value if computation fails (True for SentencePiece, False for others) | ||
|
|
||
| Returns: | ||
| bool: True if the tokenizer is space-sensitive, False otherwise | ||
|
|
||
| Example: | ||
| # A space-sensitive tokenizer (e.g., many BPE tokenizers): | ||
| # tokenize("x y") -> [87, 331] | ||
| # tokenize("x") + tokenize("y") -> [87, 379] # Different! | ||
|
|
||
| # A non-space-sensitive tokenizer would produce the same result | ||
| """ | ||
| try: | ||
| test_tokens_with_space = tokenizer_instance.tokenize("x y") | ||
| test_tokens_concat = tokenizer_instance.tokenize("x") + tokenizer_instance.tokenize("y") | ||
| return test_tokens_with_space != test_tokens_concat | ||
| except Exception: | ||
| # If tokenization fails for any reason, use the default | ||
| return default | ||
|
|
||
|
|
||
| class MegatronLegacyTokenizer(MegatronTokenizerCore): | ||
| """Base tokenizer class, extending the MegatronTokenizer from megatron core. | ||
|
|
||
| This class provides a common interface for various tokenizers used within the NeMo framework. | ||
| """ | ||
|
|
||
| def __init__(self, *args, **kwargs): | ||
| # Set legacy attribute | ||
| self.legacy = True | ||
| super().__init__(*args, **kwargs) | ||
|
|
||
| def __call__(self, *args, **kwargs): | ||
| """Makes the tokenizer instance callable, synonym for `tokenize`.""" | ||
| return self.tokenize(*args, **kwargs) | ||
|
|
||
| def text_to_ids(self, text: str) -> list[int]: | ||
| """Converts text to a list of token IDs.""" | ||
| return self.tokenize(text) | ||
|
|
||
| @property | ||
| def eod_id(self): | ||
| """ID for the end-of-document token.""" | ||
| return self.eod | ||
|
|
||
| @property | ||
| def bos_id(self): | ||
| """ID for the beginning-of-sentence token.""" | ||
| return self.bos | ||
|
|
||
| @property | ||
| def eos_id(self): | ||
| """ID for the end-of-sentence token.""" | ||
| return self.eos | ||
|
|
||
| @property | ||
| def mask_id(self): | ||
| """ID for the mask token.""" | ||
| return self.mask | ||
|
|
||
|
|
||
| def build_tokenizer(tokenizer_config: TokenizerConfig, **kwargs) -> MegatronLegacyTokenizer | MegatronTokenizer: | ||
| """Initialize tokenizer based on the provided configuration. | ||
|
|
||
| This function serves as a factory to instantiate various tokenizer types | ||
| supported by NeMo Framework, such as BERT, GPT2, SentencePiece, HuggingFace, etc. | ||
| It also handles padding the vocabulary size to be GPU-friendly. | ||
|
|
||
| Args: | ||
| tokenizer_config (TokenizerConfig): Configuration object specifying the tokenizer | ||
| config (TokenizerConfig): Configuration object specifying the tokenizer | ||
| type, paths to vocab/model files, and other | ||
| tokenizer-specific settings. | ||
| **kwargs: Additional keyword arguments that might be specific to certain tokenizers | ||
| (e.g., passed to HuggingFace AutoTokenizer). These will override any | ||
| kwargs specified in tokenizer_config.hf_tokenizer_kwargs. | ||
|
|
||
| Returns: | ||
| MegatronTokenizer: An instance of the initialized tokenizer. | ||
|
|
||
| Raises: | ||
| NotImplementedError: If the specified tokenizer_type in tokenizer_config is not supported. | ||
| ImportError: If a required library (e.g., transformers for MultimodalTokenizer) is not installed. | ||
| """ | ||
| if get_rank_safe() == 0: | ||
| print("> building {} tokenizer ...".format(tokenizer_config.tokenizer_type), flush=True) | ||
| kwargs = {} | ||
| tokenizer_library = None |
Contributor
There was a problem hiding this comment.
Don’t discard caller‑supplied kwargs.
kwargs = {} drops any kwargs passed into build_tokenizer, which can silently break callers. Preserve them and layer config‑derived values on top.
Suggested fix
- kwargs = {}
+ tokenizer_kwargs = dict(kwargs)
@@
- kwargs["additional_special_tokens"] = config.special_tokens if config.special_tokens else []
+ tokenizer_kwargs["additional_special_tokens"] = config.special_tokens if config.special_tokens else []
@@
- special_tokens = {}
- special_tokens["additional_special_tokens"] = [f"<extra_id_{i}>" for i in range(100)]
- kwargs = special_tokens
+ tokenizer_kwargs["additional_special_tokens"] = [f"<extra_id_{i}>" for i in range(100)]
@@
- kwargs["vocab_file"] = config.vocab_file
- kwargs["merges_file"] = config.merge_file
+ tokenizer_kwargs["vocab_file"] = config.vocab_file
+ tokenizer_kwargs["merges_file"] = config.merge_file
@@
- kwargs.update(config.hf_tokenizer_kwargs)
+ tokenizer_kwargs.update(config.hf_tokenizer_kwargs)
@@
- tokenizer = MegatronTokenizer.from_pretrained(tokenizer_path=tokenizer_path, metadata_path=metadata, **kwargs)
+ tokenizer = MegatronTokenizer.from_pretrained(
+ tokenizer_path=tokenizer_path,
+ metadata_path=metadata,
+ **tokenizer_kwargs,
+ )🤖 Prompt for AI Agents
In `@src/megatron/bridge/training/tokenizers/tokenizer.py` around lines 15 - 27,
The function build_tokenizer currently clobbers caller-supplied kwargs by doing
kwargs = {}, so change it to preserve incoming kwargs and merge or update it
with config-derived values instead of resetting; locate the build_tokenizer
function and remove the kwargs = {} assignment, then populate tokenizer_library
and any config-derived options by updating kwargs (e.g., kwargs.update({...}) or
creating a new dict merged = {**kwargs, **config_values}) before passing them
into MegatronTokenizer initialization to ensure caller kwargs are honored.
Comment on lines
+47
to
+57
| elif config.tokenizer_type == "TikTokenizer": | ||
| tokenizer_library = "tiktoken" | ||
| tokenizer_path = config.tokenizer_model | ||
| kwargs["chat_template"] = config.chat_template | ||
| if config.tiktoken_pattern: | ||
| kwargs["pattern"] = config.tiktoken_pattern | ||
| if config.vocab_size: | ||
| kwargs["vocab_size"] = config.vocab_size | ||
| kwargs["num_special_tokens"] = config.tiktoken_num_special_tokens | ||
| kwargs["special_tokens"] = config.special_tokens | ||
| kwargs["vocab_size"] = config.vocab_size |
Contributor
There was a problem hiding this comment.
Avoid overwriting TikTokenizer vocab_size with None.
vocab_size is assigned twice; the unconditional assignment can pass None and override a prior conditional value. Remove the duplicate and guard with is not None.
Suggested fix
- if config.vocab_size:
- kwargs["vocab_size"] = config.vocab_size
+ if config.vocab_size is not None:
+ kwargs["vocab_size"] = config.vocab_size
@@
- kwargs["vocab_size"] = config.vocab_size🤖 Prompt for AI Agents
In `@src/megatron/bridge/training/tokenizers/tokenizer.py` around lines 47 - 57,
The TikTokenizer branch sets kwargs["vocab_size"] twice which risks overwriting
a valid value with None; in the block where config.tokenizer_type ==
"TikTokenizer" remove the duplicate unconditional assignment of
kwargs["vocab_size"] and only set kwargs["vocab_size"] when config.vocab_size is
not None (i.e., keep the existing guarded assignment if config.vocab_size is
present and do not reassign later). Target the tokenizer.py TikTokenizer branch
and the kwargs dict modifications (keys "vocab_size", "num_special_tokens",
"special_tokens", "pattern", "chat_template").
Comment on lines
+125
to
+168
| @pytest.mark.timeout(30) | ||
| def test_hf_tokenizer_as_local_path_object(self, tmp_path): | ||
| # Cover the user case where a user has made a local path object of a WIP tokenizer and wants | ||
| # to use that in some megatron model at train time. | ||
|
|
||
| # First as a proxy download a tokenizer from HF and save it to a local path. A user would | ||
| # do this differently by exporting their WIP tokenizer to a local path. | ||
|
|
||
| # 1. Download a common, small tokenizer from the Hub | ||
| # "bert-base-uncased" is a safe choice as it's small and standard. | ||
| model_id = "bert-base-uncased" | ||
| tokenizer = AutoTokenizer.from_pretrained(model_id) | ||
|
|
||
| # 2. Define a local path in the temporary directory | ||
| local_model_path = tmp_path / "my_local_tokenizer" | ||
|
|
||
| # 3. Save the tokenizer to disk | ||
| # This creates tokenizer_config.json, vocab.txt, special_tokens_map.json, etc. | ||
| tokenizer.save_pretrained(str(local_model_path)) | ||
|
|
||
| # 4. Load it back using the local path | ||
| # This simulates the user providing a path to a folder instead of a Hub ID | ||
| cfg = TokenizerConfig( | ||
| tokenizer_type="HuggingFaceTokenizer", | ||
| tokenizer_model=local_model_path, | ||
| hf_tokenizer_kwargs={ | ||
| "trust_remote_code": True, | ||
| "include_special_tokens": True, | ||
| }, | ||
| ) | ||
| loaded_tokenizer = build_tokenizer(cfg) | ||
|
|
||
| # 5. Verify it functions identically | ||
| test_text = "Unit testing is important." | ||
|
|
||
| original_tokens = tokenizer.encode(test_text) | ||
| reloaded_tokens = loaded_tokenizer.tokenize(test_text) | ||
|
|
||
| assert original_tokens == reloaded_tokens | ||
| assert loaded_tokenizer.vocab_size == tokenizer.vocab_size | ||
|
|
||
| # verify that the directory actually contains files (sanity check) | ||
| assert (local_model_path / "tokenizer_config.json").exists() | ||
| assert (local_model_path / "tokenizer.json").exists() |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's locate and examine the test file
fd "test_tokenizer.py" tests/unit_tests/Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 114
🏁 Script executed:
# Check the actual content of the test file at the specified lines
cat -n tests/unit_tests/training/test_tokenizer.py | sed -n '115,180p'Repository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 2745
🏁 Script executed:
# Check the imports at the top of the file to understand what's available
head -50 tests/unit_tests/training/test_tokenizer.pyRepository: NVIDIA-NeMo/Megatron-Bridge
Length of output: 1881
🌐 Web query:
AutoTokenizer.from_pretrained network dependency HuggingFace transformers
💡 Result:
AutoTokenizer.from_pretrained(...) depends on network access only when the required tokenizer files aren’t already available locally. By default, it will try to fetch missing files from the Hugging Face Hub and cache them under ~/.cache/huggingface/hub (unless you change the cache location). [1]
How to avoid any network calls:
- Force offline (no HTTP to the Hub):
- Set
HF_HUB_OFFLINE=1. Thenfrom_pretrained()will only use locally cached files and fail if they’re missing. [2]
- Set
- Per-call: only load local files:
AutoTokenizer.from_pretrained("...or local path...", local_files_only=True)[2]
- Pre-download for offline use, then load from disk:
- Use
huggingface_hub.snapshot_download(...)ahead of time, then pointfrom_pretrained()at the downloaded directory. [1]
- Use
Cache location / controlling where files are stored:
- Hub downloads (models/tokenizers) default to
~/.cache/huggingface/huband can be redirected withHF_HUB_CACHE, or more broadly withHF_HOME. [3]
Sources:
[1] Transformers installation docs (cache setup + snapshot_download) https://huggingface.co/docs/transformers/v4.52.2/installation
[2] Transformers installation docs (offline mode + local_files_only) https://huggingface.co/docs/transformers/v4.49.0/installation
[3] Hugging Face docs (hub cache location + HF_HOME / HF_HUB_CACHE) https://huggingface.co/docs/datasets/main/en/cache
🌐 Web query:
PreTrainedTokenizerFast save_pretrained files created tokenizer_config
💡 Result:
PreTrainedTokenizerFast.save_pretrained(save_dir) always writes a tokenizer_config.json because Transformers needs a persisted “init kwargs” bundle to be able to reconstruct the tokenizer with from_pretrained().
What gets saved (typical cases):
-
Always (both fast & slow tokenizers):
tokenizer_config.json— a JSON dump ofself.init_kwargs(plusinit_inputswhen present) and an addedtokenizer_classfield; it is written unconditionally bysave_pretrained(). [1]special_tokens_map.json— serialized special-token settings. [1]
-
Fast tokenizers (
tokenizer.is_fast == True):tokenizer.json— the “full tokenizer” single-file format (FULL_TOKENIZER_FILE = "tokenizer.json"). [1]
-
Slow/legacy tokenizers (or when saved in legacy format):
- Model-specific vocab files (e.g.,
vocab.json,merges.txt,vocab.txt, etc.) and possiblyadded_tokens.jsonwhen there are added tokens. [1]
- Model-specific vocab files (e.g.,
Notes that often confuse people:
tokenizer_config.jsoncommonly containsname_or_pathbecause the base tokenizer stores it from thefrom_pretrained(...)argument (init_kwargs["name_or_path"] = pretrained_model_name_or_path). [2]- If you expected
tokenizer.jsonbut didn’t get it, it usually means you’re not saving/loading a fast tokenizer; confirm withtokenizer.is_fast. [3]
Sources: [1] [2] [3]
Add @pytest.mark.unit marker and make test network-independent by building a tiny local tokenizer.
The test currently calls AutoTokenizer.from_pretrained("bert-base-uncased"), which hits the Hugging Face Hub network on the first run and violates unit test isolation requirements. Build a small in-memory tokenizer using the tokenizers library and save it locally instead. Additionally, add the @pytest.mark.unit marker per test categorization guidelines.
Suggested implementation
- `@pytest.mark.timeout`(30)
+ `@pytest.mark.unit`
+ `@pytest.mark.timeout`(30)
def test_hf_tokenizer_as_local_path_object(self, tmp_path):
# Cover the user case where a user has made a local path object of a WIP tokenizer and wants
# to use that in some megatron model at train time.
# First as a proxy download a tokenizer from HF and save it to a local path. A user would
# do this differently by exporting their WIP tokenizer to a local path.
- # 1. Download a common, small tokenizer from the Hub
- # "bert-base-uncased" is a safe choice as it's small and standard.
- model_id = "bert-base-uncased"
- tokenizer = AutoTokenizer.from_pretrained(model_id)
+ # 1. Build a tiny local tokenizer to avoid network calls.
+ from tokenizers import Tokenizer, models, pre_tokenizers, trainers
+ from transformers import PreTrainedTokenizerFast
+
+ raw_tokenizer = Tokenizer(models.WordLevel(unk_token="[UNK]"))
+ raw_tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+ trainer = trainers.WordLevelTrainer(special_tokens=["[UNK]", "[PAD]"])
+ raw_tokenizer.train_from_iterator(["hello world"], trainer)
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=raw_tokenizer,
+ unk_token="[UNK]",
+ pad_token="[PAD]",
+ )🤖 Prompt for AI Agents
In `@tests/unit_tests/training/test_tokenizer.py` around lines 125 - 168, Add the
`@pytest.mark.unit` marker to test_hf_tokenizer_as_local_path_object and remove
the network call to AutoTokenizer.from_pretrained; instead create a tiny local
tokenizer with the tokenizers library, save it to tmp_path, and use that path in
the existing TokenizerConfig so build_tokenizer loads from disk. Specifically,
replace the AutoTokenizer.from_pretrained("bert-base-uncased") step with
constructing a small Tokenizer (e.g., a WordLevel or BPE model), train/populate
it with a tiny vocabulary or trainer, call tokenizer.save_pretrained or the
tokenizers equivalent to write tokenizer files into local_model_path, then
proceed with the existing TokenizerConfig(tokenizer_model=local_model_path) and
assertions; keep assertions comparing build_tokenizer output to the locally
saved tokenizer and verifying files exist.
Contributor
Author
|
/ok to test dfcba18 |
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Contributor
Author
|
/ok to test 61df1ee |
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Contributor
Author
|
/ok to test 3b57148 |
ko3n1g
approved these changes
Feb 11, 2026
ko3n1g
left a comment
ko3n1g
left a comment
Contributor
There was a problem hiding this comment.
tested internally, good to merge
yaoyu-33
approved these changes
Feb 11, 2026
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
3 participants
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.
What does this PR do ?
Changelog
GitHub Actions CI
See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information
Summary by CodeRabbit
Release Notes
New Features
Refactor
Tests