跳转至

llm.compat — Compatibility Layer

HuggingFace compatibility utilities for loading and publishing models. The compat group (uv sync --extra compat) provides optional dependencies (huggingface_hub, pillow, safetensors).

Overview

Module Purpose
hf_loader Load HuggingFace checkpoints
hf_publisher Publish models to HuggingFace Hub
weight_mapping Map weight names between formats

HF Loader

hf_loader

HuggingFace Model Loader.

Provides from_pretrained functionality for loading HuggingFace models into our DecoderModel format.

from_pretrained

from_pretrained(model_path, device='auto', dtype=None, trust_remote_code=False)

Load a pretrained model from HuggingFace format.

Supports loading from: - Local directory with config.json and model weights - HuggingFace Hub model ID (requires huggingface_hub)

When loading from the Hub, only *.json and *.safetensors are downloaded. *.bin files are intentionally skipped because they are pickled and can execute arbitrary code on load. Local *.bin files are still accepted (you opted into that file by putting it on disk).

参数:

名称 类型 描述 默认
model_path str | Path

Local path or HuggingFace model ID.

必需
device str | device

Device to load model on ("auto", "cuda", "cpu").

'auto'
dtype dtype | None

Data type for model weights.

None
trust_remote_code bool

Whether to trust remote code (for HF Hub).

False

返回:

类型 描述
DecoderModel

Loaded DecoderModel.

源代码位于: src/llm/compat/hf_loader.py
def from_pretrained(
    model_path: str | Path,
    device: str | torch.device = "auto",
    dtype: torch.dtype | None = None,
    trust_remote_code: bool = False,
) -> DecoderModel:
    """
    Load a pretrained model from HuggingFace format.

    Supports loading from:
    - Local directory with config.json and model weights
    - HuggingFace Hub model ID (requires huggingface_hub)

    When loading from the Hub, only ``*.json`` and ``*.safetensors`` are
    downloaded. ``*.bin`` files are intentionally skipped because they are
    pickled and can execute arbitrary code on load. Local ``*.bin`` files
    are still accepted (you opted into that file by putting it on disk).

    Args:
        model_path: Local path or HuggingFace model ID.
        device: Device to load model on ("auto", "cuda", "cpu").
        dtype: Data type for model weights.
        trust_remote_code: Whether to trust remote code (for HF Hub).

    Returns:
        Loaded DecoderModel.
    """
    model_path = Path(model_path)

    # Determine if local or Hub
    if model_path.exists():
        return _load_from_local(model_path, device, dtype)
    else:
        return _load_from_hub(str(model_path), device, dtype, trust_remote_code)

list_supported_architectures

list_supported_architectures()

List supported model architectures.

mistral is included (dense Mistral). mixtral is deliberately NOT: it is block-sparse MoE, which :func:from_pretrained rejects with a clear error rather than silently dropping every expert/router tensor (RIL ISS-144).

源代码位于: src/llm/compat/hf_loader.py
def list_supported_architectures() -> list[str]:
    """List supported model architectures.

    ``mistral`` is included (dense Mistral). ``mixtral`` is deliberately NOT:
    it is block-sparse MoE, which :func:`from_pretrained` rejects with a clear
    error rather than silently dropping every expert/router tensor (RIL
    ISS-144).
    """
    return ["llama", "llama2", "llama3", "mistral", "qwen", "qwen2"]

HF Publisher

hf_publisher

HuggingFace publish helpers — reverse of :mod:llm.compat.hf_loader.

Provides :func:save_pretrained and :func:push_to_hub so models trained with this project can be shared on HuggingFace Hub in a format that the existing :func:llm.compat.hf_loader.from_pretrained can roundtrip-load (i.e. produces config.json + model.safetensors in Llama-style naming).

Both helpers are soft-dependency-friendly:

  • safetensors is required for save — install via pip install 'llm[compat]'.
  • huggingface_hub is required for push — same install command.

Each raises a clear ImportError with the install hint when the corresponding dependency is missing, mirroring the convention used elsewhere in the project (flash_attn, huggingface_hub in the loader, etc.).

The reverse weight mapping lives in :mod:llm.compat.weight_mapping (see :func:llm.compat.weight_mapping.convert_our_weights).

save_pretrained

save_pretrained(model, save_directory)

Save a :class:DecoderModel in HuggingFace-compatible format.

Writes:

  • config.json — Llama-style config; loadable by LlamaConfig.
  • model.safetensors — state dict in Llama-style naming. Roundtrip- loadable by :func:llm.compat.hf_loader.from_pretrained.

参数:

名称 类型 描述 默认
model DecoderModel

A trained :class:DecoderModel instance.

必需
save_directory str | Path

Local directory to write into. Created if it doesn't exist.

必需

返回:

类型 描述
Path

The resolved save_directory path.

引发:

类型 描述
ImportError

If safetensors is not installed.

源代码位于: src/llm/compat/hf_publisher.py
def save_pretrained(model: DecoderModel, save_directory: str | Path) -> Path:
    """Save a :class:`DecoderModel` in HuggingFace-compatible format.

    Writes:

    * ``config.json`` — Llama-style config; loadable by ``LlamaConfig``.
    * ``model.safetensors`` — state dict in Llama-style naming. Roundtrip-
      loadable by :func:`llm.compat.hf_loader.from_pretrained`.

    Args:
        model: A trained :class:`DecoderModel` instance.
        save_directory: Local directory to write into. Created if it
            doesn't exist.

    Returns:
        The resolved ``save_directory`` path.

    Raises:
        ImportError: If ``safetensors`` is not installed.
    """
    if not SAFETENSORS_AVAILABLE:
        raise ImportError(
            "save_pretrained requires the 'safetensors' package. Install with `pip install 'llm[compat]'`."
        )

    from safetensors.torch import save_file

    save_directory = Path(save_directory)
    save_directory.mkdir(parents=True, exist_ok=True)

    # 1. Write config.json (Llama-style).
    hf_config = _build_hf_config(model)
    config_path = save_directory / "config.json"
    with config_path.open("w", encoding="utf-8") as f:
        json.dump(hf_config, f, indent=2)
    logger.info(f"Wrote HF config to {config_path}")

    # 2. Convert + write state_dict.
    num_layers = len(model.transformer_blocks)
    attn0 = model.transformer_blocks[0].self_attn
    if not isinstance(attn0, _SizedAttention):
        raise TypeError(f"attention backend {type(attn0).__name__} must expose num_heads/num_kv_heads/head_dim")
    converted = convert_our_weights(
        model.state_dict(),
        architecture="llama",
        num_layers=num_layers,
        num_heads=attn0.num_heads,
        num_kv_heads=attn0.num_kv_heads,
        head_dim=attn0.head_dim,
    )

    # Clone + contiguous for safetensors (it rejects views).
    contiguous = {k: v.detach().contiguous().clone() for k, v in converted.items()}
    weights_path = save_directory / "model.safetensors"
    save_file(contiguous, str(weights_path))
    logger.info(f"Wrote {len(contiguous)} tensors to {weights_path}")

    return save_directory

push_to_hub

push_to_hub(model, repo_id, *, token=None, private=False, commit_message='Upload model', exist_ok=True, save_directory=None)

Save the model locally and push to a HuggingFace Hub repo.

Calls :func:save_pretrained to a staging directory (or to save_directory if provided) and uploads via huggingface_hub.upload_folder. The repo is created on first push unless exist_ok=False.

参数:

名称 类型 描述 默认
model DecoderModel

A trained :class:DecoderModel.

必需
repo_id str

HF Hub repo ID (e.g. "alice/my-llm").

必需
token str | None

HF auth token. Falls back to HF_TOKEN env var if None. Use huggingface-cli login to persist.

None
private bool

Whether to create the repo as private.

False
commit_message str

Git commit message on the Hub side.

'Upload model'
exist_ok bool

Don't raise if the repo already exists.

True
save_directory str | Path | None

Optional staging dir; defaults to a temporary directory under the system temp path.

None

返回:

类型 描述
str

The HF Hub URL of the pushed repo (e.g.

str

"https://huggingface.co/alice/my-llm").

引发:

类型 描述
ImportError

If huggingface_hub (or safetensors) is not installed.

源代码位于: src/llm/compat/hf_publisher.py
def push_to_hub(
    model: DecoderModel,
    repo_id: str,
    *,
    token: str | None = None,
    private: bool = False,
    commit_message: str = "Upload model",
    exist_ok: bool = True,
    save_directory: str | Path | None = None,
) -> str:
    """Save the model locally and push to a HuggingFace Hub repo.

    Calls :func:`save_pretrained` to a staging directory (or to
    ``save_directory`` if provided) and uploads via
    ``huggingface_hub.upload_folder``. The repo is created on first
    push unless ``exist_ok=False``.

    Args:
        model: A trained :class:`DecoderModel`.
        repo_id: HF Hub repo ID (e.g. ``"alice/my-llm"``).
        token: HF auth token. Falls back to ``HF_TOKEN`` env var if
            ``None``. Use ``huggingface-cli login`` to persist.
        private: Whether to create the repo as private.
        commit_message: Git commit message on the Hub side.
        exist_ok: Don't raise if the repo already exists.
        save_directory: Optional staging dir; defaults to a temporary
            directory under the system temp path.

    Returns:
        The HF Hub URL of the pushed repo (e.g.
        ``"https://huggingface.co/alice/my-llm"``).

    Raises:
        ImportError: If ``huggingface_hub`` (or ``safetensors``) is
            not installed.
    """
    if not HF_HUB_AVAILABLE:
        raise ImportError(
            "push_to_hub requires the 'huggingface_hub' package. Install with `pip install 'llm[compat]'`."
        )

    import tempfile

    from huggingface_hub import HfApi

    if save_directory is None:
        save_directory = Path(tempfile.mkdtemp(prefix="llm-push-"))
    save_directory = Path(save_directory)

    save_pretrained(model, save_directory)

    api = HfApi(token=token)
    api.create_repo(repo_id=repo_id, private=private, exist_ok=exist_ok)
    api.upload_folder(
        folder_path=str(save_directory),
        repo_id=repo_id,
        commit_message=commit_message,
    )

    url = f"https://huggingface.co/{repo_id}"
    logger.info(f"Pushed model to {url}")
    return url

Weight Mapping

weight_mapping

Weight Mapping for HuggingFace Model Conversion.

Provides mappings from HuggingFace weight names to this project's naming convention. Supports Llama, Mistral, and Qwen architectures.

detect_architecture

detect_architecture(config)

Detect model architecture from HuggingFace config.

参数:

名称 类型 描述 默认
config dict[str, Any]

HuggingFace model config dict.

必需

返回:

类型 描述
str

Architecture name (llama, mistral, qwen, qwen2, mixtral) or

str

"unknown" for an unsupported model_type.

源代码位于: src/llm/compat/weight_mapping.py
def detect_architecture(config: dict[str, Any]) -> str:
    """
    Detect model architecture from HuggingFace config.

    Args:
        config: HuggingFace model config dict.

    Returns:
        Architecture name (llama, mistral, qwen, qwen2, mixtral) or
        ``"unknown"`` for an unsupported ``model_type``.
    """
    model_type = config.get("model_type", "").lower()

    if "llama" in model_type:
        return "llama"
    elif "mixtral" in model_type:
        # Mixtral is MoE (sparse experts + router). Our mapping is dense-only
        # and from_pretrained would build a DENSE model, silently dropping
        # every experts.*/gate tensor (RIL ISS-144). Keep it distinguishable
        # so the loader can REJECT it with a clear error instead of shipping
        # a model whose routers/experts are all at random init.
        return "mixtral"
    elif "mistral" in model_type:
        return "mistral"
    elif model_type in {"qwen2_moe", "qwen2moe", "qwen3", "qwen3_moe", "qwen3moe"}:
        # Qwen2MoE / Qwen3(MoE) are NOT supported. The substring rules below
        # would map ``qwen2_moe`` onto the DENSE Qwen2 rules (dropping every
        # expert/router tensor) and ``qwen3`` onto the Qwen1 GPT-style rules
        # (dropping most weights) — both then run from RANDOM init with
        # warnings only (the ISS-144 / round-71 anti-garbage-load philosophy).
        # Route them to "unknown" so the loader refuses loudly.
        return "unknown"
    elif "qwen2" in model_type:
        return "qwen2"
    elif "qwen" in model_type:
        return "qwen"
    else:
        # Unknown model_type (gpt2, gemma, baichuan, ...). Previously this
        # defaulted to the llama mapping and from_pretrained loaded with
        # strict=False — every unmapped weight stayed at random init and the
        # model generated garbage with only warning logs (round-71 compat
        # fix). Return a distinguishable token so the loader REFUSES instead.
        return "unknown"

get_weight_mapping

get_weight_mapping(architecture)

Get weight name mapping for an architecture.

参数:

名称 类型 描述 默认
architecture str

Architecture name.

必需

返回:

类型 描述
dict[str, str]

Dictionary mapping HF names to our names.

源代码位于: src/llm/compat/weight_mapping.py
def get_weight_mapping(architecture: str) -> dict[str, str]:
    """
    Get weight name mapping for an architecture.

    Args:
        architecture: Architecture name.

    Returns:
        Dictionary mapping HF names to our names.
    """
    return ARCHITECTURE_MAPPINGS.get(architecture, LLAMA_MAPPING)

expand_layer_mapping

expand_layer_mapping(mapping, num_layers)

Expand layer-indexed mappings for all layers.

参数:

名称 类型 描述 默认
mapping dict[str, str]

Base mapping with {layer} placeholders.

必需
num_layers int

Number of transformer layers.

必需

返回:

类型 描述
dict[str, str]

Expanded mapping with concrete layer indices.

源代码位于: src/llm/compat/weight_mapping.py
def expand_layer_mapping(mapping: dict[str, str], num_layers: int) -> dict[str, str]:
    """
    Expand layer-indexed mappings for all layers.

    Args:
        mapping: Base mapping with {layer} placeholders.
        num_layers: Number of transformer layers.

    Returns:
        Expanded mapping with concrete layer indices.
    """
    expanded = {}

    for hf_pattern, our_pattern in mapping.items():
        if "{layer}" in hf_pattern:
            for layer_idx in range(num_layers):
                hf_name = hf_pattern.format(layer=layer_idx)
                our_name = our_pattern.format(layer=layer_idx)
                expanded[hf_name] = our_name
        else:
            expanded[hf_pattern] = our_pattern

    return expanded

convert_hf_weights

convert_hf_weights(hf_state_dict, architecture, num_layers)

Convert HuggingFace state dict to our naming convention.

参数:

名称 类型 描述 默认
hf_state_dict dict[str, Any]

HuggingFace model state dict.

必需
architecture str

Model architecture.

必需
num_layers int

Number of transformer layers.

必需

返回:

类型 描述
dict[str, Any]

Converted state dict with our naming.

源代码位于: src/llm/compat/weight_mapping.py
def convert_hf_weights(
    hf_state_dict: dict[str, Any],
    architecture: str,
    num_layers: int,
) -> dict[str, Any]:
    """
    Convert HuggingFace state dict to our naming convention.

    Args:
        hf_state_dict: HuggingFace model state dict.
        architecture: Model architecture.
        num_layers: Number of transformer layers.

    Returns:
        Converted state dict with our naming.
    """
    mapping = get_weight_mapping(architecture)
    expanded_mapping = expand_layer_mapping(mapping, num_layers)

    converted = {}
    unmapped = []

    for hf_name, tensor in hf_state_dict.items():
        if hf_name in expanded_mapping:
            our_name = expanded_mapping[hf_name]
            converted[our_name] = tensor
        else:
            # Try partial match for bias terms etc.
            matched = False
            for hf_pattern, our_pattern in expanded_mapping.items():
                if hf_name.replace(".bias", ".weight") == hf_pattern:
                    our_name = our_pattern.replace(".weight", ".bias")
                    converted[our_name] = tensor
                    matched = True
                    break

            if not matched:
                unmapped.append(hf_name)

    if unmapped:
        import logging

        logger = logging.getLogger(__name__)
        logger.warning(f"Unmapped weights: {unmapped[:10]}{'...' if len(unmapped) > 10 else ''}")

    return converted

convert_gguf_weights

convert_gguf_weights(gguf_state_dict, num_layers)

Translate llama.cpp GGUF tensor names into our naming convention.

参数:

名称 类型 描述 默认
gguf_state_dict dict[str, Any]

GGUF tensors keyed by llama.cpp names (token_embd, blk.N.attn_q, blk.N.ffn_gate, ...).

必需
num_layers int

Number of transformer layers (to expand the {layer} placeholders).

必需

返回:

类型 描述
dict[str, Any]

(converted, unmapped) — converted tensors keyed by our naming

list[str]

(q/k/v still split as .*_proj, ready for

tuple[dict[str, Any], list[str]]

func:convert_hf_to_combined_qkv), and the list of GGUF tensor names

tuple[dict[str, Any], list[str]]

that had no mapping. A non-empty unmapped list means the file is

tuple[dict[str, Any], list[str]]

not a pure dense Llama-style GGUF (or carries extra tensors); the

tuple[dict[str, Any], list[str]]

caller should refuse rather than silently drop them (RIL ISS-220

tuple[dict[str, Any], list[str]]

philosophy).

源代码位于: src/llm/compat/weight_mapping.py
def convert_gguf_weights(
    gguf_state_dict: dict[str, Any],
    num_layers: int,
) -> tuple[dict[str, Any], list[str]]:
    """Translate llama.cpp GGUF tensor names into our naming convention.

    Args:
        gguf_state_dict: GGUF tensors keyed by llama.cpp names (``token_embd``,
            ``blk.N.attn_q``, ``blk.N.ffn_gate``, ...).
        num_layers: Number of transformer layers (to expand the ``{layer}``
            placeholders).

    Returns:
        ``(converted, unmapped)`` — converted tensors keyed by our naming
        (q/k/v still *split* as ``.*_proj``, ready for
        :func:`convert_hf_to_combined_qkv`), and the list of GGUF tensor names
        that had no mapping. A non-empty ``unmapped`` list means the file is
        not a pure dense Llama-style GGUF (or carries extra tensors); the
        caller should refuse rather than silently drop them (RIL ISS-220
        philosophy).
    """
    mapping = expand_layer_mapping(GGUF_MAPPING, num_layers)
    converted: dict[str, Any] = {}
    unmapped: list[str] = []
    for gguf_name, tensor in gguf_state_dict.items():
        our_name = mapping.get(gguf_name)
        if our_name is not None:
            converted[our_name] = tensor
        else:
            unmapped.append(gguf_name)
    return converted, unmapped

convert_our_weights

convert_our_weights(our_state_dict, architecture, num_layers, *, num_heads=None, num_kv_heads=None, head_dim=None)

Convert our naming convention to HuggingFace state dict.

Inverse of :func:convert_hf_weights for the supported weight names. Used by save_pretrained to publish models to HuggingFace in a format the existing from_pretrained can roundtrip-load.

Splits our combined qkv_proj projection into HF's separate q_proj / k_proj / v_proj weights so the published artifact is loadable by both our from_pretrained (which uses the reverse concat) and HF's transformers library.

参数:

名称 类型 描述 默认
our_state_dict dict[str, Any]

Our model state dict (e.g. model.state_dict()).

必需
architecture str

Target HF architecture (must match the model).

必需
num_layers int

Number of transformer layers in the model.

必需
num_heads int | None

Total attention heads. Required when the model has a combined qkv_proj so we can split Q vs. K/V.

None
num_kv_heads int | None

Number of KV heads (for GQA/MQA). Defaults to num_heads (standard MHA).

None
head_dim int | None

Per-head dimension. Defaults to hidden_size // num_heads.

None

返回:

类型 描述
dict[str, Any]

Converted state dict with HuggingFace naming.

源代码位于: src/llm/compat/weight_mapping.py
def convert_our_weights(
    our_state_dict: dict[str, Any],
    architecture: str,
    num_layers: int,
    *,
    num_heads: int | None = None,
    num_kv_heads: int | None = None,
    head_dim: int | None = None,
) -> dict[str, Any]:
    """
    Convert our naming convention to HuggingFace state dict.

    Inverse of :func:`convert_hf_weights` for the supported weight
    names. Used by ``save_pretrained`` to publish models to
    HuggingFace in a format the existing ``from_pretrained`` can
    roundtrip-load.

    Splits our **combined** ``qkv_proj`` projection into HF's separate
    ``q_proj`` / ``k_proj`` / ``v_proj`` weights so the published
    artifact is loadable by both our ``from_pretrained`` (which uses
    the reverse concat) and HF's transformers library.

    Args:
        our_state_dict: Our model state dict (e.g. ``model.state_dict()``).
        architecture: Target HF architecture (must match the model).
        num_layers: Number of transformer layers in the model.
        num_heads: Total attention heads. Required when the model has
            a combined ``qkv_proj`` so we can split Q vs. K/V.
        num_kv_heads: Number of KV heads (for GQA/MQA). Defaults to
            ``num_heads`` (standard MHA).
        head_dim: Per-head dimension. Defaults to ``hidden_size // num_heads``.

    Returns:
        Converted state dict with HuggingFace naming.
    """
    mapping = get_weight_mapping(architecture)
    expanded_mapping = expand_layer_mapping(mapping, num_layers)

    # Build the reverse map: our_name -> hf_name. The forward map is
    # hf_name -> our_name, so we invert it.
    reverse_mapping = {our_name: hf_name for hf_name, our_name in expanded_mapping.items()}

    converted = {}
    unmapped = []

    # First pass: split combined qkv_proj into q_proj/k_proj/v_proj if
    # the model uses a combined projection (the current MHA impl does).
    qkv_keys = [
        k for k in our_state_dict if k.endswith(".self_attn.qkv_proj.weight") or k.endswith(".self_attn.qkv_proj.bias")
    ]
    if qkv_keys and num_heads is not None:
        n_q = num_heads
        n_kv = num_kv_heads if num_kv_heads is not None else num_heads
        # head_dim must be supplied; defaulting here would silently
        # mis-split the projection.
        if head_dim is None:
            raise ValueError(
                "head_dim is required when splitting combined qkv_proj. Pass it explicitly from the model's MHA block."
            )
        q_size = n_q * head_dim
        kv_size = n_kv * head_dim

        for qkv_key in qkv_keys:
            tensor = our_state_dict[qkv_key]
            # Strip the trailing ".qkv_proj" suffix so we get just the
            # layer prefix, e.g. "transformer_blocks.0.self_attn".
            our_prefix = qkv_key.rsplit(".", 1)[0].rsplit(".", 1)[0]
            # Translate to HF naming: replace ``transformer_blocks.``
            # with ``model.layers.`` — the rest stays the same.
            hf_prefix = our_prefix.replace("transformer_blocks.", "model.layers.") + "."

            if "weight" in qkv_key:
                # Linear weight: shape (out_features, in_features)
                q_w, k_w, v_w = tensor.split([q_size, kv_size, kv_size], dim=0)
                converted[hf_prefix + "q_proj.weight"] = q_w.contiguous()
                converted[hf_prefix + "k_proj.weight"] = k_w.contiguous()
                converted[hf_prefix + "v_proj.weight"] = v_w.contiguous()
            else:
                q_b, k_b, v_b = tensor.split([q_size, kv_size, kv_size], dim=0)
                converted[hf_prefix + "q_proj.bias"] = q_b.contiguous()
                converted[hf_prefix + "k_proj.bias"] = k_b.contiguous()
                converted[hf_prefix + "v_proj.bias"] = v_b.contiguous()

    # Second pass: rename everything else via the reverse mapping.
    for our_name, tensor in our_state_dict.items():
        if our_name in qkv_keys:
            continue  # already handled above
        if our_name in reverse_mapping:
            converted[reverse_mapping[our_name]] = tensor
        else:
            # Try partial match for bias terms (mirrors the HF -> ours path).
            matched = False
            for our_pattern, hf_pattern in reverse_mapping.items():
                if our_name.replace(".bias", ".weight") == our_pattern:
                    converted[hf_pattern.replace(".weight", ".bias")] = tensor
                    matched = True
                    break

            if not matched:
                unmapped.append(our_name)

    if unmapped:
        import logging

        logger = logging.getLogger(__name__)
        logger.warning(f"Unmapped weights (ours -> HF): {unmapped[:10]}{'...' if len(unmapped) > 10 else ''}")

    return converted

convert_hf_to_combined_qkv

convert_hf_to_combined_qkv(our_state_dict, num_layers, *, num_heads=None, num_kv_heads=None, head_dim=None)

Concatenate our separate q_proj / k_proj / v_proj projections into the combined qkv_proj.

Used by :func:llm.compat.hf_loader.from_pretrained after :func:convert_hf_weights has renamed HF Llama's separate q/k/v projections to our naming. Our MHA stores Q/K/V in a single qkv_proj Linear — this helper fuses the three projections back together so load_state_dict finds the expected key.

参数:

名称 类型 描述 默认
our_state_dict dict[str, Any]

Our-renamed state dict (output of convert_hf_weights). Keys must use our naming (transformer_blocks.{layer}.self_attn.q_proj.weight etc.).

必需
num_layers int

Number of transformer layers.

必需
num_heads int | None

Total attention heads.

None
num_kv_heads int | None

Number of KV heads (for GQA/MQA). Defaults to num_heads (standard MHA).

None
head_dim int | None

Per-head dimension.

None

返回:

类型 描述
dict[str, Any]

State dict with combined qkv_proj projections.

源代码位于: src/llm/compat/weight_mapping.py
def convert_hf_to_combined_qkv(
    our_state_dict: dict[str, Any],
    num_layers: int,
    *,
    num_heads: int | None = None,
    num_kv_heads: int | None = None,  # noqa: ARG001 - reserved for GQA-aware paths; not used by simple concat
    head_dim: int | None = None,
) -> dict[str, Any]:
    """
    Concatenate our separate ``q_proj`` / ``k_proj`` / ``v_proj``
    projections into the combined ``qkv_proj``.

    Used by :func:`llm.compat.hf_loader.from_pretrained` after
    :func:`convert_hf_weights` has renamed HF Llama's separate
    q/k/v projections to our naming. Our MHA stores Q/K/V in a
    single ``qkv_proj`` Linear — this helper fuses the three
    projections back together so ``load_state_dict`` finds the
    expected key.

    Args:
        our_state_dict: Our-renamed state dict (output of
            ``convert_hf_weights``). Keys must use our naming
            (``transformer_blocks.{layer}.self_attn.q_proj.weight``
            etc.).
        num_layers: Number of transformer layers.
        num_heads: Total attention heads.
        num_kv_heads: Number of KV heads (for GQA/MQA). Defaults to
            ``num_heads`` (standard MHA).
        head_dim: Per-head dimension.

    Returns:
        State dict with combined ``qkv_proj`` projections.
    """
    if num_heads is None or head_dim is None:
        raise ValueError(
            "num_heads and head_dim are required to concatenate q/k/v projections. "
            "Pass them from the loaded model's MHA block."
        )

    out = dict(our_state_dict)
    for layer_idx in range(num_layers):
        prefix = f"transformer_blocks.{layer_idx}.self_attn."
        q_w = out.pop(f"{prefix}q_proj.weight", None)
        k_w = out.pop(f"{prefix}k_proj.weight", None)
        v_w = out.pop(f"{prefix}v_proj.weight", None)
        q_b = out.pop(f"{prefix}q_proj.bias", None)
        k_b = out.pop(f"{prefix}k_proj.bias", None)
        v_b = out.pop(f"{prefix}v_proj.bias", None)

        if q_w is not None and k_w is not None and v_w is not None:
            combined_w = torch.cat([q_w, k_w, v_w], dim=0)
            out[f"{prefix}qkv_proj.weight"] = combined_w.contiguous()
        if q_b is not None and k_b is not None and v_b is not None:
            combined_b = torch.cat([q_b, k_b, v_b], dim=0)
            out[f"{prefix}qkv_proj.bias"] = combined_b.contiguous()

    return out

get_config_mapping

get_config_mapping(hf_config)

Map HuggingFace config to our config format.

参数:

名称 类型 描述 默认
hf_config dict[str, Any]

HuggingFace config dict.

必需

返回:

类型 描述
dict[str, Any]

Our config dict.

源代码位于: src/llm/compat/weight_mapping.py
def get_config_mapping(hf_config: dict[str, Any]) -> dict[str, Any]:
    """
    Map HuggingFace config to our config format.

    Args:
        hf_config: HuggingFace config dict.

    Returns:
        Our config dict.
    """
    return {
        "vocab_size": hf_config.get("vocab_size", 32000),
        "hidden_size": hf_config.get("hidden_size", 4096),
        "num_layers": hf_config.get("num_hidden_layers", 32),
        "num_heads": hf_config.get("num_attention_heads", 32),
        "num_kv_heads": hf_config.get("num_key_value_heads"),
        "intermediate_size": hf_config.get("intermediate_size"),
        "max_seq_len": hf_config.get("max_position_embeddings", 4096),
        # Mistral's sliding-window attention: absent on Llama/Qwen externals
        # and defaulted to None (full-context) — wiring it prevents the silent
        # full-context attention past a 4096 window (RIL ISS-242).
        "window_size": hf_config.get("sliding_window"),
        # Sparse/streaming attention scheme (RIL TASK-244): our own publisher
        # persists ``attn_sparse`` (kind + params) so a sparse model roundtrips
        # with its scheme instead of silently rebuilding as dense on load.
        # External checkpoints carry no such key and default to None (dense).
        "attn_sparse": hf_config.get("attn_sparse"),
        "rms_norm_eps": hf_config.get("rms_norm_eps", 1e-5),
        "rope_theta": hf_config.get("rope_theta", 10000.0),
        # HF Llama/GPT-style configs carry the MLP activation as
        # ``hidden_act``; real Llama/Mistral use ``silu`` (SwiGLU).  The
        # loader maps this onto our ``mlp_activation`` so a published model
        # (or HF checkpoint) round-trips with the *same* MLP function rather
        # than silently defaulting to gelu.
        # Attention family so an MLA model round-trips as MLA instead of
        # silently rebuilding as MHA with every MLA tensor dropped at random
        # init (RIL ISS-169). Our own publisher persists it; external
        # Llama/Mistral checkpoints carry no such key and default to MHA.
        "attn_impl": hf_config.get("attn_impl", "mha"),
        "mlp_activation": hf_config.get("hidden_act", "silu"),
        # Whether the MLP is gated (SwiGLU). Real Llama/Mistral default to
        # True, but our own ``save_pretrained`` persists the actual
        # ``use_glu`` so a DEFAULT (non-GLU) model round-trips with its
        # fc1/fc2 MLP instead of being rebuilt as GLU with random gate
        # weights (RIL ISS-056). Absent (an external HF checkpoint) -> True.
        "use_glu": hf_config.get("use_glu", True),
        # Learned positional encoding flag so a learned-PE model round-trips
        # with its trained pos_embedding weights instead of silently falling
        # back to sinusoidal (RIL ISS-063). Default False (matching the
        # DecoderModel default).
        "pos_encoding_learned": hf_config.get("pos_encoding_learned", False),
        # Normalization implementation so an RMSNorm-trained model round-trips
        # with RMSNorm instead of silently rebuilding as LayerNorm (RIL
        # ISS-062). Our own publisher persists it; external checkpoints
        # default to layer_norm (the DecoderModel default).
        "norm_impl": hf_config.get("norm_impl", "layer_norm"),
        # Pre-LN vs post-LN block ordering. Our own publisher persists the
        # explicit flag so a post-LN model roundtrips as post-LN instead of
        # silently becoming the pre-LN default (RIL ISS-072).
        "norm_first": hf_config.get("norm_first", True),
        # RoPE: real Llama/Mistral/Qwen always use rotary position embedding
        # (their HF configs carry ``rope_theta`` and no ``use_rope`` key), so
        # an external checkpoint defaults to RoPE-on. Our own publisher
        # persists the explicit flag, keeping save->load self-consistent
        # (RIL ISS-062 — core.rope had zero callers before this wiring).
        "use_rope": hf_config.get("use_rope", True),
        # Bias flags: real Llama/Mistral checkpoints are bias-free (no
        # qkv/mlp/lm_head biases — ``attention_bias=False``). An external
        # checkpoint therefore defaults to bias-free; our own publisher
        # persists the actual flags so a biased model roundtrips with its
        # biases (RIL ISS-062).
        #
        # External checkpoints declare attention bias under HF's CANONICAL
        # ``attention_bias`` key (Qwen-style), not our repo-custom names —
        # falling back to those keys silently dropped every attention/MLP
        # bias for ``attention_bias: true`` checkpoints (RIL ISS-145). Prefer
        # our own persisted flags, then HF's canonical key for externals.
        "qkv_bias": hf_config.get("qkv_bias", hf_config.get("attention_bias", False)),
        "mlp_bias": hf_config.get("mlp_bias", hf_config.get("attention_bias", False)),
        "lm_head_bias": hf_config.get("lm_head_bias", hf_config.get("attention_bias", False)),
    }