VLA/wall_x/_vendor/harrix/utils/ckpt_load.py

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