For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

LoRAModel

LoRAModel

class max.pipelines.lora.LoRAModel(name, path, base_dtype, max_lora_rank, n_heads, n_kv_heads, head_dim, strict=True)

source

Bases: object

Manages LoRA weights and configuration for a single adapter.

Initializes a LoRAModel by loading its configuration and weights.

import json
import tempfile
from pathlib import Path

import numpy as np
from safetensors.numpy import save_file
from max.dtype import DType
from max.pipelines.lora.lora import LoRAModel

# layer with q/k/v/o projections) so the loader has something to
# read.
rank, n_heads, n_kv_heads, head_dim = 4, 8, 8, 16
hidden = n_heads * head_dim
kv_hidden = n_kv_heads * head_dim
tmp = tempfile.mkdtemp()
tensors = {}
for proj, out in (
    ("q_proj", hidden),
    ("k_proj", kv_hidden),
    ("v_proj", kv_hidden),
    ("o_proj", hidden),
):
    base = f"base_model.model.model.layers.0.self_attn.{proj}"
    tensors[f"{base}.lora_A.weight"] = np.zeros(
        (rank, hidden), dtype=np.float32
    )
    tensors[f"{base}.lora_B.weight"] = np.zeros(
        (out, rank), dtype=np.float32
    )
save_file(tensors, str(Path(tmp) / "adapter_model.safetensors"))
(Path(tmp) / "adapter_config.json").write_text(
    json.dumps(
        {
            "r": rank,
            "lora_alpha": 8,
            "bias": "none",
            "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"],
        }
    )
)

lora = LoRAModel(
    "my_adapter",
    tmp,
    DType.bfloat16,
    max_lora_rank=16,
    n_heads=n_heads,
    n_kv_heads=n_kv_heads,
    head_dim=head_dim,
)

Parameters:

• name (str) – A string identifier for this adapter.
• path (str) – Filesystem path is only supported
• base_dtype (DType) – The base model dtype.
• max_lora_rank (int) – The maximum LoRA rank supported by the system.
• n_heads (int) – Number of attention heads in the base model.
• n_kv_heads (int) – Number of key-value heads in the base model.
• head_dim (int) – Dimension of each attention head.
• strict (bool) – Whether to enforce strict validation while loading the adapter.

Raises:

ValueError – If weight files are not in the supported safetensors format, or if the keys in the weights are malformed or incomplete.

adapter_config

property adapter_config: dict[str, Any]

source

A dictionary containing metadata/configuration for the LoRA adapter.

get()

get(key)

source

Gets the WeightData from the key. If key doesn’t exist in model, then None is returned.

Parameters:

key (str) – Key of LoRA

Returns:

WeightData for the key or None if it doesn’t exist.

Return type:

WeightData | None

lora_A

property lora_A: dict[str, WeightData]

source

A dictionary mapping weight keys to LoRA A WeightData.

lora_B

property lora_B: dict[str, WeightData]

source

A dictionary mapping weight keys to LoRA B WeightData.

lora_bias

property lora_bias: dict[str, WeightData]

source

A dictionary mapping weight keys to LoRA bias WeightData.