File size: 6,871 Bytes
f0e7192 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """Freeze a selected adapter or full model as a standard HF sequence checkpoint."""
# ruff: noqa: PLC0415
import argparse
import hashlib
import json
from pathlib import Path
from .data import load_contract
from .model import configuration_overrides, load_model, load_task_config
from .optimization import validate_method, verify_full_checkpoint
from .runtime_mapping import RUNTIME_TASKS, write_runtime_mappings
from .task_head import (
assert_task_head_preserved,
task_head_scope,
verify_saved_task_head,
)
def validate_provenance(args):
validate_method(args.method, args.adapter, False)
run = json.loads(args.run_manifest.read_text())
if run.get("test_used") is not False:
raise ValueError("Candidate provenance must explicitly exclude test selection")
if run.get("base_revision") != args.base_revision:
raise ValueError("Export base revision differs from the training receipt")
if run.get("base_model", args.base_id) != args.base_id:
raise ValueError("Export base model differs from the training receipt")
if run.get("method", "lora") != args.method:
raise ValueError("Export method differs from the training receipt")
if args.method == "full":
origin = json.loads((Path(args.base) / "training-origin.json").read_text())
expected = {
"method": "full",
"base_model": args.base_id,
"base_revision": args.base_revision,
}
if origin != expected:
raise ValueError("Full checkpoint origin differs from the training receipt")
config = json.loads((Path(args.base) / "config.json").read_text())
labels, ids = load_contract(args.contract)
contract = configuration_overrides(args.contract)
contract.update(label2id=labels, id2label={str(k): v for k, v in ids.items()})
if any(config.get(key) != value for key, value in contract.items()):
raise ValueError("Export contract differs from the trained full checkpoint")
def main():
import torch
parser = argparse.ArgumentParser()
parser.add_argument("--base", required=True)
parser.add_argument("--base-id", required=True)
parser.add_argument("--base-revision", required=True)
parser.add_argument("--method", choices=["lora", "full"], default="lora")
parser.add_argument("--adapter", type=Path)
parser.add_argument("--contract", required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--run-manifest", type=Path, required=True)
parser.add_argument("--runtime-task", choices=RUNTIME_TASKS, required=True)
args = parser.parse_args()
if args.output.exists() and any(args.output.iterdir()):
raise ValueError("Refusing to overwrite a frozen candidate")
validate_provenance(args)
torch.set_num_threads(8)
model, tokenizer, labels, _ = load_model(args.base, args.contract, args.adapter)
model.eval()
if args.method == "full":
verify_full_checkpoint(model, args.base)
head = task_head_scope(model)
else:
head = verify_saved_task_head(model, args.adapter / "adapter_model.safetensors")
tokens = tokenizer(
["Numerical checkpoint equivalence probe.", "用于核对模型数值的输入。"],
padding=True,
return_tensors="pt",
)
with torch.inference_mode():
before = model(**tokens).logits
exported = (
model.merge_and_unload(safe_merge=True) if args.method == "lora" else model
)
after = exported(**tokens).logits
assert_task_head_preserved(head, exported)
torch.testing.assert_close(before, after, atol=2e-4, rtol=2e-4)
if type(exported).__name__ != "ModernBertForSequenceClassification":
raise ValueError(
"Export must preserve the standard supported HF sequence architecture"
)
if exported.config.label2id != labels:
raise ValueError("Export changed label IDs")
exported.config._name_or_path = args.base_id
args.output.mkdir(parents=True, exist_ok=True)
exported.save_pretrained(args.output, safe_serialization=True)
verify_full_checkpoint(exported, args.output)
tokenizer.save_pretrained(args.output)
restored = (
type(exported)
.from_pretrained(
args.output,
config=load_task_config(args.output),
torch_dtype=torch.float32,
attn_implementation="sdpa",
)
.eval()
)
verify_full_checkpoint(restored, args.output)
assert_task_head_preserved(head, restored)
with torch.inference_mode():
reloaded_logits = restored(**tokens).logits
torch.testing.assert_close(after, reloaded_logits, atol=0, rtol=0)
mapping_files = write_runtime_mappings(args.output, labels, args.runtime_task)
receipt = {
"method": args.method,
"base_model": args.base_id,
"base_revision": args.base_revision,
"run_manifest_sha256": hashlib.sha256(
args.run_manifest.read_bytes()
).hexdigest(),
"parameter_count": sum(
parameter.numel() for parameter in exported.parameters()
),
"encoder_parameter_count": sum(
parameter.numel() for parameter in exported.model.parameters()
),
"label2id": labels,
"runtime_task": args.runtime_task,
"runtime_mapping_files": mapping_files,
"task": exported.config.problem_type,
"classifier_pooling": exported.config.classifier_pooling,
"max_position_embeddings": exported.config.max_position_embeddings,
"position_capacity_is_quality_evidence": False,
"weight_dtype": "float32",
"task_head": head,
"task_head_preserved_after_reload": True,
"all_tensors_preserved_after_reload": True,
"test_used_for_selection": False,
"files": {
path.name: hashlib.sha256(path.read_bytes()).hexdigest()
for path in sorted(args.output.iterdir())
if path.is_file()
},
}
if args.method == "lora":
receipt.update(
adapter_sha256=hashlib.sha256(
(args.adapter / "adapter_model.safetensors").read_bytes()
).hexdigest(),
task_head_preserved_after_merge_and_reload=True,
merge_probe_max_absolute_difference=float((before - after).abs().max()),
)
(args.output / "candidate-lock.json").write_text(
json.dumps(receipt, indent=2) + "\n"
)
print(
json.dumps(
{
key: receipt[key]
for key in (
"method",
"parameter_count",
"all_tensors_preserved_after_reload",
)
}
),
flush=True,
)
if __name__ == "__main__":
main()
|