Skip to content

vllm_mlx.utils.tokenizer

Tokenizer utilities with fallback support for non-standard tokenizers.

View the complete module source at #L1-L280.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.utils.tokenizer

Tokenizer utilities with fallback support for non-standard tokenizers.

Some models (e.g., Nemotron) use non-standard tokenizer configurations that transformers doesn't recognize. This module provides fallback loading directly from tokenizer.json.

vllm_mlx.utils.tokenizer.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.utils.tokenizer.FALLBACK_MODELS module-attribute

FALLBACK_MODELS = ['nemotron', 'NVIDIA-Nemotron']

vllm_mlx.utils.tokenizer._needs_tokenizer_fallback

_needs_tokenizer_fallback(model_name: str) -> bool

Check if model needs tokenizer fallback.

Source code in vllm_mlx/utils/tokenizer.py
def _needs_tokenizer_fallback(model_name: str) -> bool:
    """Check if model needs tokenizer fallback."""
    model_lower = model_name.lower()
    return any(pattern.lower() in model_lower for pattern in FALLBACK_MODELS)

vllm_mlx.utils.tokenizer._needs_strict_false

_needs_strict_false(model_name: str) -> bool

Check if model needs strict=False loading (VLM models with extra weights).

VLM models (e.g., Qwen3.5) have vision_tower weights that don't match the text-only model class. Loading with strict=True fails and wastes memory by loading all weights (~100 GB) before raising ValueError. Detect these models up-front to avoid the double-load penalty.

Source code in vllm_mlx/utils/tokenizer.py
def _needs_strict_false(model_name: str) -> bool:
    """Check if model needs strict=False loading (VLM models with extra weights).

    VLM models (e.g., Qwen3.5) have vision_tower weights that don't match
    the text-only model class.  Loading with strict=True fails and wastes
    memory by loading all weights (~100 GB) before raising ValueError.
    Detect these models up-front to avoid the double-load penalty.
    """
    from mlx_lm.utils import _download, load_config

    try:
        model_path = _download(model_name)
        config = load_config(model_path)
    except Exception:
        return False
    # VLM models have vision_config or text_config with a separate model_type
    if "vision_config" in config and "text_config" in config:
        return True
    return False

vllm_mlx.utils.tokenizer.load_model_with_fallback

load_model_with_fallback(model_name: str, tokenizer_config: dict = None)

Load model and tokenizer with fallback for non-standard tokenizers.

Parameters:

  • model_name (str) –

    HuggingFace model name or local path

  • tokenizer_config (dict, default: None ) –

    Optional tokenizer configuration

Returns:

  • Tuple of (model, tokenizer)

Source code in vllm_mlx/utils/tokenizer.py
def load_model_with_fallback(model_name: str, tokenizer_config: dict = None):
    """
    Load model and tokenizer with fallback for non-standard tokenizers.

    Args:
        model_name: HuggingFace model name or local path
        tokenizer_config: Optional tokenizer configuration

    Returns:
        Tuple of (model, tokenizer)
    """
    from mlx_lm import load

    tokenizer_config = tokenizer_config or {}

    # Check if model needs fallback (e.g., Nemotron)
    if _needs_tokenizer_fallback(model_name):
        logger.info(
            f"Model {model_name} requires tokenizer fallback, loading directly..."
        )
        return _load_with_tokenizer_fallback(model_name)

    # VLM models (e.g., Qwen3.5) have extra vision weights that cause
    # strict=True to fail.  Skip the first load attempt to avoid loading
    # ~100 GB of weights twice (which can cause OOM on 256 GB systems).
    if _needs_strict_false(model_name):
        logger.info(
            f"Model {model_name} detected as VLM, loading directly with strict=False"
        )
        return _load_strict_false(model_name, tokenizer_config)

    try:
        model, tokenizer = load(model_name, tokenizer_config=tokenizer_config)
    except ValueError as e:
        # Fallback for models with non-standard tokenizers
        if "TokenizersBackend" in str(e) or "Tokenizer class" in str(e):
            logger.warning(f"Standard tokenizer loading failed, using fallback: {e}")
            return _load_with_tokenizer_fallback(model_name)
        # Fallback for models with extra weights (e.g., vision tower, MTP layers).
        # Retry with strict=False to discard extra weights.
        elif "parameters not in model" in str(e):
            logger.warning(
                f"Extra parameters found (e.g., vision tower / MTP weights), "
                f"retrying with strict=False: {e}"
            )
            # Clear traceback references to free memory from the failed first load.
            # Without this, large models (200GB+) cause OOM during retry because
            # the traceback holds references to the first load's weight tensors.
            e.__traceback__ = None
            del e
            import gc

            gc.collect()
            return _load_strict_false(model_name, tokenizer_config)
        else:
            raise

    # After successful load, check if MTP weights exist but were stripped by sanitize()
    _try_inject_mtp_post_load(model, model_name)
    return model, tokenizer

vllm_mlx.utils.tokenizer._load_strict_false

_load_strict_false(model_name: str, tokenizer_config: dict = None)

Load model with strict=False to discard extra weights.

Handles models with extra parameters that the text-only model class doesn't define (e.g., vision tower weights in VLM models like Qwen3.5, or MTP layers). The model's own sanitize() handles key remapping (e.g., language_model.* prefix), and strict=False silently drops unmatched keys.

Source code in vllm_mlx/utils/tokenizer.py
def _load_strict_false(model_name: str, tokenizer_config: dict = None):
    """Load model with strict=False to discard extra weights.

    Handles models with extra parameters that the text-only model class
    doesn't define (e.g., vision tower weights in VLM models like Qwen3.5,
    or MTP layers).  The model's own sanitize() handles key remapping
    (e.g., language_model.* prefix), and strict=False silently drops
    unmatched keys.
    """
    import mlx.core as mx
    from mlx_lm.utils import _download, load_model, load_tokenizer

    model_path = _download(model_name)
    model, config = load_model(model_path, strict=False)

    # Verify weights loaded correctly
    from mlx.utils import tree_flatten

    params = tree_flatten(model.parameters())
    total_params = len(params)
    zero_params = sum(1 for _, v in params if mx.all(v == 0).item())
    logger.info(
        f"[strict=False] Loaded {total_params} parameters, "
        f"{zero_params} all-zero tensors"
    )
    # Spot-check embedding weights
    if hasattr(model, "language_model"):
        emb = model.language_model.model.embed_tokens.weight
        logger.info(
            f"[strict=False] embed_tokens: shape={emb.shape}, "
            f"dtype={emb.dtype}, mean={mx.mean(emb.astype(mx.float32)).item():.4f}"
        )

    tokenizer = load_tokenizer(
        model_path,
        tokenizer_config or {},
        eos_token_ids=config.get("eos_token_id", None),
    )
    _try_inject_mtp(model, model_path, config)
    return model, tokenizer

vllm_mlx.utils.tokenizer._try_inject_mtp

_try_inject_mtp(model, model_path, config)

Inject MTP support if model has MTP config + weights.

Source code in vllm_mlx/utils/tokenizer.py
def _try_inject_mtp(model, model_path, config):
    """Inject MTP support if model has MTP config + weights."""
    # Qwen3-Next: flat num_nextn_predict_layers
    if config.get("num_nextn_predict_layers", 0) > 0:
        # Detect Qwen3.5 vs Qwen3-Next by checking text_config or model_type
        text_config = config.get("text_config", config)
        model_type = text_config.get("model_type", config.get("model_type", ""))
        if "qwen3_5" in model_type:
            from ..patches.qwen3_5_mtp import inject_mtp_support
        else:
            from ..patches.qwen3_next_mtp import inject_mtp_support
        inject_mtp_support(model, model_path, config)
        return

    # Qwen3.5: mtp_num_hidden_layers in text_config
    text_config = config.get("text_config", config)
    num_mtp = text_config.get("mtp_num_hidden_layers", 0)
    if num_mtp > 0:
        from ..patches.qwen3_5_mtp import inject_mtp_support

        inject_mtp_support(model, model_path, config)

vllm_mlx.utils.tokenizer._try_inject_mtp_post_load

_try_inject_mtp_post_load(model, model_name)

Check if MTP weights exist but were stripped by sanitize(), and inject.

Source code in vllm_mlx/utils/tokenizer.py
def _try_inject_mtp_post_load(model, model_name):
    """Check if MTP weights exist but were stripped by sanitize(), and inject."""
    import json

    from mlx_lm.utils import _download

    model_path = _download(model_name)
    config_path = Path(model_path) / "config.json"
    if not config_path.exists():
        return
    with open(config_path) as f:
        config = json.load(f)
    # Check for MTP in flat config and nested text_config
    text_config = config.get("text_config", {})
    num_mtp = config.get("num_nextn_predict_layers", 0)
    if num_mtp == 0:
        num_mtp = text_config.get("num_nextn_predict_layers", 0)
    if num_mtp == 0:
        num_mtp = text_config.get("mtp_num_hidden_layers", 0)
    # Also check mtp attribute on language_model for VLM wrappers
    check_model = model
    if hasattr(model, "language_model"):
        check_model = model.language_model
    if num_mtp > 0 and getattr(check_model, "mtp", None) is None:
        mtp_file = Path(model_path) / "mtp" / "weights.safetensors"
        if not mtp_file.exists():
            mtp_file = Path(model_path) / "model-mtp.safetensors"
        if mtp_file.exists():
            logger.info(
                f"[MTP] Found MTP config (layers={num_mtp}) and weights, injecting..."
            )
            _try_inject_mtp(model, model_path, config)
        else:
            logger.info(
                f"[MTP] Config has num_nextn_predict_layers={num_mtp} "
                "but MTP weights not found, skipping MTP."
            )

vllm_mlx.utils.tokenizer._load_with_tokenizer_fallback

_load_with_tokenizer_fallback(model_name: str)

Load model with fallback tokenizer for non-standard models like Nemotron.

Source code in vllm_mlx/utils/tokenizer.py
def _load_with_tokenizer_fallback(model_name: str):
    """Load model with fallback tokenizer for non-standard models like Nemotron."""
    from mlx_lm.utils import load_model

    from .download import ensure_model_downloaded

    logger.info("Loading with tokenizer fallback...")

    # Get model path (with retry/timeout support)
    model_path = ensure_model_downloaded(model_name, is_mllm=False)

    # Load model
    model, _ = load_model(model_path)

    # Try to load tokenizer from tokenizer.json directly
    tokenizer_json = model_path / "tokenizer.json"
    if tokenizer_json.exists():
        from tokenizers import Tokenizer
        from transformers import PreTrainedTokenizerFast

        logger.info("Loading tokenizer from tokenizer.json")
        base_tokenizer = Tokenizer.from_file(str(tokenizer_json))

        # Read tokenizer_config.json for special tokens and chat template
        tokenizer_config_path = model_path / "tokenizer_config.json"
        bos_token = "<s>"
        eos_token = "</s>"
        unk_token = "<unk>"
        chat_template = None

        if tokenizer_config_path.exists():
            with open(tokenizer_config_path) as f:
                config = json.load(f)
                bos_token = config.get("bos_token", bos_token)
                eos_token = config.get("eos_token", eos_token)
                unk_token = config.get("unk_token", unk_token)
                chat_template = config.get("chat_template")

        tokenizer = PreTrainedTokenizerFast(
            tokenizer_object=base_tokenizer,
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            pad_token="<pad>",
        )

        # Set chat template if available
        if chat_template:
            tokenizer.chat_template = chat_template
            logger.info("Chat template loaded from tokenizer_config.json")
        elif _needs_tokenizer_fallback(model_name):
            # Use official Nemotron chat template with thinking support
            tokenizer.chat_template = NEMOTRON_CHAT_TEMPLATE
            logger.info("Using official Nemotron chat template with thinking support")
        else:
            # Default simple ChatML format for other models
            tokenizer.chat_template = DEFAULT_CHATML_TEMPLATE
            logger.info("Using default ChatML chat template")

        logger.info("Tokenizer loaded via fallback successfully")
        return model, tokenizer
    else:
        raise ValueError(f"No tokenizer.json found in {model_path}")

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.utils.tokenizer._needs_tokenizer_fallback · function
vllm_mlx.utils.tokenizer._needs_tokenizer_fallback(model_name: str) -> bool

Check if model needs tokenizer fallback.

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.

Returns

  • Type: bool
  • Direct return expressions: any((pattern.lower() in model_lower for pattern in FALLBACK_MODELS))

Exceptions and behavior

Function _needs_tokenizer_fallback calls model_name.lower, any, pattern.lower; returns any((pattern.lower() in model_lower for pattern in FALLBACK_MODELS)). No direct raise statement appears in this definition.

View source #L25-L28.

vllm_mlx.utils.tokenizer._needs_strict_false · function
vllm_mlx.utils.tokenizer._needs_strict_false(model_name: str) -> bool

Check if model needs strict=False loading (VLM models with extra weights).

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.

Returns

  • Type: bool
  • Direct return expressions: False; True

Exceptions and behavior

Function _needs_strict_false calls _download, load_config; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L31-L49.

vllm_mlx.utils.tokenizer.load_model_with_fallback · function
vllm_mlx.utils.tokenizer.load_model_with_fallback(model_name: str, tokenizer_config: dict = None) -> not annotated

Load model and tokenizer with fallback for non-standard tokenizers.

Parameters

Name Type Required Default Description
model_name str yes none HuggingFace model name or local path
tokenizer_config dict no None Optional tokenizer configuration

Returns

  • Type: not annotated
  • Direct return expressions: _load_with_tokenizer_fallback(model_name); _load_strict_false(model_name, tokenizer_config); (model, tokenizer)

Exceptions and behavior

Function load_model_with_fallback calls _needs_tokenizer_fallback, logger.info, _load_with_tokenizer_fallback, _needs_strict_false; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L52-L111.

vllm_mlx.utils.tokenizer._load_strict_false · function
vllm_mlx.utils.tokenizer._load_strict_false(model_name: str, tokenizer_config: dict = None) -> not annotated

Load model with strict=False to discard extra weights.

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.
tokenizer_config dict no None Optional positional or keyword input; defaults to None.

Returns

  • Type: not annotated
  • Direct return expressions: (model, tokenizer)

Exceptions and behavior

Function _load_strict_false calls _download, load_model, tree_flatten, model.parameters; returns (model, tokenizer). No direct raise statement appears in this definition.

View source #L114-L153.

vllm_mlx.utils.tokenizer._try_inject_mtp · function
vllm_mlx.utils.tokenizer._try_inject_mtp(model, model_path, config) -> not annotated

Inject MTP support if model has MTP config + weights.

Parameters

Name Type Required Default Description
model not annotated yes none Required positional or keyword input.
model_path not annotated yes none Required positional or keyword input.
config not annotated yes none Required positional or keyword input.

Returns

  • Type: not annotated
  • Direct return expressions: None

Exceptions and behavior

Function _try_inject_mtp calls config.get, text_config.get, inject_mtp_support; returns None. No direct raise statement appears in this definition.

View source #L156-L176.

vllm_mlx.utils.tokenizer._try_inject_mtp_post_load · function
vllm_mlx.utils.tokenizer._try_inject_mtp_post_load(model, model_name) -> not annotated

Check if MTP weights exist but were stripped by sanitize(), and inject.

Parameters

Name Type Required Default Description
model not annotated yes none Required positional or keyword input.
model_name not annotated yes none Required positional or keyword input.

Returns

  • Type: not annotated
  • Direct return expressions: None

Exceptions and behavior

Function _try_inject_mtp_post_load calls _download, Path, config_path.exists, open; returns None. No direct raise statement appears in this definition.

View source #L179-L215.

vllm_mlx.utils.tokenizer._load_with_tokenizer_fallback · function
vllm_mlx.utils.tokenizer._load_with_tokenizer_fallback(model_name: str) -> not annotated

Load model with fallback tokenizer for non-standard models like Nemotron.

Parameters

Name Type Required Default Description
model_name str yes none Required positional or keyword input.

Returns

  • Type: not annotated
  • Direct return expressions: (model, tokenizer)

Exceptions and behavior

Function _load_with_tokenizer_fallback calls logger.info, ensure_model_downloaded, load_model, tokenizer_json.exists; can raise ValueError; returns (model, tokenizer). Directly raised exceptions: ValueError.

View source #L218-L280.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol Kind Signature and inputs What it does Source
_needs_tokenizer_fallback function _needs_tokenizer_fallback(model_name: str) -> bool Check if model needs tokenizer fallback. #L25-L28
_needs_strict_false function _needs_strict_false(model_name: str) -> bool Check if model needs strict=False loading (VLM models with extra weights). #L31-L49
load_model_with_fallback function load_model_with_fallback(model_name: str, tokenizer_config: dict = None) -> not annotated Load model and tokenizer with fallback for non-standard tokenizers. #L52-L111
_load_strict_false function _load_strict_false(model_name: str, tokenizer_config: dict = None) -> not annotated Load model with strict=False to discard extra weights. #L114-L153
_try_inject_mtp function _try_inject_mtp(model, model_path, config) -> not annotated Inject MTP support if model has MTP config + weights. #L156-L176
_try_inject_mtp_post_load function _try_inject_mtp_post_load(model, model_name) -> not annotated Check if MTP weights exist but were stripped by sanitize(), and inject. #L179-L215
_load_with_tokenizer_fallback function _load_with_tokenizer_fallback(model_name: str) -> not annotated Load model with fallback tokenizer for non-standard models like Nemotron. #L218-L280