refactor(convert): move shared checkpoint utils to convert.utils

Move download_checkpoint_from_hub, ensure_checkpoint_is_local,
load_checkpoint_config, and load_checkpoint_weights from
convert/eagle/utils.py to convert/utils.py so they can be shared
across converter implementations without cross-algorithm imports.

Update all consumers (eagle, eagle3, mtp converters and their tests)
to import from the canonical location.

Signed-off-by: Rahul-Tuli <rtuli@redhat.com>

Author: rahul-tuli

src/speculators/convert/eagle/eagle3_converter.py

@@ -13,8 +13,10 @@
 from speculators.convert.eagle.utils import (
     build_llama_config_dtype_kwarg,
     build_llama_config_rope_kwargs,
-    ensure_checkpoint_is_local,
     find_vocab_size,
+)
+from speculators.convert.utils import (
+    ensure_checkpoint_is_local,
     load_checkpoint_config,
     load_checkpoint_weights,
 )

src/speculators/convert/eagle/eagle_converter.py

@@ -16,6 +16,8 @@
 from speculators.convert.eagle.utils import (
     build_llama_config_rope_kwargs,
     detect_fusion_bias_and_layernorms,
+)
+from speculators.convert.utils import (
     ensure_checkpoint_is_local,
     load_checkpoint_config,
     load_checkpoint_weights,

src/speculators/convert/eagle/utils.py

@@ -1,18 +1,14 @@
 """
-Utility functions for checkpoint conversion operations.
+Utility functions for Eagle checkpoint conversion operations.
 """
 
 from __future__ import annotations
 
 import inspect
-import json
-from pathlib import Path
 from typing import Any
 
 import torch
-from huggingface_hub import snapshot_download
 from loguru import logger
-from safetensors import safe_open
 from transformers import LlamaConfig
 
 _LLAMA_CONFIG_PARAMS = set(inspect.signature(LlamaConfig.__init__).parameters)
@@ -69,139 +65,6 @@ def find_vocab_size(config_dict: dict) -> int | None:
     return None
 
 
-def download_checkpoint_from_hub(model_id: str, cache_dir: str | None = None) -> Path:
-    """
-    Download a checkpoint from HuggingFace Hub.
-
-    :param model_id: HuggingFace model ID
-    :param cache_dir: Optional directory to cache downloads
-    :return: Local path to the downloaded checkpoint
-    :raises FileNotFoundError: If the checkpoint cannot be downloaded
-
-    :Example:
-
-        >>> path = download_checkpoint_from_hub("yuhuili/EAGLE-LLaMA3.1-Instruct-8B")
-        >>> print(path)
-        /home/user/.cache/huggingface/hub/models--yuhuili--EAGLE-LLaMA3.1-Instruct-8B/snapshots/...
-    """
-    logger.info(f"Downloading checkpoint from HuggingFace: {model_id}")
-    try:
-        local_path = snapshot_download(
-            repo_id=model_id,
-            allow_patterns=["*.json", "*.safetensors", "*.bin", "*.index.json"],
-            cache_dir=cache_dir,
-        )
-        logger.debug(f"Downloaded to: {local_path}")
-        return Path(local_path)
-    except Exception as hf_exception:
-        logger.error(f"Failed to download checkpoint: {hf_exception}")
-        raise FileNotFoundError(f"Checkpoint not found: {model_id}") from hf_exception
-
-
-def ensure_checkpoint_is_local(
-    checkpoint_path: str | Path, cache_dir: str | Path | None = None
-) -> Path:
-    """
-    Ensure we have a local copy of the checkpoint.
-
-    If the path exists locally, return it. Otherwise, treat it as a
-    HuggingFace model ID and download it.
-
-    :param checkpoint_path: Local path or HuggingFace model ID
-    :param cache_dir: Optional cache directory for downloads
-    :return: Path to local checkpoint directory
-
-    :Example:
-
-        >>> # Local path - returned as-is
-        >>> local = ensure_checkpoint_is_local("./my_checkpoint")
-
-        >>> # HuggingFace ID - downloaded first
-        >>> downloaded = ensure_checkpoint_is_local(
-        ...     "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
-        ... )
-    """
-    checkpoint_path = Path(checkpoint_path)
-
-    if checkpoint_path.exists():
-        logger.debug(f"Using local checkpoint: {checkpoint_path}")
-        return checkpoint_path
-
-    return download_checkpoint_from_hub(
-        model_id=str(checkpoint_path), cache_dir=str(cache_dir) if cache_dir else None
-    )
-
-
-def load_checkpoint_config(checkpoint_dir: Path) -> dict:
-    """
-    Load the config.json from a checkpoint directory.
-
-    :param checkpoint_dir: Path to checkpoint directory
-    :return: Config dictionary
-    :raises FileNotFoundError: If config.json is not found
-
-    :Example:
-
-        >>> config = load_checkpoint_config(Path("./checkpoint"))
-        >>> print(config["model_type"])
-        llama
-    """
-    config_path = checkpoint_dir / "config.json"
-    if not config_path.exists():
-        raise FileNotFoundError(f"No config.json found at {checkpoint_dir}")
-
-    logger.debug(f"Loading config from: {config_path}")
-    with config_path.open() as f:
-        return json.load(f)
-
-
-def load_checkpoint_weights(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
-    """
-    Load model weights from a checkpoint directory.
-
-    Supports both safetensors and PyTorch bin formats.
-
-    :param checkpoint_dir: Path to checkpoint directory
-    :return: Dictionary mapping weight names to tensors
-    :raises FileNotFoundError: If no weights are found
-    :raises NotImplementedError: If checkpoint is sharded
-
-    :Example:
-
-        >>> weights = load_checkpoint_weights(Path("./checkpoint"))
-        >>> print(f"Loaded {len(weights)} weights")
-        Loaded 50 weights
-    """
-    weights = {}
-
-    safetensors_path = checkpoint_dir / "model.safetensors"
-    if safetensors_path.exists():
-        logger.debug(f"Loading safetensors weights from: {safetensors_path}")
-        with safe_open(safetensors_path, framework="pt") as f:
-            # safetensors requires iterating over keys() method
-            for key in f.keys():  # noqa: SIM118
-                weights[key] = f.get_tensor(key)
-        return weights
-
-    pytorch_path = checkpoint_dir / "pytorch_model.bin"
-    if pytorch_path.exists():
-        logger.debug(f"Loading PyTorch weights from: {pytorch_path}")
-        return torch.load(pytorch_path, map_location="cpu")
-
-    index_paths = [
-        checkpoint_dir / "model.safetensors.index.json",
-        checkpoint_dir / "pytorch_model.bin.index.json",
-    ]
-    for index_path in index_paths:
-        if index_path.exists():
-            raise NotImplementedError(
-                f"Sharded checkpoint detected: {index_path}. "
-                "Please use a single-file checkpoint."
-            )
-
-    raise FileNotFoundError(f"No weights found at {checkpoint_dir}")
-
-
 def detect_fusion_bias_and_layernorms(
     weights: dict[str, torch.Tensor],
 ) -> tuple[bool, bool]:

src/speculators/convert/mtp/converter.py

@@ -20,7 +20,7 @@
 from transformers import PretrainedConfig
 
 from speculators.config import SpeculatorsConfig, VerifierConfig
-from speculators.convert.eagle.utils import (
+from speculators.convert.utils import (
     ensure_checkpoint_is_local,
     load_checkpoint_config,
 )

src/speculators/convert/utils.py

@@ -0,0 +1,143 @@
+"""Shared utility functions for checkpoint conversion operations."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from loguru import logger
+from safetensors import safe_open
+
+
+def download_checkpoint_from_hub(model_id: str, cache_dir: str | None = None) -> Path:
+    """
+    Download a checkpoint from HuggingFace Hub.
+
+    :param model_id: HuggingFace model ID
+    :param cache_dir: Optional directory to cache downloads
+    :return: Local path to the downloaded checkpoint
+    :raises FileNotFoundError: If the checkpoint cannot be downloaded
+
+    :Example:
+
+        >>> path = download_checkpoint_from_hub("yuhuili/EAGLE-LLaMA3.1-Instruct-8B")
+        >>> print(path)
+        /home/user/.cache/huggingface/hub/models--yuhuili--EAGLE-LLaMA3.1-Instruct-8B/snapshots/...
+    """
+    logger.info(f"Downloading checkpoint from HuggingFace: {model_id}")
+    try:
+        local_path = snapshot_download(
+            repo_id=model_id,
+            allow_patterns=["*.json", "*.safetensors", "*.bin", "*.index.json"],
+            cache_dir=cache_dir,
+        )
+        logger.debug(f"Downloaded to: {local_path}")
+        return Path(local_path)
+    except Exception as hf_exception:
+        logger.error(f"Failed to download checkpoint: {hf_exception}")
+        raise FileNotFoundError(f"Checkpoint not found: {model_id}") from hf_exception
+
+
+def ensure_checkpoint_is_local(
+    checkpoint_path: str | Path, cache_dir: str | Path | None = None
+) -> Path:
+    """
+    Ensure we have a local copy of the checkpoint.
+
+    If the path exists locally, return it. Otherwise, treat it as a
+    HuggingFace model ID and download it.
+
+    :param checkpoint_path: Local path or HuggingFace model ID
+    :param cache_dir: Optional cache directory for downloads
+    :return: Path to local checkpoint directory
+
+    :Example:
+
+        >>> # Local path - returned as-is
+        >>> local = ensure_checkpoint_is_local("./my_checkpoint")
+
+        >>> # HuggingFace ID - downloaded first
+        >>> downloaded = ensure_checkpoint_is_local(
+        ...     "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
+        ... )
+    """
+    checkpoint_path = Path(checkpoint_path)
+
+    if checkpoint_path.exists():
+        logger.debug(f"Using local checkpoint: {checkpoint_path}")
+        return checkpoint_path
+
+    return download_checkpoint_from_hub(
+        model_id=str(checkpoint_path), cache_dir=str(cache_dir) if cache_dir else None
+    )
+
+
+def load_checkpoint_config(checkpoint_dir: Path) -> dict:
+    """
+    Load the config.json from a checkpoint directory.
+
+    :param checkpoint_dir: Path to checkpoint directory
+    :return: Config dictionary
+    :raises FileNotFoundError: If config.json is not found
+
+    :Example:
+
+        >>> config = load_checkpoint_config(Path("./checkpoint"))
+        >>> print(config["model_type"])
+        llama
+    """
+    config_path = checkpoint_dir / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"No config.json found at {checkpoint_dir}")
+
+    logger.debug(f"Loading config from: {config_path}")
+    with config_path.open() as f:
+        return json.load(f)
+
+
+def load_checkpoint_weights(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
+    """
+    Load model weights from a checkpoint directory.
+
+    Supports both safetensors and PyTorch bin formats.
+
+    :param checkpoint_dir: Path to checkpoint directory
+    :return: Dictionary mapping weight names to tensors
+    :raises FileNotFoundError: If no weights are found
+    :raises NotImplementedError: If checkpoint is sharded
+
+    :Example:
+
+        >>> weights = load_checkpoint_weights(Path("./checkpoint"))
+        >>> print(f"Loaded {len(weights)} weights")
+        Loaded 50 weights
+    """
+    weights = {}
+
+    safetensors_path = checkpoint_dir / "model.safetensors"
+    if safetensors_path.exists():
+        logger.debug(f"Loading safetensors weights from: {safetensors_path}")
+        with safe_open(safetensors_path, framework="pt") as f:
+            for key in f.keys():  # noqa: SIM118
+                weights[key] = f.get_tensor(key)
+        return weights
+
+    pytorch_path = checkpoint_dir / "pytorch_model.bin"
+    if pytorch_path.exists():
+        logger.debug(f"Loading PyTorch weights from: {pytorch_path}")
+        return torch.load(pytorch_path, map_location="cpu")
+
+    index_paths = [
+        checkpoint_dir / "model.safetensors.index.json",
+        checkpoint_dir / "pytorch_model.bin.index.json",
+    ]
+    for index_path in index_paths:
+        if index_path.exists():
+            raise NotImplementedError(
+                f"Sharded checkpoint detected: {index_path}. "
+                "Please use a single-file checkpoint."
+            )
+
+    raise FileNotFoundError(f"No weights found at {checkpoint_dir}")

tests/unit/convert/test_eagle3_converter.py

@@ -14,9 +14,7 @@
 import torch
 
 from speculators.convert.eagle.eagle3_converter import Eagle3Converter
-from speculators.convert.eagle.utils import (
-    load_checkpoint_config,
-)
+from speculators.convert.utils import load_checkpoint_config
 
 
 class TestEagle3ConverterFixes: