vllm.model_executor.model_loader.weight_utils ¶

def _get_available_ram_bytes() -> int:
    """Return available RAM, honoring cgroup limits."""
    import psutil

    host_available = psutil.virtual_memory().available

    from vllm.utils.cpu_resource_utils import (
        get_cgroup_memory_limit,
        get_cgroup_memory_usage,
    )

    cgroup_limit = get_cgroup_memory_limit()
    if cgroup_limit is None:
        return host_available
    cgroup_usage = get_cgroup_memory_usage()
    cgroup_available = (
        cgroup_limit if cgroup_usage is None else max(0, cgroup_limit - cgroup_usage)
    )
    return min(host_available, cgroup_available)

def _get_checkpoints_size_bytes(files: list[str]) -> int:
    """Return the total size of the checkpoint files in bytes."""
    if not files:
        return 0
    return sum(os.path.getsize(f) for f in files)

def _get_fs_type(files: list[str]) -> str:
    """Get the filesystem type of the first file in *files* (Linux only)."""
    if not files:
        return ""
    try:
        # Only the first file is checked — all checkpoint shards reside
        # in the same directory and therefore on the same filesystem.
        resolved = os.path.realpath(files[0])
        best_mount = ""
        best_fstype = ""
        # /proc/mounts may contain nested mount points (e.g. "/" -> ext4,
        # "/data" -> nfs4, "/data/local" -> ext4).  We pick the entry with
        # the longest matching mount_point — the same "longest prefix match"
        # rule the kernel uses to decide which filesystem serves a path.
        with open("/proc/mounts") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 3:
                    continue
                mount_point, fstype = parts[1], parts[2]
                if (
                    resolved == mount_point
                    or resolved.startswith(os.path.join(mount_point, ""))
                ) and len(mount_point) > len(best_mount):
                    best_mount = mount_point
                    best_fstype = fstype
        return best_fstype
    except Exception:
        # /proc/mounts is Linux-specific; on other OSes (or if the read
        # fails for any reason) we fall back to an empty string.
        return ""

def _mapped_weight_name(weights_mapper: "WeightsMapper | None", key: str) -> str | None:
    """Apply ``WeightsMapper.map_name`` when present; else identity."""
    if weights_mapper is None:
        return key
    return weights_mapper.map_name(key)

def _natural_sort_key(filepath: str) -> list:
    """Natural sort key for filenames with numeric components, such as
    model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']"""
    return [
        int(s) if s.isdigit() else s
        for s in re.split(r"(\d+)", os.path.basename(filepath))
    ]

def _prefetch_all_checkpoints(
    sorted_files: list[str],
    num_prefetch_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
    block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
) -> None:
    """Start prefetching checkpoint files into page cache in a background thread."""
    if num_prefetch_threads < 1:
        raise ValueError("safetensors prefetch num threads must be >= 1")
    if block_size < 1:
        raise ValueError("safetensors prefetch block size must be >= 1")

    if torch.distributed.is_initialized():
        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()
    else:
        rank = 0
        world_size = 1
    paths_to_prefetch = sorted_files[rank::world_size]
    total_for_rank = len(paths_to_prefetch)

    async def _prefetch_all() -> None:
        loop = asyncio.get_running_loop()
        completed = 0
        next_log_pct = 10

        async def prefetch_one(
            path: str,
            executor: concurrent.futures.ThreadPoolExecutor,
        ) -> None:
            nonlocal completed, next_log_pct
            try:
                await loop.run_in_executor(
                    executor, _prefetch_checkpoint, path, block_size
                )
                completed += 1
                if total_for_rank > 0 and next_log_pct <= 100:
                    pct = 100 * completed / total_for_rank
                    if pct >= next_log_pct:
                        logger.info(
                            "Prefetching checkpoint files: %d%% (%d/%d)",
                            next_log_pct,
                            completed,
                            total_for_rank,
                        )
                        next_log_pct += 10
            except Exception:
                logger.warning(
                    "Failed to prefetch checkpoint file %r.", path, exc_info=True
                )

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=num_prefetch_threads
        ) as executor:
            await asyncio.gather(
                *(prefetch_one(p, executor) for p in paths_to_prefetch)
            )

    def _run_prefetch() -> None:
        start = time.perf_counter()
        asyncio.run(_prefetch_all())
        elapsed = time.perf_counter() - start
        logger.info(
            "Prefetching checkpoint files into page cache finished in %.2fs",
            elapsed,
        )

    logger.info(
        "Prefetching checkpoint files into page cache started "
        "(in background, num_threads=%d, block_size=%d bytes)",
        num_prefetch_threads,
        block_size,
    )
    threading.Thread(target=_run_prefetch, daemon=True).start()

def _prefetch_checkpoint(
    file_path: str,
    block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
) -> None:
    """Prefetch a checkpoint file into the OS page cache.

    Reads the file in blocks so the kernel caches its pages before workers load
    the same file.
    """
    if block_size < 1:
        raise ValueError("safetensors prefetch block size must be >= 1")

    with open(file_path, "rb") as f:
        while f.read(block_size):
            pass

def _shared_hf_root_module_prefixes(
    lm_module_prefixes: tuple[str, ...],
    weights_mapper: "WeightsMapper | None",
) -> tuple[str, ...]:
    """Module prefixes that are also an HF root for non-LM weights.

    E.g. Molmo/Phi-4-MM/Muse mark LM as ``model`` while vision HF keys stay
    under ``model.vision_*`` / ``model.embed_tokens_extend.*``. Detected via
    ``WeightsMapper.orig_to_new_prefix`` (no hard-coded name list).
    """
    if weights_mapper is None:
        return ()
    mapping = weights_mapper.orig_to_new_prefix
    if not mapping:
        return ()

    non_lm_origs: list[str] = []
    for orig, new in mapping.items():
        orig_p = orig if orig.endswith(".") else f"{orig}."
        if not _hf_prefix_maps_to_lm(new, lm_module_prefixes):
            non_lm_origs.append(orig_p)

    if not non_lm_origs:
        return ()

    shared = [
        p
        for p in lm_module_prefixes
        if any(orig.startswith(p) for orig in non_lm_origs)
    ]
    return tuple(shared)

@contextmanager
def atomic_writer(
    filepath: str | Path, mode: str = "w", encoding: str | None = None
) -> Generator[IO]:
    """Context manager that provides an atomic file writing routine.

    The context manager writes to a temporary file and, if successful,
    atomically replaces the original file.

    Args:
        filepath (str or Path): The path to the file to write.
        mode (str): The file mode for the temporary file (e.g., 'w', 'wb').
        encoding (str): The encoding for text mode.

    Yields:
        file object: A handle to the temporary file.

    """
    # Create a temporary file in the same directory as the target file
    # to ensure it's on the same filesystem for an atomic replace.
    temp_dir = os.path.dirname(filepath)
    temp_fd, temp_path = tempfile.mkstemp(dir=temp_dir)

    try:
        # Open the temporary file for writing
        with os.fdopen(temp_fd, mode=mode, encoding=encoding) as temp_file:
            yield temp_file

        # If the 'with' block completes successfully,
        # perform the atomic replace.
        os.replace(temp_path, filepath)

    except Exception:
        logger.exception(
            "Error during atomic write. Original file '%s' not modified", filepath
        )
        raise
    finally:
        # Clean up the temporary file if it still exists.
        if os.path.exists(temp_path):
            os.remove(temp_path)

def composed_weight_loader(
    loader: LoaderFunction, fn: Callable[[torch.Tensor], torch.Tensor]
) -> LoaderFunction:
    """Create a weight loader that post-processes the weights after loading"""

    def composed_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
        loader(param, loaded_weight)
        param.data.copy_(fn(param))
        return

    return composed_loader

def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
    """Default weight loader."""
    try:
        if param.numel() == 1 and loaded_weight.numel() == 1:
            # Sometimes scalar values aren't considered tensors with shapes
            # so if both param and loaded_weight are a scalar,
            # reshape to match before copying
            param.data.copy_(loaded_weight.view(param.shape))
        else:
            assert param.size() == loaded_weight.size(), (
                f"Attempted to load weight ({loaded_weight.size()}) "
                f"into parameter ({param.size()})"
            )

            param.data.copy_(loaded_weight)
    except Exception:
        # NOTE: This exception is added for the purpose of setting breakpoint to
        # debug weight loading issues.
        raise

def download_safetensors_index_file_from_hf(
    model_name_or_path: str,
    index_file: str,
    cache_dir: str | None,
    subfolder: str | None = None,
    revision: str | None = None,
) -> None:
    """Download hf safetensors index file from Hugging Face Hub.

    Args:
        model_name_or_path (str): The model name or path.
        index_file (str): The safetensors index file name
        cache_dir (Optional[str]): The cache directory to store the model
            weights. If None, will use HF defaults.
        subfolder (Optional[str]): The subfolder within the model repository
            to download weights from.
        revision (Optional[str]): The revision of the model.

    """
    # Use file lock to prevent multiple processes from
    # downloading the same model weights at the same time.
    with get_lock(model_name_or_path, cache_dir):
        try:
            # Download the safetensors index file.
            hf_api().hf_hub_download(
                repo_id=model_name_or_path,
                filename=index_file,
                cache_dir=cache_dir,
                revision=revision,
                subfolder=subfolder,
                local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
            )
        # If file not found on remote or locally, we should not fail since
        # only some models will have index_file.
        except huggingface_hub.utils.LocalEntryNotFoundError:
            logger.info("No %s found in local cache.", index_file)
        except huggingface_hub.utils.EntryNotFoundError:
            logger.info("No %s found in remote.", index_file)

@instrument(span_name="Download weights - HF")
def download_weights_from_hf(
    model_name_or_path: str,
    cache_dir: str | None,
    allow_patterns: list[str],
    revision: str | None = None,
    subfolder: str | None = None,
    ignore_patterns: str | list[str] | None = None,
) -> str:
    """Download model weights from Hugging Face Hub.

    Args:
        model_name_or_path (str): The model name or path.
        cache_dir (Optional[str]): The cache directory to store the model
            weights. If None, will use HF defaults.
        allow_patterns (list[str]): The allowed patterns for the
            weight files. Files matched by any of the patterns will be
            downloaded.
        revision (Optional[str]): The revision of the model.
        subfolder (Optional[str]): The subfolder within the model repository
            to download weights from.
        ignore_patterns (Optional[Union[str, list[str]]]): The patterns to
            filter out the weight files. Files matched by any of the patterns
            will be ignored.

    Returns:
        str: The path to the downloaded model weights.

    """
    assert len(allow_patterns) > 0
    local_only = huggingface_hub.constants.HF_HUB_OFFLINE
    if not local_only:
        # Attempt to reduce allow_patterns to a single pattern
        # so we only have to call snapshot_download once.
        try:
            fs = hf_fs()
            file_list = fs.ls(
                os.path.join(model_name_or_path, subfolder or ""),
                detail=False,
                revision=revision,
            )

            # If downloading safetensors and an index file exists, use the
            # specific file names from the index to avoid downloading
            # unnecessary files (e.g., from subdirectories like "original/").
            index_file = f"{model_name_or_path}/{SAFE_WEIGHTS_INDEX_NAME}"
            if "*.safetensors" in allow_patterns and index_file in file_list:
                index_path = hf_api().hf_hub_download(
                    repo_id=model_name_or_path,
                    filename=SAFE_WEIGHTS_INDEX_NAME,
                    cache_dir=cache_dir,
                    revision=revision,
                    subfolder=subfolder,
                )
                with open(index_path) as f:
                    weight_map = json.load(f)["weight_map"]
                if weight_map:
                    # Extra [] so that weight_map files are treated as a
                    # single allow_pattern in the loop below
                    allow_patterns = [list(set(weight_map.values()))]  # type: ignore[list-item]
                else:
                    allow_patterns = ["*.safetensors"]
            else:
                # Use the first pattern found in the HF repo's files.
                for pattern in allow_patterns:
                    if fnmatch.filter(file_list, pattern):
                        allow_patterns = [pattern]
                        break
        except Exception as e:
            logger.warning(
                "Failed to get file list for '%s'. Trying each pattern in "
                "allow_patterns individually until weights have been "
                "downloaded. Error: %s",
                model_name_or_path,
                e,
            )

    logger.debug("Using model weights format %s", allow_patterns)
    # Use file lock to prevent multiple processes from
    # downloading the same model weights at the same time.
    with get_lock(model_name_or_path, cache_dir):
        start_time = time.perf_counter()
        for allow_pattern in allow_patterns:
            hf_folder = hf_api().snapshot_download(
                model_name_or_path,
                allow_patterns=allow_pattern,
                ignore_patterns=ignore_patterns,
                cache_dir=cache_dir,
                tqdm_class=DisabledTqdm,
                revision=revision,
                local_files_only=local_only,
            )
            # If we have downloaded weights for this allow_pattern,
            # we don't need to check the rest.
            # allow_pattern can be a list (from weight_map) or str (glob)
            if isinstance(allow_pattern, list):
                break
            if any(Path(hf_folder).glob(allow_pattern)):
                break
        time_taken = time.perf_counter() - start_time
        if time_taken > 0.5:
            logger.info(
                "Time spent downloading weights for %s: %.6f seconds",
                model_name_or_path,
                time_taken,
            )
    return hf_folder

def enable_xet_high_performance():
    """Automatically activates xet high performance mode"""
    if "HF_XET_HIGH_PERFORMANCE" not in os.environ:
        huggingface_hub.constants.HF_XET_HIGH_PERFORMANCE = True

def fastsafetensors_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files
    using fastsafetensor library.

    Uses ParallelLoader for pipelined loading: the producer thread
    prepares metadata for the next shard while the consumer yields
    tensors from the current shard.
    """
    from fastsafetensors.parallel_loader import ParallelLoader

    if torch.distributed.is_initialized():
        pg = torch.distributed.group.WORLD
    else:
        pg = SingleGroup()

    device = torch.device(f"cuda:{current_platform.current_device()}")
    hf_weights_files = sorted(hf_weights_files, key=_natural_sort_key)

    # Use nogds=True for TP > 1 to avoid cuFileDriverOpen() which
    # initializes the GDS DMA subsystem for all visible GPUs, creating
    # unwanted CUDA contexts on every device.
    nogds = pg.size() > 1

    queue_size = envs.VLLM_FASTSAFETENSORS_QUEUE_SIZE
    tqdm_enabled = enable_tqdm(use_tqdm_on_load)

    def _make_loader(nogds: bool) -> "ParallelLoader":
        return ParallelLoader(
            pg=pg,
            hf_weights_files=hf_weights_files,
            queue_size=queue_size,
            use_tqdm_on_load=tqdm_enabled,
            device=str(device),
            nogds=nogds,
        )

    # GDS can fail either at construction or lazily inside the producer
    # thread during iteration (e.g. cuFileHandleRegister returning
    # CU_FILE_HANDLE_NOT_REGISTERED on a filesystem without GDS support).
    # Catch both and fall back to nogds, but only before yielding any
    # tensor -- restarting mid-stream would reload earlier shards.
    pl = None
    yielded = False
    try:
        try:
            pl = _make_loader(nogds)
            for name, tensor in pl.iterate_weights():
                yielded = True
                yield name, tensor
        except RuntimeError as e:
            if nogds or yielded or "gds" not in str(e):
                raise
            logger.warning_once(
                "GDS not enabled, setting `nogds=True`.\n"
                "For more information, see: https://github.com/foundation-model-stack/"
                "fastsafetensors?tab=readme-ov-file#basic-api-usages"
            )
            if pl is not None:
                pl.close()
            pl = _make_loader(nogds=True)
            yield from pl.iterate_weights()
    finally:
        if pl is not None:
            pl.close()

def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]:
    """Exclude files that are not needed for inference.

    See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
    """
    blacklist = [
        "training_args.bin",
        "optimizer.bin",
        "optimizer.pt",
        "scheduler.pt",
        "scaler.pt",
    ]
    hf_weights_files = [
        f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist)
    ]
    return hf_weights_files

def filter_mm_encoder_only_safetensors_files(
    hf_weights_files: list[str],
    hf_folder: str,
    index_file: str,
    language_model_prefixes: Iterable[str],
    *,
    weights_mapper: "WeightsMapper | None" = None,
) -> list[str]:
    """Drop safetensors shards that only contain language-model weights.

    Used with ``--mm-encoder-only`` so Encoder-only EPD instances avoid reading
    pure LM shards from disk/DRAM.

    Each HF index key is classified by mapping through ``weights_mapper`` (when
    provided) and testing the *vLLM* name against ``language_model_prefixes``.
    Without a mapper, HF names are compared directly (identity checkpoint
    layout, e.g. Kimi ``language_model.*``).

    A shard is kept if it contains any non-LM key. Without an index file,
    returns ``hf_weights_files`` unchanged.
    """
    prefixes = tuple(language_model_prefixes)
    if not prefixes:
        return hf_weights_files

    index_path = os.path.join(hf_folder, index_file)
    if not os.path.isfile(index_path):
        logger.warning_once(
            "mm-encoder-only shard filter skipped: index file %s not found",
            index_path,
        )
        return hf_weights_files

    with open(index_path) as f:
        weight_map: dict[str, str] = json.load(f)["weight_map"]

    keys_by_file: dict[str, list[str]] = {}
    for weight_name, weight_file in weight_map.items():
        keys_by_file.setdefault(weight_file, []).append(weight_name)

    def _is_lm_key(key: str) -> bool:
        mapped = _mapped_weight_name(weights_mapper, key)
        if mapped is None:
            return False
        return any(mapped.startswith(p) for p in prefixes)

    def _is_lm_only(weight_file: str) -> bool:
        keys = keys_by_file.get(weight_file)
        if not keys:
            # Not listed in the index (e.g. extra matched glob) — keep.
            return False
        return all(_is_lm_key(key) for key in keys)

    kept: list[str] = []
    skipped = 0
    for path in hf_weights_files:
        rel = os.path.relpath(path, hf_folder)
        # Index weight_map values are usually basenames; tolerate either form.
        basename = os.path.basename(path)
        if _is_lm_only(rel) or _is_lm_only(basename):
            skipped += 1
            continue
        kept.append(path)

    if skipped:
        logger.info_once(
            "mm-encoder-only: skipped %d/%d safetensors shard(s) that only "
            "contain language-model weights (prefixes=%s)",
            skipped,
            skipped + len(kept),
            prefixes,
        )
    return kept

def instanttensor_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files
    using instanttensor library."""
    try:
        import instanttensor
    except ImportError as e:
        raise ImportError(
            "Please install instanttensor via `pip install vllm[instanttensor]`"
        ) from e

    if not current_platform.is_cuda():
        raise ValueError("InstantTensor requires NVIDIA GPUs")

    try:
        world_group = get_world_group()
    except AssertionError:
        # Entering here only in unit tests where the world group is not initialized.
        process_group = None
    else:
        process_group = world_group.device_group if world_group.world_size > 1 else None

    device = current_platform.current_device()

    # copy=True yields tensors that own their memory, staying valid after the
    # context exits or InstantTensor reuses its buffer.
    with instanttensor.safe_open(
        hf_weights_files,
        framework="pt",
        device=device,
        process_group=process_group,
        copy=True,
    ) as f:
        # Track bytes so the bar reports load throughput (GB/s).
        pbar = tqdm(
            total=f.total_tensor_size,
            desc="Loading safetensors using InstantTensor loader",
            disable=not enable_tqdm(use_tqdm_on_load),
            bar_format=_BAR_FORMAT,
            position=tqdm._get_free_pos(),
            unit="B",
            unit_scale=True,
            unit_divisor=1024,
            mininterval=1.0,
        )
        try:
            for name, tensor in f.tensors():
                pbar.update(tensor.numel() * tensor.element_size())
                yield name, tensor
        finally:
            pbar.close()

def maybe_download_from_modelscope(
    model: str,
    revision: str | None = None,
    download_dir: str | None = None,
    ignore_patterns: str | list[str] | None = None,
    allow_patterns: list[str] | str | None = None,
) -> str | None:
    """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.

    Returns the path to the downloaded model, or None if the model is not
    downloaded from ModelScope."""
    if envs.VLLM_USE_MODELSCOPE:
        # download model from ModelScope hub,
        # lazy import so that modelscope is not required for normal use.
        # pylint: disable=C.
        from modelscope.hub.snapshot_download import snapshot_download

        # Use file lock to prevent multiple processes from
        # downloading the same model weights at the same time.
        with get_lock(model, download_dir):
            if not os.path.exists(model):
                model_path = snapshot_download(
                    model_id=model,
                    cache_dir=download_dir,
                    local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
                    revision=revision,
                    ignore_file_pattern=ignore_patterns,
                    allow_patterns=allow_patterns,
                )
            else:
                model_path = model
        return model_path
    return None

def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
    """Remap the name of FP8 k/v_scale parameters.

    This function handles the remapping of FP8 k/v_scale parameter names.
    It detects if the given name ends with a suffix and attempts to remap
    it to the expected name format in the model. If the remapped name is not
    found in the params_dict, a warning is printed and None is returned.

    Args:
        name (str): The original loaded checkpoint parameter name.
        params_dict (dict): Dictionary containing the model's named parameters.

    Returns:
        str: The remapped parameter name if successful, or the original name
             if no remapping is needed.
        None: If the remapped name is not found in params_dict.

    """
    # Already in vLLM's expected form (e.g. weights pre-renamed by a
    # `WeightsMapper` from the quant config). Skip the regex remap, which
    # would otherwise double-apply the `.attn` prefix and drop the weight.
    if name in params_dict:
        return name
    if name.endswith(".kv_scale"):
        logger.warning_once(
            "DEPRECATED. Found kv_scale in the checkpoint. "
            "This format is deprecated in favor of separate k_scale and "
            "v_scale tensors and will be removed in a future release. "
            "Functionally, we will remap kv_scale to k_scale and duplicate "
            "k_scale to v_scale"
        )
        # NOTE: we remap the deprecated kv_scale to k_scale
        remapped_name = name.replace(".kv_scale", ".attn.k_scale")
        if remapped_name not in params_dict:
            logger.warning_once(
                "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.",  #  noqa: E501
                name,
                remapped_name,
            )
            return None
        return remapped_name

    if any("mla_attn" in key for key in params_dict):
        attn_str = "mla_attn.mla_attn"
        logger.debug_once(
            f"Found mla_attn with k_scale and v_scale in "
            f"the checkpoint, using {attn_str} as attn_str"
        )
    else:
        attn_str = "attn"
    # Define scale name mapping patterns in order of precedence
    scale_mapping_patterns = [
        # ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale ->
        # .self_attn.attn.{k,v}_scale
        (
            r"\.self_attn\.([kv])_proj\.([kv])_scale$",
            rf".self_attn.{attn_str}.\2_scale",
        ),
        # QKV proj format: .self_attn.qkv_proj.{k,v}_scale ->
        # .self_attn.attn.{k,v}_scale
        (r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),
        # Qwen3 MoE format: .self_attn.qkqkv_proj.{k,v}_scale ->
        # .self_attn.attn.{k,v}_scale
        (r"\.self_attn\.qkqkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),
        # NemotronH format: .mixer.{k,v}_proj.{k,v}_scale ->
        # .mixer.attn.{k,v}_scale
        (r"\.mixer\.[kv]_proj\.([kv])_scale$", r".mixer.attn.\1_scale"),
        # HYV3 format: .self_attn.q.scale -> .self_attn.attn.q_scale
        (r"\.self_attn\.q\.scale$", r".self_attn.attn.q_scale"),
        # HYV3 format: .self_attn.{k,v}_cache.scale ->
        # .self_attn.attn.{k,v}_scale
        (r"\.self_attn\.([kv])_cache\.scale$", r".self_attn.attn.\1_scale"),
        # Default format: .{k,v}_scale -> .attn.{k,v}_scale
        (r"\.([qkv])_scale$", r".attn.\1_scale"),
        (r"\.([qkv])_zero_point$", r".attn.\1_zero_point"),
    ]

    # Check if name ends with k_scale or v_scale
    if name.endswith(
        (
            ".k_scale",
            ".v_scale",
            ".q_scale",
            ".k_zero_point",
            ".v_zero_point",
            ".q_zero_point",
            ".q.scale",
            ".k_cache.scale",
            ".v_cache.scale",
        )
    ):
        import regex as re

        for pattern, replacement in scale_mapping_patterns:
            if re.search(pattern, name):
                remapped_name = re.sub(pattern, replacement, name)
                if remapped_name not in params_dict:
                    scale_type = name.split(".")[-1]
                    logger.warning_once(
                        "Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.",  # noqa: E501
                        scale_type,
                        name,
                        remapped_name,
                        scale_type,
                    )
                    return None
                return remapped_name

    # If there were no matches, return the untouched param name
    return name

def maybe_remap_moe_expert_param_name(
    name: str,
    params_dict: dict[str, torch.nn.Parameter],
) -> str:
    """Remap MoE expert parameter names to account for routed_experts hierarchy.

    This handles the transition from the old FusedMoE structure where weights
    were directly in the experts module, to the new MoERunner → RoutedExperts
    structure.

    Checkpoint weights have names like:
        layers.0.mlp.experts.w13_weight
        layers.0.feed_forward.experts.w2_input_scale
    But actual parameters are now:
        layers.0.mlp.experts.routed_experts.w13_weight
        layers.0.feed_forward.experts.routed_experts.w2_input_scale

    This function inserts 'routed_experts.' into the path when needed.

    Args:
        name: Parameter name from checkpoint
        params_dict: Dictionary of model parameters (from named_parameters())

    Returns:
        Remapped parameter name if routed_experts hierarchy exists,
        otherwise the original name

    """
    # Only remap if this looks like an expert parameter
    if ".experts." not in name:
        return name

    # Skip if already has routed_experts
    if ".experts.routed_experts." in name:
        return name

    # Expert parameter patterns to check
    expert_param_suffixes = [
        "w13_weight",
        "w2_weight",
        "w13_weight_scale",
        "w2_weight_scale",
        "w13_input_scale",
        "w2_input_scale",
        "w13_bias",
        "w2_bias",
        "w13_scale",
        "w2_scale",
        "w13_qweight",
        "w2_qweight",
        "w13_qzeros",
        "w2_qzeros",
        "w13_weight_shape",
        "w2_weight_shape",
    ]

    # Check if this is an expert weight parameter
    is_expert_param = any(
        f".{suffix}" in name or name.endswith(suffix)
        for suffix in expert_param_suffixes
    )

    if not is_expert_param:
        return name

    # Try inserting routed_experts after .experts.
    new_name = name.replace(".experts.", ".experts.routed_experts.", 1)

    # Only use the new name if it exists in the model
    if new_name in params_dict:
        return new_name

    # Otherwise return original name (old checkpoint format or different structure)
    return name

def multi_thread_pt_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
    pt_load_map_location: str | dict[str, str] = "cpu",
    max_workers: int = 4,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Multi-Thread iterate over the weights in the model bin/pt files."""

    def _load_file(bin_file: str):
        return torch.load(
            bin_file, map_location=pt_load_map_location, weights_only=True
        )

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(_load_file, bin_file) for bin_file in hf_weights_files
        ]
        futures_iter = tqdm(
            concurrent.futures.as_completed(futures),
            total=len(hf_weights_files),
            desc="Multi-thread loading pt checkpoint shards",
            disable=not enable_tqdm(use_tqdm_on_load),
            bar_format=_BAR_FORMAT,
        )

        for future in futures_iter:
            state = future.result()
            yield from state.items()
            del state

def multi_thread_safetensors_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
    max_workers: int = 4,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Multi-Thread iterate over the weights in the model safetensor files."""

    def _load_file(st_file: str):
        result = load_file(st_file, device="cpu")
        return result

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Note to use generator here so we do not store all the loaded files in memory
        # at the same time, which can cause OOM for large models.
        futures = (executor.submit(_load_file, st_file) for st_file in hf_weights_files)
        futures_iter = tqdm(
            concurrent.futures.as_completed(futures),
            total=len(hf_weights_files),
            desc="Multi-thread loading shards",
            disable=not enable_tqdm(use_tqdm_on_load),
            bar_format=_BAR_FORMAT,
        )

        for future in futures_iter:
            state_dict = future.result()
            del future
            for key in list(state_dict):
                yield key, state_dict.pop(key)