diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index e8d7f81450c..eea649c2355 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1206,11 +1206,10 @@ def setup_model_and_optimizer( # Context used for persisting some state between checkpoint saves. checkpointing_context = init_checkpointing_context(megatron_cfg.checkpoint) - # Tokenizer - if megatron_cfg.tokenizer.hf_tokenizer_kwargs is None: - megatron_cfg.tokenizer.hf_tokenizer_kwargs = {} - megatron_cfg.tokenizer.hf_tokenizer_kwargs["trust_remote_code"] = True - megatron_cfg.tokenizer.hf_tokenizer_kwargs["use_fast"] = True + # Set the attribute directly instead of updating hf_tokenizer_kwargs, because + # Megatron-Bridge's TokenizerConfig snapshots hf_tokenizer_kwargs into plain + # attributes at __post_init__ and never re-reads the dict afterwards. + megatron_cfg.tokenizer.trust_remote_code = True build_tokenizer( megatron_cfg.tokenizer, make_vocab_size_divisible_by=megatron_cfg.model.make_vocab_size_divisible_by diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 55f511bb569..8f2448e4e76 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -18,6 +18,55 @@ from nemo_rl.utils.checkpoint import PretrainedCheckpointConfig +def _patch_transformers_tokenizer_class_set(): + """Undo the transformers block on deepseek_v3 tokenizers. + + Root cause: transformers 5.4-5.11 lists "deepseek_v3" in two internal + registries -- MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS (a set) and + TOKENIZER_MAPPING_NAMES (a dict pinning it to "TokenizersBackend"). Together + they force the fast tokenizer backend and suppress trust_remote_code, so + AutoTokenizer can only load via a local tokenizer.json. Models like + Moonlight-16B-A3B ship no tokenizer.json (only tiktoken.model + a remote-code + TikTokenTokenizer), so offline loading fails. + + Removing both entries restores the trust_remote_code / auto_map path. + discard/pop-with-default are no-ops when the entries are absent, so this is + safe on any transformers version in the currently-supported range. + + Placed here (nemo_rl/models/policy/__init__.py) so it fires exactly once + per process the first time any policy code is imported -- covers the driver + (via nemo_rl.algorithms.grpo) and every policy worker (Megatron / DTensor / + DTensor v2 all import from nemo_rl.models.policy) without polluting nemo_rl + consumers that don't touch tokenizers. + """ + import transformers + from packaging.version import Version as PkgVersion + + # This whole patch exists only because Megatron-Bridge caps the transformers + # upper bound below 5.9 today, which forces us onto a transformers version + # that still has the deepseek_v3 tokenizer-blocklist bug. Once MBridge relaxes + # its transformers upper bound to >=5.12, we can drop this workaround. + # TODO: remove this patch (and the assert below) once MBridge relaxes its + # transformers upper bound past the deepseek_v3 fix (~transformers 5.12). + # https://github.com/NVIDIA-NeMo/RL/issues/2764 + assert PkgVersion(transformers.__version__) < PkgVersion("5.12.0"), ( + f"transformers {transformers.__version__} detected. " + "The deepseek_v3 tokenizer-blocklist patch was written for <5.12. " + "Check if the upstream fix now applies and remove this patch if so." + ) + + from transformers.models.auto.tokenization_auto import ( + MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS, + TOKENIZER_MAPPING_NAMES, + ) + + MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS.discard("deepseek_v3") + TOKENIZER_MAPPING_NAMES.pop("deepseek_v3", None) + + +_patch_transformers_tokenizer_class_set() + + class LoRAConfigDisabled(TypedDict): enabled: Literal[False]