311 lines
12 KiB
Python
311 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Component weight loading for MiniMax Music 3.
|
|
|
|
The checkpoint is a multi-component repo. The Qwen3 backbone is a normal
|
|
``Qwen3ForCausalLM`` under ``language_model/`` and loads through vLLM. The
|
|
audio embedding table and the depth decoder live in ``rvq_depth_decoder/``,
|
|
and the acoustic stage's three components live in their own folders. None of
|
|
those are things vLLM's loader knows how to map, so they are read directly.
|
|
|
|
Every loader here is strict. A silently partial load would produce audio that
|
|
sounds plausible but is wrong, which is far worse than a startup failure.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from torch import nn
|
|
from vllm.config import VllmConfig
|
|
from vllm.logger import init_logger
|
|
|
|
from .constants import AUDIO_VOCAB_SIZE, NUM_CODEBOOKS
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_AUDIO_EMBEDDING_KEY = "audio_embeddings.weight"
|
|
_COMPONENT_WEIGHT_NAME = "diffusion_pytorch_model.safetensors"
|
|
|
|
# Component folders, relative to the repository root.
|
|
DEPTH_DECODER_DIR = "rvq_depth_decoder"
|
|
CONDITION_ENCODER_DIR = "condition_encoder"
|
|
TRANSFORMER_DIR = "transformer"
|
|
VOCODER_DIR = "vocoder"
|
|
|
|
# Folder names that identify a repository root. This is deliberately a
|
|
# name-only probe and NOT the set a stage must be able to load: see
|
|
# ``_REQUIRED_COMPONENTS``.
|
|
_ROOT_MARKERS = (DEPTH_DECODER_DIR, TRANSFORMER_DIR, VOCODER_DIR)
|
|
|
|
# Every component a stage actually reads. ``condition_encoder/`` is loaded by
|
|
# the acoustic stage, so a snapshot missing it is incomplete even though the
|
|
# root markers above are all present.
|
|
_REQUIRED_COMPONENTS = (DEPTH_DECODER_DIR, CONDITION_ENCODER_DIR, TRANSFORMER_DIR, VOCODER_DIR)
|
|
|
|
# The subset a stage reading the repo ROOT needs on disk. The AR backbone under
|
|
# ``language_model/`` is pulled by vLLM's own loader, and the repo's top-level
|
|
# ``.pth`` files and ``qwen_7B/`` folder belong to the reference implementation,
|
|
# so neither is downloaded here.
|
|
_COMPONENT_PATTERNS = (
|
|
"*.json",
|
|
f"{DEPTH_DECODER_DIR}/*",
|
|
f"{CONDITION_ENCODER_DIR}/*",
|
|
f"{TRANSFORMER_DIR}/*",
|
|
f"{VOCODER_DIR}/*",
|
|
)
|
|
|
|
|
|
def _component_is_complete(folder: Path) -> bool:
|
|
"""True when *folder* holds a weight set ``load_component_state`` can read.
|
|
|
|
A single shard satisfies a bare glob, so an interrupted multi-shard
|
|
download would otherwise pass as complete. Hub shards are named
|
|
``<stem>-<index>-of-<total>.safetensors``, which states how many pieces
|
|
to expect, so that run can be checked for gaps.
|
|
|
|
Only that naming is provably incomplete. ``load_component_state`` loads
|
|
whatever the same glob returns, so any other shard-named file is accepted
|
|
here: judging it incomplete would force a needless re-download and fail
|
|
closed offline on weights that are present and readable.
|
|
"""
|
|
stem = Path(_COMPONENT_WEIGHT_NAME).stem
|
|
if (folder / _COMPONENT_WEIGHT_NAME).is_file():
|
|
return True
|
|
|
|
shards = list(folder.glob(f"{stem}-*.safetensors"))
|
|
if not shards:
|
|
return False
|
|
|
|
totals: set[int] = set()
|
|
seen: set[int] = set()
|
|
for shard in shards:
|
|
marker = shard.name[len(stem) + 1 : -len(".safetensors")]
|
|
index_text, separator, total_text = marker.partition("-of-")
|
|
if not separator or not index_text.isdigit() or not total_text.isdigit():
|
|
continue
|
|
totals.add(int(total_text))
|
|
seen.add(int(index_text))
|
|
if not totals:
|
|
# No piece announces a total, so there is nothing to be short of.
|
|
return True
|
|
if len(totals) != 1:
|
|
return False
|
|
return seen == set(range(1, totals.pop() + 1))
|
|
|
|
|
|
def _snapshot_has_components(root: Path) -> bool:
|
|
"""True when every required component folder holds its complete weights."""
|
|
return all(_component_is_complete(root / component) for component in _REQUIRED_COMPONENTS)
|
|
|
|
|
|
def _download_component_snapshot(
|
|
model: str,
|
|
*,
|
|
revision: str | None = None,
|
|
download_dir: str | None = None,
|
|
) -> Path:
|
|
"""Resolve an HF repo id to a snapshot holding the component folders.
|
|
|
|
Cache-first, so a warm cache needs no Hub round trip. Depending on the
|
|
huggingface_hub version, ``local_files_only=True`` either raises on a
|
|
partial cache or returns the partial snapshot root unchanged; both
|
|
variants must fall through to the online call, or a cold cache fails
|
|
startup without ever consulting the Hub.
|
|
|
|
``revision`` and ``download_dir`` must be threaded through: resolving the
|
|
components against the default branch and the default cache while the AR
|
|
backbone was pinned elsewhere would silently mix weights from two commits.
|
|
"""
|
|
from vllm.transformers_utils.repo_utils import hf_api
|
|
|
|
kwargs: dict[str, object] = {"allow_patterns": list(_COMPONENT_PATTERNS)}
|
|
if revision is not None:
|
|
kwargs["revision"] = revision
|
|
if download_dir is not None:
|
|
kwargs["cache_dir"] = download_dir
|
|
|
|
try:
|
|
cached = Path(hf_api().snapshot_download(model, local_files_only=True, **kwargs))
|
|
except Exception:
|
|
cached = None
|
|
if cached is not None and _snapshot_has_components(cached):
|
|
return cached
|
|
return Path(hf_api().snapshot_download(model, **kwargs))
|
|
|
|
|
|
def _repo_id_and_revision_from_cache_path(path: Path) -> tuple[str | None, str | None]:
|
|
"""Recover the Hub repo id and commit when *path* is inside an HF cache.
|
|
|
|
Stage init hands the talker the snapshot's ``language_model`` folder. On a
|
|
cold cache the sibling component folders were never downloaded, so the
|
|
marker walk fails even though the missing pieces are one snapshot_download
|
|
away. The cache layout is ``models--{org}--{name}/snapshots/{revision}/...``,
|
|
so the exact commit is recoverable too and must be reused: re-resolving the
|
|
repo id alone would fetch the default branch instead of the snapshot the
|
|
caller is already sitting in.
|
|
"""
|
|
parts = path.parts
|
|
for position, part in enumerate(parts):
|
|
if not part.startswith("models--"):
|
|
continue
|
|
repo_id = part.removeprefix("models--").replace("--", "/")
|
|
revision = None
|
|
if position + 2 < len(parts) and parts[position + 1] == "snapshots":
|
|
revision = parts[position + 2]
|
|
return repo_id, revision
|
|
return None, None
|
|
|
|
|
|
def resolve_repo_root(
|
|
model_path: str | os.PathLike[str],
|
|
*,
|
|
revision: str | None = None,
|
|
download_dir: str | None = None,
|
|
) -> Path:
|
|
"""Find the multi-component repository root from a stage's model path.
|
|
|
|
A stage that declares ``model_subdir`` sees the subfolder as its model
|
|
path, so the root is usually the parent. Both are checked, and the search
|
|
walks up a couple of levels for snapshot layouts.
|
|
|
|
The acoustic stage declares no ``model_subdir`` because it reads the root
|
|
itself, so on a hub-id deployment its model path is the repo id rather
|
|
than a directory. Resolve that to a local snapshot instead of treating it
|
|
as a relative path.
|
|
|
|
A cold cache is recoverable too: stage init pre-downloads only
|
|
``language_model/`` and ``tokenizer/``, so the talker's model path exists
|
|
while the component folders do not. Fetch them here instead of failing.
|
|
|
|
A candidate is accepted only when its components are actually loadable.
|
|
Matching on folder names alone would accept a snapshot left weightless by
|
|
an interrupted download and defer the failure to ``load_component_state``,
|
|
past the point where the repair below could still run.
|
|
|
|
Raises:
|
|
FileNotFoundError: If no ancestor looks like the repository root.
|
|
"""
|
|
path = Path(model_path).expanduser().resolve()
|
|
candidates = [path, *path.parents[:3]]
|
|
marker_roots = [
|
|
candidate for candidate in candidates if all((candidate / marker).is_dir() for marker in _ROOT_MARKERS)
|
|
]
|
|
for candidate in marker_roots:
|
|
if _snapshot_has_components(candidate):
|
|
return candidate
|
|
|
|
# Either no ancestor carries the markers, or one does but is incomplete.
|
|
# Three recoverable cases: the path is a Hub repo id (the acoustic stage
|
|
# passes the id through), or it is a cache-snapshot subfolder whose sibling
|
|
# component folders were never downloaded, or a cached snapshot lost part
|
|
# of its weights. ``_download_component_snapshot`` is cache-first, so a
|
|
# warm cache stays offline either way.
|
|
cached_repo_id, cached_revision = _repo_id_and_revision_from_cache_path(path)
|
|
reference = str(model_path) if not path.exists() else cached_repo_id
|
|
if reference is not None:
|
|
try:
|
|
resolved = _download_component_snapshot(
|
|
reference,
|
|
revision=revision or cached_revision,
|
|
download_dir=download_dir,
|
|
)
|
|
except Exception as exc:
|
|
raise FileNotFoundError(
|
|
f"MiniMax Music 3 could not resolve {model_path!r} to a local repository root; "
|
|
f"expected folders {_REQUIRED_COMPONENTS}"
|
|
) from exc
|
|
if _snapshot_has_components(resolved):
|
|
return resolved
|
|
|
|
if marker_roots:
|
|
raise FileNotFoundError(
|
|
f"MiniMax Music 3 found a repository root at {marker_roots[0]} but its components are incomplete; "
|
|
f"expected loadable weights under {_REQUIRED_COMPONENTS}"
|
|
)
|
|
raise FileNotFoundError(
|
|
f"MiniMax Music 3 could not locate the repository root from {path}; "
|
|
f"expected sibling folders {_REQUIRED_COMPONENTS}"
|
|
)
|
|
|
|
|
|
def load_component_state(root: Path, component: str, *, device: str = "cpu") -> dict[str, torch.Tensor]:
|
|
"""Read one component's safetensors, sharded or not.
|
|
|
|
Raises:
|
|
FileNotFoundError: If the component has no weight files.
|
|
"""
|
|
from safetensors.torch import load_file
|
|
|
|
folder = root / component
|
|
single = folder / _COMPONENT_WEIGHT_NAME
|
|
if single.exists():
|
|
return load_file(str(single), device=device)
|
|
|
|
shards = sorted(folder.glob(f"{Path(_COMPONENT_WEIGHT_NAME).stem}-*.safetensors"))
|
|
if not shards:
|
|
raise FileNotFoundError(f"MiniMax Music 3 component {component!r} has no weights under {folder}")
|
|
state: dict[str, torch.Tensor] = {}
|
|
for shard in shards:
|
|
state.update(load_file(str(shard), device=device))
|
|
return state
|
|
|
|
|
|
def load_depth_decoder_weights(
|
|
vllm_config: VllmConfig,
|
|
*,
|
|
audio_embeddings: nn.Embedding,
|
|
rvq_decoder: nn.Module,
|
|
) -> set[str]:
|
|
"""Load the audio embedding table and the RVQ depth decoder in place.
|
|
|
|
Args:
|
|
vllm_config: The stage's config, used to locate the checkpoint.
|
|
audio_embeddings: The ``c1..c7`` embedding table to populate.
|
|
rvq_decoder: The depth decoder to populate.
|
|
|
|
Returns:
|
|
The model-side parameter names that were loaded.
|
|
|
|
Raises:
|
|
ValueError: If the audio embedding is absent or the wrong shape.
|
|
"""
|
|
root = resolve_repo_root(
|
|
vllm_config.model_config.model,
|
|
revision=vllm_config.model_config.revision,
|
|
download_dir=vllm_config.load_config.download_dir,
|
|
)
|
|
state = load_component_state(root, DEPTH_DECODER_DIR)
|
|
|
|
weight = state.pop(_AUDIO_EMBEDDING_KEY, None)
|
|
expected = (AUDIO_VOCAB_SIZE * (NUM_CODEBOOKS - 1), audio_embeddings.embedding_dim)
|
|
if weight is None or tuple(weight.shape) != expected:
|
|
raise ValueError(
|
|
f"MiniMax Music 3 audio embedding mismatch: expected {expected}, got "
|
|
f"{None if weight is None else tuple(weight.shape)}"
|
|
)
|
|
target = audio_embeddings.weight
|
|
target.data.copy_(weight.to(device=target.device, dtype=target.dtype))
|
|
|
|
rvq_decoder.load_checkpoint_state(state)
|
|
logger.info(
|
|
"MiniMax Music 3 loaded audio embedding and depth decoder from %s",
|
|
root / DEPTH_DECODER_DIR,
|
|
)
|
|
loaded = {"audio_embeddings.weight"}
|
|
loaded.update(f"rvq_decoder.{name}" for name, _ in rvq_decoder.named_parameters())
|
|
return loaded
|
|
|
|
|
|
__all__ = [
|
|
"CONDITION_ENCODER_DIR",
|
|
"DEPTH_DECODER_DIR",
|
|
"TRANSFORMER_DIR",
|
|
"VOCODER_DIR",
|
|
"load_component_state",
|
|
"load_depth_decoder_weights",
|
|
"resolve_repo_root",
|
|
]
|