vllm.model_executor.model_loader.weight_utils

Utilities for downloading and initializing model weights.

Functions:

• atomic_writer –

  Context manager that provides an atomic file writing routine.

• composed_weight_loader –

  Create a weight loader that post-processes the weights after loading

• default_weight_loader –

  Default weight loader.

• download_safetensors_index_file_from_hf –

  Download hf safetensors index file from Hugging Face Hub.

• download_weights_from_hf –

  Download model weights from Hugging Face Hub.

• enable_xet_high_performance –

  Automatically activates xet high performance mode

• fastsafetensors_weights_iterator –

  Iterate over the weights in the model safetensor files

• filter_files_not_needed_for_inference –

  Exclude files that are not needed for inference.

• filter_mm_encoder_only_safetensors_files –

  Drop safetensors shards that only contain language-model weights.

• filter_safetensors_files_by_weight_name –

  Drop safetensors shards in which is_unused_weight accepts every tensor.

• instanttensor_weights_iterator –

  Iterate over the weights in the model safetensor files

• maybe_download_from_modelscope –

  Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.

• maybe_remap_kv_scale_name –

  Remap the name of FP8 k/v_scale parameters.

• maybe_remap_moe_expert_param_name –

  Remap MoE expert parameter names to account for routed_experts hierarchy.

• multi_thread_pt_weights_iterator –

  Multi-Thread iterate over the weights in the model bin/pt files.

• multi_thread_safetensors_weights_iterator –

  Multi-Thread iterate over the weights in the model safetensor files.

• np_cache_weights_iterator –

  Iterate over the weights in the model np files.

• pt_weights_iterator –

  Iterate over the weights in the model bin/pt files.

• remap_moe_expert_weights –

  Remap MoE expert parameter names for backward compatibility.

• resolve_mm_encoder_only_lm_prefixes –

  Resolve vLLM LM module prefixes for --mm-encoder-only shard skip.

• row_parallel_weight_loader –

  Load weights that are row-parallelized.

• runai_safetensors_weights_iterator –

  Iterate over the weights in the model safetensor files.

• safetensors_weights_iterator –

  Iterate over the weights in the model safetensor files.

• sharded_weight_loader –

  Create a weight loader that shards the weights along the given axis

_get_available_ram_bytes()

Return available RAM, honoring cgroup limits.

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

_get_checkpoints_size_bytes(files)

Return the total size of the checkpoint files in bytes.

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

_get_fs_type(files)

Get the filesystem type of the first file in files (Linux only).

Source code in vllm/model_executor/model_loader/weight_utils.py

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 ""

_mapped_weight_name(weights_mapper, key)

Apply WeightsMapper.map_name when present; else identity.

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

_natural_sort_key(filepath)

Natural sort key for filenames with numeric components, such as model-00001-of-00005.safetensors -> ['model-', 1, '-of-', 5, '.safetensors']

Source code in vllm/model_executor/model_loader/weight_utils.py

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))
    ]

_prefetch_all_checkpoints(sorted_files, num_prefetch_threads=DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, block_size=DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE)

Start prefetching checkpoint files into page cache in a background thread.

Source code in vllm/model_executor/model_loader/weight_utils.py

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()

_prefetch_checkpoint(file_path, block_size=DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE)

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.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

_shared_hf_root_module_prefixes(lm_module_prefixes, weights_mapper)

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).

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

atomic_writer(filepath, mode='w', encoding=None)

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.

Parameters:

• filepath

  (str or Path) –

  The path to the file to write.

• mode

  (str, default: 'w' ) –

  The file mode for the temporary file (e.g., 'w', 'wb').

• encoding

  (str, default: None ) –

  The encoding for text mode.

Yields:

• Generator[IO] –

  file object: A handle to the temporary file.

Source code in vllm/model_executor/model_loader/weight_utils.py

@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)

composed_weight_loader(loader, fn)

Create a weight loader that post-processes the weights after loading

Source code in vllm/model_executor/model_loader/weight_utils.py

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

default_weight_loader(param, loaded_weight)

Default weight loader.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

download_safetensors_index_file_from_hf(model_name_or_path, index_file, cache_dir, subfolder=None, revision=None)

Download hf safetensors index file from Hugging Face Hub.

Parameters:

• 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], default: None ) –

  The subfolder within the model repository to download weights from.

• revision

  (Optional[str], default: None ) –

  The revision of the model.

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

download_weights_from_hf(model_name_or_path, cache_dir, allow_patterns, revision=None, subfolder=None, ignore_patterns=None)

Download model weights from Hugging Face Hub.

Parameters:

• 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], default: None ) –

  The revision of the model.

• subfolder

  (Optional[str], default: None ) –

  The subfolder within the model repository to download weights from.

• ignore_patterns

  (Optional[Union[str, list[str]]], default: None ) –

  The patterns to filter out the weight files. Files matched by any of the patterns will be ignored.

Returns:

• str ( str ) –

  The path to the downloaded model weights.

Source code in vllm/model_executor/model_loader/weight_utils.py

@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

enable_xet_high_performance()

Automatically activates xet high performance mode

Source code in vllm/model_executor/model_loader/weight_utils.py

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

fastsafetensors_weights_iterator(hf_weights_files, use_tqdm_on_load)

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.

Source code in vllm/model_executor/model_loader/weight_utils.py

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()

filter_files_not_needed_for_inference(hf_weights_files)

Exclude files that are not needed for inference.

See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233

Source code in vllm/model_executor/model_loader/weight_utils.py

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

filter_mm_encoder_only_safetensors_files(hf_weights_files, hf_folder, index_file, language_model_prefixes, *, weights_mapper=None)

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.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

filter_safetensors_files_by_weight_name(hf_weights_files, is_unused_weight)

Drop safetensors shards in which is_unused_weight accepts every tensor.

Loaders that read whole files (InstantTensor, fastsafetensors, multi-thread, eager, prefetch) cannot skip single tensors, so dropping shards is their only way to avoid reading weights the model does not load.

Parameters:

• hf_weights_files

  (list[str]) –

  Safetensors shard paths.

• is_unused_weight

  (Callable[[str], bool]) –

  Returns True for checkpoint weight names the model does not load.

Returns:

• list[str] –

  The shards holding at least one wanted tensor, or hf_weights_files

• list[str] –

  unchanged when no shard does.

Source code in vllm/model_executor/model_loader/weight_utils.py

def filter_safetensors_files_by_weight_name(
    hf_weights_files: list[str],
    is_unused_weight: Callable[[str], bool],
) -> list[str]:
    """Drop safetensors shards in which `is_unused_weight` accepts every tensor.

    Loaders that read whole files (InstantTensor, fastsafetensors, multi-thread,
    eager, prefetch) cannot skip single tensors, so dropping shards is their
    only way to avoid reading weights the model does not load.

    Args:
        hf_weights_files: Safetensors shard paths.
        is_unused_weight: Returns True for checkpoint weight names the model does
            not load.

    Returns:
        The shards holding at least one wanted tensor, or `hf_weights_files`
        unchanged when no shard does.

    """
    kept: list[str] = []
    for st_file in hf_weights_files:
        with safe_open(st_file, framework="pt") as f:
            if not all(is_unused_weight(name) for name in f.keys()):  # noqa: SIM118
                kept.append(st_file)
    if not kept:
        return hf_weights_files
    if len(kept) < len(hf_weights_files):
        logger.info_once(
            "Skipped %d/%d safetensors shard(s) holding no weights the model loads",
            len(hf_weights_files) - len(kept),
            len(hf_weights_files),
        )
    return kept

instanttensor_weights_iterator(hf_weights_files, use_tqdm_on_load)

Iterate over the weights in the model safetensor files using instanttensor library.

Source code in vllm/model_executor/model_loader/weight_utils.py

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()

maybe_download_from_modelscope(model, revision=None, download_dir=None, ignore_patterns=None, allow_patterns=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.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

maybe_remap_kv_scale_name(name, params_dict)

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.

Parameters:

• name

  (str) –

  The original loaded checkpoint parameter name.

• params_dict

  (dict) –

  Dictionary containing the model's named parameters.

Returns:

• str ( str | None ) –

  The remapped parameter name if successful, or the original name if no remapping is needed.

• None ( str | None ) –

  If the remapped name is not found in params_dict.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

maybe_remap_moe_expert_param_name(name, params_dict)

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.

Parameters:

• name

  (str) –

  Parameter name from checkpoint

• params_dict

  (dict[str, Parameter]) –

  Dictionary of model parameters (from named_parameters())

Returns:

• str –

  Remapped parameter name if routed_experts hierarchy exists,

• str –

  otherwise the original name

Source code in vllm/model_executor/model_loader/weight_utils.py

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

multi_thread_pt_weights_iterator(hf_weights_files, use_tqdm_on_load, pt_load_map_location='cpu', max_workers=4)

Multi-Thread iterate over the weights in the model bin/pt files.

Source code in vllm/model_executor/model_loader/weight_utils.py

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

multi_thread_safetensors_weights_iterator(hf_weights_files, use_tqdm_on_load, max_workers=4)

Multi-Thread iterate over the weights in the model safetensor files.

Source code in vllm/model_executor/model_loader/weight_utils.py

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)

np_cache_weights_iterator(model_name_or_path, cache_dir, hf_folder, hf_weights_files, use_tqdm_on_load)

Iterate over the weights in the model np files.

Will dump the model weights to numpy files if they are not already dumped.

Source code in vllm/model_executor/model_loader/weight_utils.py

def np_cache_weights_iterator(
    model_name_or_path: str,
    cache_dir: str | None,
    hf_folder: str,
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model np files.

    Will dump the model weights to numpy files if they are not already dumped.
    """
    # Convert the model weights from torch tensors to numpy arrays for
    # faster loading.
    np_folder = os.path.join(hf_folder, "np")
    os.makedirs(np_folder, exist_ok=True)
    weight_names_file = os.path.join(np_folder, "weight_names.json")
    # Use file lock to prevent multiple processes from
    # dumping the same model weights to numpy at the same time.
    with get_lock(model_name_or_path, cache_dir):
        if not os.path.exists(weight_names_file):
            weight_names: list[str] = []
            for bin_file in tqdm(
                hf_weights_files,
                desc="Loading np_cache checkpoint shards",
                disable=not enable_tqdm(use_tqdm_on_load),
                bar_format=_BAR_FORMAT,
            ):
                state = torch.load(bin_file, map_location="cpu", weights_only=True)
                for name, param in state.items():
                    param_path = os.path.join(np_folder, name)
                    with open(param_path, "wb") as f:
                        np.save(f, param.cpu().detach().numpy())
                    weight_names.append(name)
            with open(weight_names_file, "w") as f:
                json.dump(weight_names, f)

    with open(weight_names_file) as f:
        weight_names = json.load(f)

    for name in weight_names:
        param_path = os.path.join(np_folder, name)
        with open(param_path, "rb") as f:
            param = np.load(f)
        yield name, torch.from_numpy(param)

pt_weights_iterator(hf_weights_files, use_tqdm_on_load, pt_load_map_location='cpu')

Iterate over the weights in the model bin/pt files.

Source code in vllm/model_executor/model_loader/weight_utils.py

def pt_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
    pt_load_map_location: str | dict[str, str] = "cpu",
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model bin/pt files."""
    for bin_file in tqdm(
        hf_weights_files,
        desc="Loading pt checkpoint shards",
        disable=not enable_tqdm(use_tqdm_on_load),
        bar_format=_BAR_FORMAT,
    ):
        state = torch.load(
            bin_file, map_location=pt_load_map_location, weights_only=True
        )
        yield from state.items()
        del state

remap_moe_expert_weights(weights, params_dict)

Remap MoE expert parameter names for backward compatibility.

This allows models with custom weight loading to automatically handle both old and new checkpoint formats without needing model-specific remapping code.

Usage

params_dict = dict(model.named_parameters()) for name, weight in remap_moe_expert_weights(weights, params_dict): # name is automatically remapped if needed param = params_dict[name] ...

Parameters:

• weights

  (Iterable[tuple[str, Tensor]]) –

  Iterator of (name, tensor) tuples from checkpoint

• params_dict

  (dict[str, Parameter]) –

  Dictionary of model parameters (from named_parameters())

Yields:

• tuple[str, Tensor] –

  (remapped_name, tensor) tuples

Source code in vllm/model_executor/model_loader/weight_utils.py

def remap_moe_expert_weights(
    weights: Iterable[tuple[str, torch.Tensor]],
    params_dict: dict[str, torch.nn.Parameter],
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Remap MoE expert parameter names for backward compatibility.

    This allows models with custom weight loading to automatically handle both old
    and new checkpoint formats without needing model-specific remapping code.

    Usage:
        params_dict = dict(model.named_parameters())
        for name, weight in remap_moe_expert_weights(weights, params_dict):
            # name is automatically remapped if needed
            param = params_dict[name]
            ...

    Args:
        weights: Iterator of (name, tensor) tuples from checkpoint
        params_dict: Dictionary of model parameters (from named_parameters())

    Yields:
        (remapped_name, tensor) tuples

    """
    for name, weight in weights:
        remapped_name = maybe_remap_moe_expert_param_name(name, params_dict)
        yield (remapped_name, weight)

resolve_mm_encoder_only_lm_prefixes(language_model_names, *, weights_mapper=None)

Resolve vLLM LM module prefixes for --mm-encoder-only shard skip.

Prefixes come from _language_model_names. Classification of HF index keys is done later via optional weights_mapper (see filter_mm_encoder_only_safetensors_files), so this helper does not hard-code Qwen/HF nest lists.

Returns None (leave the safetensors file list unchanged) when _language_model_names is empty/missing, or when a module prefix is a shared HF checkpoint root for both LM and non-LM weights — fail-closed so Molmo / Phi-4-MM / Muse cannot under-load the encoder.

Source code in vllm/model_executor/model_loader/weight_utils.py

def resolve_mm_encoder_only_lm_prefixes(
    language_model_names: Iterable[str] | None,
    *,
    weights_mapper: "WeightsMapper | None" = None,
) -> tuple[str, ...] | None:
    """Resolve vLLM LM *module* prefixes for ``--mm-encoder-only`` shard skip.

    Prefixes come from ``_language_model_names``. Classification of HF index
    keys is done later via optional ``weights_mapper`` (see
    ``filter_mm_encoder_only_safetensors_files``), so this helper does **not**
    hard-code Qwen/HF nest lists.

    Returns ``None`` (leave the safetensors file list unchanged) when
    ``_language_model_names`` is empty/missing, or when a module prefix is a
    shared HF checkpoint root for both LM and non-LM weights — fail-closed so
    Molmo / Phi-4-MM / Muse cannot under-load the encoder.
    """
    prefixes = _normalize_module_prefixes(language_model_names)
    if not prefixes:
        return None
    shared = _shared_hf_root_module_prefixes(prefixes, weights_mapper)
    if shared:
        logger.warning_once(
            "mm-encoder-only whole-shard filter not applied: LM module "
            "prefix(es) %s share an HF checkpoint root with non-LM weights "
            "(via WeightsMapper). Encoder-only load keeps the full "
            "safetensors file list; finer HF prefixes or key-level "
            "filtering can enable skip later.",
            shared,
        )
        return None
    return prefixes

row_parallel_weight_loader(param, loaded_weight)

Load weights that are row-parallelized.

Source code in vllm/model_executor/model_loader/weight_utils.py

def row_parallel_weight_loader(
    param: torch.Tensor, loaded_weight: torch.Tensor
) -> None:
    """Load weights that are row-parallelized."""
    tp_rank = get_tensor_model_parallel_rank()
    shard_dim = 0 if param.dim() != 1 else None

    if shard_dim is not None:
        shard_size = param.data.shape[shard_dim]
        start_idx = tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)

    return default_weight_loader(param, loaded_weight)

runai_safetensors_weights_iterator(hf_weights_files, use_tqdm_on_load, is_distributed=False)

Iterate over the weights in the model safetensor files.

Source code in vllm/model_executor/model_loader/weight_utils.py

def runai_safetensors_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
    is_distributed: bool = False,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files."""
    with SafetensorsStreamer() as streamer:
        is_cuda_alike = current_platform.is_cuda_alike()
        device = (
            f"cuda:{current_platform.current_device()}"
            if is_distributed and is_cuda_alike
            else "cpu"
        )

        streamer.stream_files(
            hf_weights_files,
            device=device,
            is_distributed=is_distributed,
        )
        total_tensors = sum(
            len(tensors_meta)
            for tensors_meta in streamer.files_to_tensors_metadata.values()
        )

        tensor_iter = tqdm(
            streamer.get_tensors(),
            total=total_tensors,
            desc="Loading safetensors using Runai Model Streamer",
            bar_format=_BAR_FORMAT,
            disable=not enable_tqdm(use_tqdm_on_load),
            mininterval=2,
        )

        for name, tensor in tensor_iter:
            yield name, tensor.clone()

safetensors_weights_iterator(hf_weights_files, use_tqdm_on_load, safetensors_load_strategy=None, local_expert_ids=None, *, safetensors_prefetch_num_threads=DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS, safetensors_prefetch_block_size=DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE)

Iterate over the weights in the model safetensor files.

When local_expert_ids is provided, expert weights not belonging to this rank are skipped before reading from disk, which drastically reduces storage I/O for MoE models under EP.

Source code in vllm/model_executor/model_loader/weight_utils.py

def safetensors_weights_iterator(
    hf_weights_files: list[str],
    use_tqdm_on_load: bool,
    safetensors_load_strategy: str | None = None,
    local_expert_ids: set[int] | None = None,
    *,
    safetensors_prefetch_num_threads: int = DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS,
    safetensors_prefetch_block_size: int = DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE,
) -> Generator[tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files.

    When *local_expert_ids* is provided, expert weights not belonging to
    this rank are skipped **before** reading from disk, which drastically
    reduces storage I/O for MoE models under EP.
    """
    loading_desc = "Loading safetensors checkpoint shards"
    if safetensors_load_strategy == "eager":
        loading_desc += " (eager)"

    sorted_files = sorted(hf_weights_files, key=_natural_sort_key)

    fs_type = _get_fs_type(sorted_files)
    is_net_fs = fs_type in ("nfs", "nfs4", "lustre")
    total_bytes = _get_checkpoints_size_bytes(sorted_files)
    avail_bytes = _get_available_ram_bytes()
    ram_threshold_pct = 90
    fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes
    fs_name = fs_type.upper() if fs_type else "unknown"

    logger.info_once(
        "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. "
        "Available RAM: %.2f GiB.",
        fs_name,
        total_bytes / 1024**3,
        avail_bytes / 1024**3,
    )

    should_prefetch = safetensors_load_strategy == "prefetch"
    if safetensors_load_strategy is None:
        if is_net_fs and fits_in_ram:
            should_prefetch = True
        elif is_net_fs and not fits_in_ram:
            logger.warning_once(
                "Network filesystem (%s) detected but checkpoint total size "
                "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). "
                "Skipping auto-prefetch.",
                fs_name,
                total_bytes / 1024**3,
                ram_threshold_pct,
                avail_bytes / 1024**3,
            )
        elif not is_net_fs and fits_in_ram:
            logger.info_once(
                "Auto-prefetch is disabled because the filesystem (%s) is not a "
                "recognized network FS (NFS/Lustre). If you want to force "
                "prefetching, start vLLM with --safetensors-load-strategy=prefetch.",
                fs_name,
            )
        elif not is_net_fs and not fits_in_ram:
            logger.info_once(
                "Auto-prefetch is disabled because the filesystem (%s) is not a "
                "recognized network FS (NFS/Lustre) and the checkpoint size "
                "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).",
                fs_name,
                total_bytes / 1024**3,
                ram_threshold_pct,
                avail_bytes / 1024**3,
            )
    elif should_prefetch and not fits_in_ram:
        logger.warning_once(
            "safetensors_load_strategy='prefetch' was explicitly specified, but "
            "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM "
            "(%.2f GiB). This may cause out-of-memory errors.",
            total_bytes / 1024**3,
            ram_threshold_pct,
            avail_bytes / 1024**3,
        )

    if should_prefetch:
        _prefetch_all_checkpoints(
            sorted_files,
            num_prefetch_threads=safetensors_prefetch_num_threads,
            block_size=safetensors_prefetch_block_size,
        )

    leftover_state_dict: dict[str, torch.Tensor] = {}
    for st_file in tqdm(
        sorted_files,
        desc=loading_desc,
        disable=not enable_tqdm(use_tqdm_on_load),
        bar_format=_BAR_FORMAT,
    ):
        if safetensors_load_strategy == "eager":
            with open(st_file, "rb") as f:
                state_dict = load(f.read())
            for name, param in state_dict.items():
                if not should_skip_weight(name, local_expert_ids):
                    yield name, param
        elif safetensors_load_strategy == "torchao":
            # we can't load flattened torchao tensor subclasses directly into the model
            # instead we reconstruct the subclasses here before returning
            if not torchao_version_at_least("0.15.0"):
                raise ValueError(
                    "Please use torchao version >= 0.15.0 "
                    "to load torchao safetensors checkpoint"
                )
            from torchao.prototype.safetensors.safetensors_support import (
                unflatten_tensor_state_dict,
            )

            with safe_open(st_file, framework="pt") as f:
                state_dict = {}
                for name in f.keys():  # noqa: SIM118
                    if should_skip_weight(name, local_expert_ids):
                        continue
                    state_dict[name] = f.get_tensor(name)

                # update with leftover tensor data from previous iteration, if any
                state_dict.update(leftover_state_dict)
                metadata = f.metadata()
                # due to sharded checkpoints, we are not guaranteed that we have all
                # tensor subclass data on one file
                # state_dict has the leftover data from this step and we wait for
                # missing information to be provided in a future iteration
                unflattened_state_dict, leftover_state_dict = (
                    unflatten_tensor_state_dict(state_dict, metadata)
                )
            yield from unflattened_state_dict.items()
        else:
            with safe_open(st_file, framework="pt") as f:
                for name in f.keys():  # noqa: SIM118
                    if should_skip_weight(name, local_expert_ids):
                        continue
                    param = f.get_tensor(name)
                    yield name, param

sharded_weight_loader(shard_axis)

Create a weight loader that shards the weights along the given axis

Source code in vllm/model_executor/model_loader/weight_utils.py

def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
    """Create a weight loader that shards the weights along the given axis"""

    def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
        tp_rank = get_tensor_model_parallel_rank()

        shard_size = param.data.shape[shard_axis]
        start_idx = tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)

        return default_weight_loader(param, loaded_weight)

    return loader