262 lines
9.1 KiB
Python
262 lines
9.1 KiB
Python
"""Load checkpoint weights and apply fused-format conversion when needed.
|
|
|
|
Model-instance operations such as ``load_state_dict`` and ``set_normalizer``
|
|
are intentionally left to the adapter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import json
|
|
import math
|
|
from typing import Callable, Optional
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
def _noop_log(_msg: str, **_kw) -> None:
|
|
pass
|
|
|
|
|
|
def resolve_lora_scale(
|
|
checkpoint_path: str,
|
|
train_config: dict,
|
|
state_dict: dict,
|
|
log_fn: Optional[Callable] = None,
|
|
) -> float:
|
|
"""Read PEFT alpha/r from an exported config and validate it against weights."""
|
|
log_fn = log_fn or _noop_log
|
|
ranks = {
|
|
int(tensor.shape[0])
|
|
for name, tensor in state_dict.items()
|
|
if name.endswith(".lora_A.default.weight")
|
|
}
|
|
if not ranks:
|
|
return 1.0
|
|
if len(ranks) != 1:
|
|
raise ValueError(f"LoRA weights contain multiple ranks: {sorted(ranks)}")
|
|
rank = ranks.pop()
|
|
|
|
candidates = [
|
|
os.path.join(checkpoint_path, "lora_config.json"),
|
|
os.path.join(checkpoint_path, "adapter_config.json"),
|
|
(train_config.get("model") or {}).get("lora_config_path"),
|
|
]
|
|
source = next((path for path in candidates if path and os.path.isfile(path)), None)
|
|
if source is None:
|
|
log_fn(
|
|
"LoRA scale metadata is missing; using legacy alpha/r=2.0. "
|
|
"Export the training LoRA JSON as checkpoint/lora_config.json "
|
|
"before deploying a checkpoint with changed LoRA settings."
|
|
)
|
|
return 2.0
|
|
|
|
with open(source, encoding="utf-8") as stream:
|
|
config = json.load(stream)
|
|
configured_rank = int(config.get("lora_r", config.get("r", rank)))
|
|
if configured_rank != rank:
|
|
raise ValueError(
|
|
f"LoRA rank mismatch: {source} has {configured_rank}, weights have {rank}"
|
|
)
|
|
if config.get("rank_pattern") or config.get("alpha_pattern"):
|
|
raise ValueError(f"Per-layer LoRA scaling in {source} is unsupported")
|
|
alpha = config.get("lora_alpha")
|
|
if alpha is None:
|
|
raise ValueError(f"LoRA config {source} lacks lora_alpha")
|
|
scale = float(alpha) / (math.sqrt(rank) if config.get("use_rslora") else rank)
|
|
if not math.isfinite(scale) or scale <= 0:
|
|
raise ValueError(f"Invalid LoRA scale {scale} from {source}")
|
|
log_fn(f"LoRA scale={scale:g} from {source} (alpha={alpha}, rank={rank})")
|
|
return scale
|
|
|
|
|
|
def _align_checkpoint_tensor(
|
|
param: torch.Tensor,
|
|
target: torch.Tensor,
|
|
name: str,
|
|
log_fn: Callable,
|
|
) -> torch.Tensor | None:
|
|
"""Crop or pad a checkpoint tensor to match the current model parameter."""
|
|
if param.shape == target.shape:
|
|
return param
|
|
if param.ndim != target.ndim:
|
|
log_fn(
|
|
f"Skipping '{name}': ndim mismatch "
|
|
f"checkpoint={param.ndim} model={target.ndim}"
|
|
)
|
|
return None
|
|
|
|
overlap = tuple(
|
|
slice(0, min(src, dst)) for src, dst in zip(param.shape, target.shape)
|
|
)
|
|
|
|
if all(src >= dst for src, dst in zip(param.shape, target.shape)):
|
|
aligned = param[overlap].contiguous()
|
|
log_fn(
|
|
f"Cropped '{name}': checkpoint {tuple(param.shape)} "
|
|
f"-> model {tuple(aligned.shape)}"
|
|
)
|
|
return aligned
|
|
|
|
if all(src <= dst for src, dst in zip(param.shape, target.shape)):
|
|
aligned = target.detach().clone()
|
|
aligned[overlap] = param[overlap]
|
|
log_fn(
|
|
f"Padded '{name}': checkpoint {tuple(param.shape)} "
|
|
f"-> model {tuple(aligned.shape)} (tail keeps model init)"
|
|
)
|
|
return aligned
|
|
|
|
aligned = target.detach().clone()
|
|
aligned[overlap] = param[overlap]
|
|
log_fn(
|
|
f"Partially aligned '{name}': checkpoint {tuple(param.shape)} "
|
|
f"-> model {tuple(aligned.shape)} (non-overlap keeps model init)"
|
|
)
|
|
return aligned
|
|
|
|
|
|
def reshape_compatible_state_dict(
|
|
state_dict: dict,
|
|
model_sd: dict,
|
|
log_fn: Optional[Callable] = None,
|
|
lora_scale: float = 2.0,
|
|
) -> dict:
|
|
"""Align checkpoint tensors to the target model shapes via crop / pad."""
|
|
log_fn = log_fn or _noop_log
|
|
# Training checkpoints keep PEFT modules (base_layer/lora_A/lora_B),
|
|
# while inference uses the plain fused projection tensors. Materialize
|
|
# those tensors here so serving does not silently drop the LoRA update.
|
|
merged = dict(state_dict)
|
|
for name, base in list(state_dict.items()):
|
|
suffix = ".base_layer.weight"
|
|
if not name.endswith(suffix):
|
|
continue
|
|
stem = name[: -len(suffix)]
|
|
a_name = stem + ".lora_A.default.weight"
|
|
b_name = stem + ".lora_B.default.weight"
|
|
if a_name not in state_dict or b_name not in state_dict:
|
|
continue
|
|
a = state_dict[a_name].to(dtype=torch.float32)
|
|
b = state_dict[b_name].to(dtype=torch.float32)
|
|
merged[name.replace(".base_layer.weight", ".weight")] = (
|
|
base.to(dtype=torch.float32) + lora_scale * (b @ a)
|
|
).to(dtype=base.dtype)
|
|
del merged[name]
|
|
del merged[a_name]
|
|
del merged[b_name]
|
|
log_fn(f"Merged LoRA weights for {stem}")
|
|
state_dict = merged
|
|
out = {}
|
|
for name, param in state_dict.items():
|
|
# Training with PEFT wraps the backbone under ``base_model.model``;
|
|
# serving instantiates the unwrapped model. Normalize that prefix so
|
|
# fine-tuned backbone weights are applied instead of being ignored.
|
|
target_name = name
|
|
if ".base_layer." in target_name:
|
|
target_name = target_name.replace(".base_layer.", ".")
|
|
if target_name not in model_sd and target_name.startswith(
|
|
"model.base_model.model."
|
|
):
|
|
target_name = "model." + target_name[len("model.base_model.model.") :]
|
|
if target_name not in model_sd:
|
|
log_fn(f"Not used parameter: {name}")
|
|
continue
|
|
target = model_sd[target_name]
|
|
if param.shape == target.shape:
|
|
out[target_name] = param
|
|
continue
|
|
aligned = _align_checkpoint_tensor(param, target, target_name, log_fn)
|
|
if aligned is not None:
|
|
out[target_name] = aligned
|
|
return out
|
|
|
|
|
|
def load_state_dict(checkpoint_path: str, model_class) -> dict:
|
|
"""Load a state dict from a checkpoint directory.
|
|
|
|
Supported formats:
|
|
- pytorch_model_fsdp.bin, optionally wrapped as {"state_dict": ...}
|
|
- model.safetensors
|
|
|
|
If the model class reports that the state dict is not fused, it is converted
|
|
through ``model_class.convert_to_fused``.
|
|
"""
|
|
fsdp_ckpt = os.path.join(checkpoint_path, "pytorch_model_fsdp.bin")
|
|
safetensor_ckpt = os.path.join(checkpoint_path, "model.safetensors")
|
|
|
|
if os.path.exists(fsdp_ckpt):
|
|
state_dict = torch.load(fsdp_ckpt, map_location="cpu")
|
|
if isinstance(state_dict, dict) and "state_dict" in state_dict:
|
|
state_dict = state_dict["state_dict"]
|
|
elif os.path.exists(safetensor_ckpt):
|
|
state_dict = load_file(safetensor_ckpt, device="cpu")
|
|
else:
|
|
raise FileNotFoundError(
|
|
"checkpoint contains neither pytorch_model_fsdp.bin nor model.safetensors: "
|
|
f"{checkpoint_path}"
|
|
)
|
|
|
|
if not model_class.is_fused(state_dict):
|
|
state_dict = model_class.convert_to_fused(state_dict)
|
|
|
|
return state_dict
|
|
|
|
|
|
def read_global_step(checkpoint_path: str) -> int | None:
|
|
"""Read ``global_step.pth`` when present."""
|
|
p = os.path.join(checkpoint_path, "global_step.pth")
|
|
if not os.path.exists(p):
|
|
return None
|
|
payload = torch.load(p)
|
|
return int(payload["global_step"])
|
|
|
|
|
|
def _dir_has_weights(path: str) -> bool:
|
|
return os.path.exists(
|
|
os.path.join(path, "pytorch_model_fsdp.bin")
|
|
) or os.path.exists(os.path.join(path, "model.safetensors"))
|
|
|
|
|
|
def resolve_checkpoint_dir(checkpoint_path: str) -> str:
|
|
"""Return a directory that directly contains model weights.
|
|
|
|
Training saves under a root such as ``libero6/`` with step subdirs
|
|
``libero6/0/``, ``libero6/3/``, etc. Inference callers may pass either the
|
|
root or a concrete step directory.
|
|
"""
|
|
if os.path.isfile(checkpoint_path):
|
|
checkpoint_path = os.path.dirname(checkpoint_path)
|
|
|
|
if _dir_has_weights(checkpoint_path):
|
|
return checkpoint_path
|
|
|
|
if not os.path.isdir(checkpoint_path):
|
|
raise FileNotFoundError(f"checkpoint path does not exist: {checkpoint_path}")
|
|
|
|
candidates: list[tuple[int, float, str]] = []
|
|
for entry in os.listdir(checkpoint_path):
|
|
sub = os.path.join(checkpoint_path, entry)
|
|
if not os.path.isdir(sub) or not _dir_has_weights(sub):
|
|
continue
|
|
step = read_global_step(sub)
|
|
sort_step = step if step is not None else -1
|
|
candidates.append((sort_step, os.path.getmtime(sub), sub))
|
|
|
|
if not candidates:
|
|
return checkpoint_path
|
|
|
|
candidates.sort()
|
|
resolved = candidates[-1][2]
|
|
if resolved != checkpoint_path:
|
|
import logging
|
|
|
|
logging.getLogger(__name__).info(
|
|
"Resolved checkpoint root %s -> %s (global_step=%s)",
|
|
checkpoint_path,
|
|
resolved,
|
|
read_global_step(resolved),
|
|
)
|
|
return resolved
|