File size: 8,773 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 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | """Explicit adapter or full-encoder optimization and checkpoint ownership."""
# ruff: noqa: PLC0415
import hashlib
import json
import math
from pathlib import Path
from .model import load_model, save_adapter
from .task_head import task_head_scope, tensor_digest
def configure_trainable_tail(model, last_layers):
"""Restrict a complete ModernBERT checkpoint to its tail and entire task head."""
if (
type(model).__name__ != "ModernBertForSequenceClassification"
or hasattr(model, "peft_config")
or hasattr(model, "get_base_model")
):
raise ValueError("Tail training requires a complete ModernBERT classifier")
layers = model.model.layers
if type(last_layers) is not int or not 1 <= last_layers <= len(layers):
raise ValueError("Trainable tail must contain one through all encoder layers")
start = len(layers) - last_layers
modules = [*layers[start:], model.head, model.classifier]
selected = {
id(parameter) for module in modules for parameter in module.parameters()
}
named = dict(model.named_parameters())
if not selected or not selected <= {id(parameter) for parameter in named.values()}:
raise ValueError("Trainable modules do not belong to the classifier")
if any(
id(parameter) in selected for parameter in model.model.final_norm.parameters()
):
raise ValueError("Encoder final normalization must remain frozen")
model.requires_grad_(False)
for module in modules:
module.requires_grad_(True)
tensors = {
name: {"shape": list(parameter.shape), "numel": parameter.numel()}
for name, parameter in named.items()
if parameter.requires_grad
}
return {
"architecture": type(model).__name__,
"layer_indices": list(range(start, len(layers))),
"trainable_tensors": tensors,
"trainable_parameters": sum(item["numel"] for item in tensors.values()),
"encoder_final_norm_trainable": False,
}
def frozen_parameter_receipt(model):
"""Hash frozen parameters once so an optional restricted update is auditable."""
return {
name: {
"shape": list(parameter.shape),
"dtype": str(parameter.dtype),
"sha256": tensor_digest(parameter),
}
for name, parameter in model.named_parameters()
if not parameter.requires_grad
}
def assert_frozen_parameters_preserved(expected, model):
if frozen_parameter_receipt(model) != expected:
raise ValueError("Frozen parameter scope or tensor values changed")
def validate_method(method, adapter, fresh_head):
if method not in {"lora", "full"}:
raise ValueError("Training method must be lora or full")
if method == "lora" and (adapter is None or fresh_head):
raise ValueError("LoRA continuation requires its initialized adapter")
if method == "full" and adapter is not None:
raise ValueError(
"Full training starts from a complete checkpoint, not an adapter"
)
def load_trainable_model(base, contract, adapter, method, fresh_head=False):
validate_method(method, adapter, fresh_head)
model, tokenizer, labels, ids = load_model(
base, contract, adapter, trainable=method == "lora"
)
if method == "full":
if type(model).__name__ != "ModernBertForSequenceClassification":
raise ValueError("Full sequence training requires a ModernBERT classifier")
model.requires_grad_(True)
if fresh_head:
from torch import nn
head_parameter = next(model.head.parameters())
model.head = type(model.head)(model.config).to(
device=head_parameter.device, dtype=head_parameter.dtype
)
classifier = model.classifier
model.classifier = nn.Linear(
classifier.in_features,
classifier.out_features,
bias=classifier.bias is not None,
device=classifier.weight.device,
dtype=classifier.weight.dtype,
)
# Fresh modules let HF apply its native initialization without
# mutating the initialization guards on loaded checkpoint tensors.
model.head.apply(model._init_weights)
# ModernBERT initializes the final classifier on its owning task
# module, not when visiting the bare Linear child.
model._init_weights(model)
return model, tokenizer, labels, ids
def optimizer_groups(model, learning_rate, head_learning_rate=None):
rates = [learning_rate]
if head_learning_rate is not None:
rates.append(head_learning_rate)
if any(not math.isfinite(rate) or rate <= 0 for rate in rates):
raise ValueError("Learning rates must be finite and positive")
base = model.get_base_model() if hasattr(model, "get_base_model") else model
head_ids = {
id(parameter)
for module in (base.head, base.classifier)
for parameter in module.parameters()
if parameter.requires_grad
}
encoder, head = [], []
for parameter in model.parameters():
if parameter.requires_grad:
(head if id(parameter) in head_ids else encoder).append(parameter)
groups = []
for name, parameters, rate in (
("encoder", encoder, learning_rate),
("task_head", head, head_learning_rate or learning_rate),
):
if parameters:
groups.append(
{"name": name, "params": parameters, "lr": rate, "initial_lr": rate}
)
if not groups:
raise ValueError("No trainable parameters")
return groups
def initial_artifact_receipt(base, adapter):
"""Hash the actual local initialization, including every weight shard."""
root = Path(base)
if not root.is_dir():
raise ValueError("Training requires a downloaded local base checkpoint")
names = {"config.json", "tokenizer.json", "tokenizer_config.json"}
for path in root.iterdir():
if (
path.suffix == ".safetensors"
or path.name
in {
"model.safetensors.index.json",
"pytorch_model.bin",
"pytorch_model.bin.index.json",
}
or path.name.startswith("pytorch_model-")
):
names.add(path.name)
files = {}
for name in sorted(names):
path = root / name
if path.is_file():
with path.open("rb") as stream:
files[name] = hashlib.file_digest(stream, "sha256").hexdigest()
result = {"base_files": files}
if adapter is not None:
for filename in ("adapter_model.safetensors", "adapter_config.json"):
path = Path(adapter) / filename
if filename == "adapter_config.json" and not path.is_file():
result[filename] = None
continue
with path.open("rb") as stream:
result[filename] = hashlib.file_digest(stream, "sha256").hexdigest()
return result
def verify_full_checkpoint(model, checkpoint):
"""Require every tensor of a complete sequence checkpoint to be preserved."""
if type(model).__name__ != "ModernBertForSequenceClassification":
raise ValueError("Full checkpoints must own the actual encoder weights")
from safetensors.torch import load_file
saved = load_file(str(Path(checkpoint) / "model.safetensors"), device="cpu")
state = model.state_dict()
if saved.keys() != state.keys():
raise ValueError("Full checkpoint omitted or introduced model tensors")
import torch
for name, value in state.items():
if saved[name].dtype != value.dtype or not torch.equal(
saved[name], value.detach().cpu()
):
raise ValueError(f"Full checkpoint changed tensor: {name}")
def save_training_checkpoint(model, tokenizer, destination, method, base_id, revision):
if method == "lora":
save_adapter(model, tokenizer, destination, base_id, revision)
return
if method != "full" or hasattr(model, "peft_config"):
raise ValueError("Full checkpoints must own the actual encoder weights")
destination.mkdir(parents=True, exist_ok=True)
expected = task_head_scope(model)
model.save_pretrained(destination, safe_serialization=True)
verify_full_checkpoint(model, destination)
tokenizer.save_pretrained(destination)
(destination / "task-head.json").write_text(json.dumps(expected, indent=2) + "\n")
(destination / "training-origin.json").write_text(
json.dumps(
{"method": method, "base_model": base_id, "base_revision": revision},
indent=2,
)
+ "\n"
)
|