| """Merge a LoRA adapter (+ trained readout) into a d3 package, keeping the package layout. |
| |
| python -m d3train.merge --base <d3 package> --adapter <run>/final --out <merged package> [--name d3-mini-ft] |
| |
| The base shards are rewritten tensor by tensor under the same names, shard files and index: an adapted weight |
| becomes ``(W.float() + scale * B @ A).to(W.dtype)``, every other tensor is copied unchanged (so a zero-step adapter, |
| B = 0, reproduces every tensor bit for bit). ``readout.safetensors`` comes from the run (FP32, [255, H]); |
| config.json, decision_config.json (temperature optionally replaced), tokenizer, chat template and the runtime code |
| are copied. MODEL_MANIFEST.json is regenerated for the new files, so ``verify="fast"`` keeps working. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import re |
| import shutil |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import load_file, save_file |
|
|
| WEIGHTS = re.compile(r"^model(-\d{5}-of-\d{5})?\.safetensors$") |
| REGENERATED = {"MODEL_MANIFEST.json", "readout.safetensors", "model.safetensors.index.json", "decision_config.json"} |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with open(path, "rb") as stream: |
| for block in iter(lambda: stream.read(16 << 20), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def lora_deltas(adapter: Path) -> dict[str, torch.Tensor]: |
| """{checkpoint weight name: FP32 delta} from a PEFT LoRA adapter saved by d3train.train.""" |
| config = json.loads((adapter / "adapter_config.json").read_text(encoding="utf-8")) |
| if config.get("peft_type") != "LORA" or config.get("use_dora") or config.get("fan_in_fan_out"): |
| raise ValueError("only plain LoRA adapters (no DoRA, fan_in_fan_out False) are supported") |
| r, alpha = config["r"], config["lora_alpha"] |
| if config.get("rank_pattern") or config.get("alpha_pattern"): |
| raise ValueError("rank/alpha patterns are not supported") |
| scale = alpha / (r ** 0.5) if config.get("use_rslora") else alpha / r |
| tensors = load_file(str(adapter / "adapter_model.safetensors")) |
| a_suffix, b_suffix = ".lora_A.weight", ".lora_B.weight" |
| deltas = {} |
| for name, a in tensors.items(): |
| if not name.endswith(a_suffix): |
| continue |
| module = name[: -len(a_suffix)] |
| b = tensors[module + b_suffix] |
| key = module.removeprefix("base_model.model.") + ".weight" |
| deltas[key] = (b.float() @ a.float()) * scale |
| unknown = [n for n in tensors if not n.endswith((a_suffix, b_suffix))] |
| if unknown: |
| raise ValueError(f"adapter holds tensors other than LoRA A/B: {unknown[:4]}") |
| return deltas |
|
|
|
|
| def check_llamacpp_acceptance(root: Path) -> list[str]: |
| """What llama.cpp's conversion/vllm_d3.py + pplx_decider.py check before converting; [] when accepted.""" |
| problems = [] |
| for name in ("config.json", "decision_config.json", "readout.safetensors"): |
| if not (root / name).is_file(): |
| problems.append(f"missing {name}") |
| if problems: |
| return problems |
| decision = json.loads((root / "decision_config.json").read_text(encoding="utf-8")) |
| if decision.get("format_id") != "d3-code-readout-v1": |
| problems.append("format_id is not d3-code-readout-v1 (not detected as vllm-d3)") |
| if decision.get("prompt") != "d3": |
| problems.append("prompt != d3") |
| if decision.get("attention_mode") != "noncausal_full_attention": |
| problems.append("attention_mode != noncausal_full_attention") |
| if decision.get("pooling") != "last": |
| problems.append("pooling != last") |
| if not isinstance(decision.get("temperature"), (int, float)): |
| problems.append("temperature missing (written for choice/score/noul)") |
| with safe_open(str(root / "readout.safetensors"), framework="pt") as handle: |
| if list(handle.keys()) != ["weight"]: |
| problems.append("readout.safetensors must hold exactly 'weight'") |
| else: |
| shape = handle.get_slice("weight").get_shape() |
| ids = decision.get("token_ids") or [] |
| if not (shape[0] == len(ids) == len(set(ids))): |
| problems.append("readout rows != number of distinct token_ids") |
| config = json.loads((root / "config.json").read_text(encoding="utf-8")) |
| vocab = (config.get("text_config") or config).get("vocab_size") |
| if any(not 0 <= i < (vocab or 0) for i in decision.get("token_ids") or []): |
| problems.append("token id outside vocab") |
| if config.get("model_type") != "qwen3_5": |
| problems.append("model_type != qwen3_5 (converter maps vllm-d3 to MODEL_ARCH.QWEN35)") |
| return problems |
|
|
|
|
| def merge( |
| base: Path, adapter: Path, out: Path, *, name: str | None = None, temperature: float | None = None, |
| ) -> dict[str, Any]: |
| base, adapter, out = Path(base), Path(adapter), Path(out) |
| if out.exists() and any(out.iterdir()): |
| raise FileExistsError(f"{out} is not empty") |
| out.mkdir(parents=True, exist_ok=True) |
| deltas = lora_deltas(adapter) if (adapter / "adapter_config.json").is_file() else {} |
| readout_file = adapter / "readout.safetensors" |
| readout = load_file(str(readout_file if readout_file.is_file() else base / "readout.safetensors"))["weight"] |
| if readout.dtype != torch.float32: |
| raise ValueError("the readout must be FP32") |
|
|
| |
| for path in sorted(base.rglob("*")): |
| rel = path.relative_to(base) |
| if path.is_dir() or any(p.startswith(".") or p == "__pycache__" for p in rel.parts): |
| continue |
| if rel.as_posix() in REGENERATED or WEIGHTS.match(rel.name) or rel.name.endswith(".header.json"): |
| continue |
| target = out / rel |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy2(path, target) |
|
|
| |
| index_path = base / "model.safetensors.index.json" |
| if index_path.is_file(): |
| index = json.loads(index_path.read_text(encoding="utf-8")) |
| shards = sorted(set(index["weight_map"].values())) |
| else: |
| index, shards = None, ["model.safetensors"] |
| applied, changed = set(), 0 |
| for shard in shards: |
| tensors, metadata = {}, None |
| with safe_open(str(base / shard), framework="pt") as handle: |
| metadata = handle.metadata() |
| for key in handle.keys(): |
| tensor = handle.get_tensor(key) |
| if key in deltas: |
| delta = deltas[key] |
| if delta.shape != tensor.shape: |
| raise ValueError(f"{key}: delta {tuple(delta.shape)} vs weight {tuple(tensor.shape)}") |
| merged = (tensor.float() + delta).to(tensor.dtype) |
| changed += int(not torch.equal(merged, tensor)) |
| tensor = merged |
| applied.add(key) |
| tensors[key] = tensor.contiguous() |
| save_file(tensors, str(out / shard), metadata=metadata) |
| missing = sorted(set(deltas) - applied) |
| if missing: |
| raise ValueError(f"adapter weights not found in the base checkpoint: {missing[:4]}") |
| if index is not None: |
| |
| shutil.copy2(index_path, out / "model.safetensors.index.json") |
|
|
| save_file({"weight": readout.contiguous()}, str(out / "readout.safetensors")) |
| decision = json.loads((base / "decision_config.json").read_text(encoding="utf-8")) |
| if temperature is not None: |
| decision["temperature"] = float(temperature) |
| (out / "decision_config.json").write_text(json.dumps(decision, indent=2) + "\n", encoding="utf-8") |
| notice = out / "NOTICE" |
| if notice.is_file(): |
| notice.write_text( |
| notice.read_text(encoding="utf-8").rstrip("\n") |
| + f"\n\nModified: LoRA fine-tune merged into the text backbone and readout ({datetime.now(timezone.utc).date()}).\n", |
| encoding="utf-8", |
| ) |
|
|
| problems = check_llamacpp_acceptance(out) |
| report = {"adapted_tensors": len(applied), "changed_tensors": changed, "llamacpp_acceptance": problems or "ok"} |
| (out / "MERGE-REPORT.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") |
|
|
| |
| old = {} |
| if (base / "MODEL_MANIFEST.json").is_file(): |
| old = json.loads((base / "MODEL_MANIFEST.json").read_text(encoding="utf-8")) |
| files = sorted( |
| p.relative_to(out).as_posix() for p in out.rglob("*") |
| if p.is_file() and p.name != "MODEL_MANIFEST.json" |
| and not any(part.startswith(".") or part == "__pycache__" for part in p.relative_to(out).parts) |
| ) |
| manifest = { |
| "schema": "d3-package-manifest/1", |
| **{k: old[k] for k in ("format_id", "prompt", "attention_mode", "max_input_tokens", "base_model", "backbone", |
| "licence", "parameters") if k in old}, |
| "model_name": name or f"{old.get('model_name', base.name)}-ft", |
| "derived_from": {"model_name": old.get("model_name"), "repo_id": old.get("repo_id"), |
| "built_utc": old.get("built_utc"), "adapter": adapter.name}, |
| "built_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), |
| "files_sha256": {f: sha256_file(out / f) for f in files}, |
| "files_bytes": {f: (out / f).stat().st_size for f in files}, |
| } |
| (out / "MODEL_MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") |
| return report |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| ap.add_argument("--base", required=True, type=Path) |
| ap.add_argument("--adapter", required=True, type=Path, help="run checkpoint dir (adapter + readout.safetensors)") |
| ap.add_argument("--out", required=True, type=Path) |
| ap.add_argument("--name") |
| ap.add_argument("--temperature", type=float, help="write a recalibrated temperature into decision_config.json") |
| args = ap.parse_args() |
| print(json.dumps(merge(args.base, args.adapter, args.out, name=args.name, temperature=args.temperature))) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
