Path-based huggingface tokenizers do not load from the checkpoint when using `load_tokenizer(...)`

[jstjohn]

Describe the bug

If you use a path based hugging face checkpoint, the object saved into the mbridge checkpoint will not load from the config saved into the checkpoint. The problem is that the config is saved with the tokenizer_model set to PosixPath(".") which would only work if you ran the load function from within the iter_XXXXX/tokenizer directory.

A loaded config from the checkpoint looks like:

(Pdb) cfg
TokenizerConfig(vocab_size=None, vocab_file=None, merge_file=None, vocab_extra_ids=0, tokenizer_type='HuggingFaceTokenizer', tokenizer_model=PosixPath('.'), tiktoken_pattern=None, tiktoken_num_special_tokens=1000, tiktoken_special_tokens=None, tokenizer_prompt_format=None, special_tokens=None, image_tag_type=None, hf_tokenizer_kwargs={'trust_remote_code': True})

Where originally the tokenizer_model=PosixPath("/path/to/original/tokenizer/dir"). Leaving it as is would not be portable. However maybe we could add special handling of paths to tokenizer = load_tokenizer(mbridge_ckpt_path) so that it sets tokenizer_model=Path(mbridge_ckpt_path)/"tokenizer") at runtime if some condition is met? Maybe when isinstance(cfg.tokenizer_model, Path) then override that probably invalid path with ckpt_dir/"tokenizer"?

To be clear the following works still with the tokenizer saved into the checkpoint, so the issue is that the path saved in the config does not point to the right place (it couldn't though with a portable config):

tokenizer = _HuggingFaceTokenizer(mbridge_ckpt_path / "tokenizer")

Another option would be to use some magic special symbol like __CKPT_BASEDIR__/tokenizer and then resolve that checkpoint basedir at runtime from the config to a path? I think the default Path(".") that is in there now is probably enough of a clue that we can change the value safely to something valid at runtime, unless we have other places we want to resolve paths that are relative to the base of a checkpoint.

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

Add any other context about the problem here.

[JRD971000]

Hi @dimapihtar , could you please TAL this bug?

[dimapihtar]

Hi @jstjohn, thank you for reporting this bug.

Could you please provide a simple script which reproduces this issue? Thanks.

[jstjohn]

I just had cursor make this repro script @dimapihtar.

"""
Script to reproduce issue #1900: Path-based HuggingFace tokenizers do not load
from checkpoint when using load_tokenizer(...).

This script demonstrates that:
1. When you save a tokenizer config using a local path (not a HF model ID), the path
   is not properly serialized by the mbridge config system.
2. The `load_tokenizer()` function will fail to load the tokenizer from the
   checkpoint when run from a different directory.

Repro steps:
1. AutoTokenizer.from_pretrained to download a tokenizer from HF
2. tokenizer.save_pretrained() to save it to a local path on disk
3. Configure a TokenizerConfig to use the tokenizer from step 2 via PATH (not HF model ID)
4. Build the tokenizer and use ACTUAL mbridge save mechanisms:
   - save_tokenizer_assets() to save tokenizer files
   - ConfigContainer.to_yaml() to save run_config.yaml
5. Try to load the tokenizer from the checkpoint using mbridge tools from TWO different directories
   - It will fail because the saved path is not portable

Issue: https://github.com/NVIDIA-NeMo/Megatron-Bridge/issues/1900
"""

import os
import tempfile
from dataclasses import dataclass
from pathlib import Path

import yaml
from transformers import AutoTokenizer

from megatron.bridge.training.checkpointing import save_tokenizer_assets
from megatron.bridge.training.tokenizers.config import TokenizerConfig
from megatron.bridge.training.tokenizers.tokenizer import _HuggingFaceTokenizer, build_tokenizer
from megatron.bridge.training.utils.config_utils import _ConfigContainerBase
from megatron.bridge.utils.instantiate_utils import instantiate


@dataclass
class MockRunConfig(_ConfigContainerBase):
    """Mock run config container that mimics what mbridge saves to run_config.yaml."""

    tokenizer: TokenizerConfig | None = None


def try_load_tokenizer_from_checkpoint(checkpoint_path: Path, run_config_path: Path, description: str) -> bool:
    """Try to load tokenizer from checkpoint and report success/failure.

    Args:
        checkpoint_path: Path to the checkpoint directory
        run_config_path: Path to the run_config.yaml file
        description: Description of the test scenario

    Returns:
        True if tokenizer loaded successfully, False otherwise
    """
    print(f"\n--- Testing from: {description} ---")
    print(f"Current working directory: {os.getcwd()}")

    # Read the config from checkpoint (what load_tokenizer does internally)
    with open(run_config_path) as f:
        loaded_config = yaml.safe_load(f)

    tokenizer_cfg_dict = loaded_config["tokenizer"]
    print(f"Loaded tokenizer_model from config: {tokenizer_cfg_dict.get('tokenizer_model')}")
    print(f"Type: {type(tokenizer_cfg_dict.get('tokenizer_model'))}")

    # Instantiate the TokenizerConfig
    cfg = instantiate(tokenizer_cfg_dict)
    print(f"After instantiate: cfg.tokenizer_model = '{cfg.tokenizer_model}'")

    # Try to build the tokenizer
    try:
        tokenizer = build_tokenizer(cfg)
        print(f"SUCCESS: Loaded tokenizer with vocab_size={tokenizer.vocab_size}")
        return True
    except Exception as e:
        print(f"FAILURE: {e}")
        return False


def reproduce_issue():
    """Reproduce the tokenizer loading issue using actual mbridge save mechanisms."""
    # Create temporary directories for our test
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)

        # ======================================================================
        # Step 1: Download a tokenizer from HuggingFace and save it locally
        # ======================================================================
        print("=" * 70)
        print("STEP 1: Download tokenizer from HuggingFace and save to local path")
        print("=" * 70)

        local_tokenizer_path = tmpdir / "local_tokenizer"
        local_tokenizer_path.mkdir(parents=True, exist_ok=True)

        # Download a small tokenizer (GPT-2 is small and commonly available)
        hf_tokenizer = AutoTokenizer.from_pretrained("gpt2")
        hf_tokenizer.save_pretrained(str(local_tokenizer_path))
        print(f"Downloaded and saved tokenizer to: {local_tokenizer_path}")
        print(f"Tokenizer files: {[f.name for f in local_tokenizer_path.iterdir()]}")

        # ======================================================================
        # Step 2: Create a TokenizerConfig using the LOCAL PATH (not HF model ID)
        # ======================================================================
        print("\n" + "=" * 70)
        print("STEP 2: Create TokenizerConfig using LOCAL PATH (not HF model ID)")
        print("=" * 70)

        # This is the problematic case: using a Path object instead of HF model ID
        tokenizer_config = TokenizerConfig(
            tokenizer_type="HuggingFaceTokenizer",
            tokenizer_model=local_tokenizer_path,  # <-- Using local Path, NOT "gpt2"
            hf_tokenizer_kwargs={"trust_remote_code": False},
        )
        print(f"TokenizerConfig.tokenizer_model = {tokenizer_config.tokenizer_model}")
        print(f"Type: {type(tokenizer_config.tokenizer_model)}")

        # ======================================================================
        # Step 3: Build the tokenizer (needed for save_tokenizer_assets)
        # ======================================================================
        print("\n" + "=" * 70)
        print("STEP 3: Build tokenizer from config")
        print("=" * 70)

        tokenizer_instance = build_tokenizer(tokenizer_config)
        print(f"Built tokenizer: {tokenizer_instance}")
        print(f"Vocab size: {tokenizer_instance.vocab_size}")

        # ======================================================================
        # Step 4: Use ACTUAL mbridge save mechanisms
        # ======================================================================
        print("\n" + "=" * 70)
        print("STEP 4: Save checkpoint using ACTUAL mbridge mechanisms")
        print("=" * 70)

        checkpoint_path = tmpdir / "checkpoints" / "iter_0000001"
        checkpoint_path.mkdir(parents=True, exist_ok=True)

        # 4a. Use save_tokenizer_assets (the actual mbridge function)
        print("\n4a. Calling save_tokenizer_assets()...")
        save_tokenizer_assets(tokenizer_instance, tokenizer_config, str(checkpoint_path))

        tokenizer_in_ckpt = checkpoint_path / "tokenizer"
        if tokenizer_in_ckpt.exists():
            print(f"Tokenizer saved to: {tokenizer_in_ckpt}")
            print(f"Files: {[f.name for f in tokenizer_in_ckpt.iterdir()]}")
        else:
            print("WARNING: tokenizer directory was not created!")

        # 4b. Use ConfigContainer.to_yaml (the actual mbridge config save)
        print("\n4b. Calling ConfigContainer.to_yaml()...")
        run_config = MockRunConfig(tokenizer=tokenizer_config)
        run_config_path = checkpoint_path / "run_config.yaml"
        run_config.to_yaml(str(run_config_path))

        print(f"Saved run_config.yaml to: {run_config_path}")
        print("\nActual file content:")
        with open(run_config_path) as f:
            file_content = f.read()
            print(file_content)

        # ======================================================================
        # Step 5: Try to load tokenizer from TWO different directories
        # ======================================================================
        print("=" * 70)
        print("STEP 5: Test loading from TWO different directories")
        print("=" * 70)

        original_cwd = os.getcwd()
        results = {}

        # Test 1: Load from /tmp
        try:
            os.chdir("/tmp")
            results["from /tmp"] = try_load_tokenizer_from_checkpoint(
                checkpoint_path, run_config_path, "directory /tmp"
            )
        finally:
            os.chdir(original_cwd)

        # Test 2: Load from /home (another different directory)
        try:
            os.chdir("/home")
            results["from /home"] = try_load_tokenizer_from_checkpoint(
                checkpoint_path, run_config_path, "directory /home"
            )
        finally:
            os.chdir(original_cwd)

        # Test 3: Load from the checkpoint directory itself (this might work due to ".")
        try:
            os.chdir(checkpoint_path)
            results["from checkpoint dir"] = try_load_tokenizer_from_checkpoint(
                checkpoint_path, run_config_path, f"checkpoint directory {checkpoint_path}"
            )
        finally:
            os.chdir(original_cwd)

        # ======================================================================
        # Summary and Analysis
        # ======================================================================
        print("\n" + "=" * 70)
        print("SUMMARY")
        print("=" * 70)

        print("\nTest Results:")
        for location, success in results.items():
            status = "PASS" if success else "FAIL"
            print(f"  {location}: {status}")

        all_passed = all(results.values())
        any_passed = any(results.values())

        if not all_passed:
            print("\n>>> BUG CONFIRMED! <<<")
            if any_passed:
                print("The tokenizer loads from SOME directories but not others.")
                print("This proves the checkpoint is NOT PORTABLE.")
            else:
                print("The tokenizer failed to load from ALL test directories.")

            print("\nRoot cause: PosixPath objects are serialized incorrectly by")
            print("_safe_object_representer in yaml_utils.py as:")
            print("    _target_: pathlib.PosixPath")
            print("    _call_: true")
            print("This loses the actual path value! When instantiated, it becomes PosixPath('.').")

        # ======================================================================
        # Workaround demonstration
        # ======================================================================
        print("\n" + "=" * 70)
        print("WORKAROUND: Load directly from checkpoint/tokenizer directory")
        print("=" * 70)

        os.chdir("/tmp")  # Change to a different directory
        try:
            direct_tokenizer = _HuggingFaceTokenizer(tokenizer_in_ckpt)
            print(f"Direct loading from {tokenizer_in_ckpt} works!")
            print(f"Vocab size: {direct_tokenizer.vocab_size}")
        finally:
            os.chdir(original_cwd)

        # ======================================================================
        # Proposed fixes
        # ======================================================================
        print("\n" + "=" * 70)
        print("PROPOSED FIXES")
        print("=" * 70)

        print("""
1. Add a custom YAML representer for Path objects in yaml_utils.py:

    def _path_representer(dumper, data):
        return dumper.represent_str(str(data))

    yaml.SafeDumper.add_multi_representer(Path, _path_representer)

2. In load_tokenizer(), detect invalid paths and resolve to checkpoint/tokenizer:

    if isinstance(cfg.tokenizer_model, Path) or cfg.tokenizer_model == '.':
        cfg.tokenizer_model = Path(checkpoint_path) / 'tokenizer'

3. Convert Path to str in TokenizerConfig before serialization (in to_dict or
   _convert_value_to_dict).
""")

        print("=" * 70)
        print("Done!")
        print("=" * 70)


if __name__ == "__main__":
    reproduce_issue()

[jstjohn]

Note how running try_load_tokenizer works from the tokenizer directory (highlighting the issue where it's saved as path("."))

But also note that we want it to work from anywhere, and ideally without resorting to an abspath, like we should be able to copy the checkpoint folder to a new location and it should still work when you load from anywhere

[dimapihtar]

PR-fix was merged to main: #2047