Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/_vendor/harrix/utils/ckpt_load.py

@@ -7,6 +7,8 @@ are intentionally left to the adapter.
 from __future__ import annotations
 
 import os
+import json
+import math
 from typing import Callable, Optional
 
 import torch
@@ -17,6 +19,58 @@ 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,
@@ -64,22 +118,58 @@ def _align_checkpoint_tensor(
 
 
 def reshape_compatible_state_dict(
-    state_dict: dict, model_sd: dict, log_fn: Optional[Callable] = None
+    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():
-        if name not in model_sd:
+        # 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[name]
+        target = model_sd[target_name]
         if param.shape == target.shape:
-            out[name] = param
+            out[target_name] = param
             continue
-        aligned = _align_checkpoint_tensor(param, target, name, log_fn)
+        aligned = _align_checkpoint_tensor(param, target, target_name, log_fn)
         if aligned is not None:
-            out[name] = aligned
+            out[target_name] = aligned
     return out