File size: 3,393 Bytes
f9d0103 | 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 | """Extend existing adapters without changing deployment encoding or adding inference modules."""
import torch
from veyra.option_model import LoRALinear, install_adapters
def layer_number(name):
return int(name.split("language_model.layers.", 1)[1].split(".", 1)[0])
def expand_depth(model):
previous = tuple(model.adapter_names)
total = len(model.encoder.model.language_model.layers)
old_depth = model.adaptation["layers"]
boundary = total - old_depth
if total != 24 or old_depth != 12 or not previous:
raise ValueError("this feasibility design requires the declared 12-of-24-layer parent")
if {layer_number(name) for name in previous} != set(range(boundary, total)):
raise ValueError("parent adapter placement changed")
added = install_adapters(
model.encoder.model, total, model.adaptation["rank"], model.adaptation["alpha"]
)
if set(added) & set(previous) or {layer_number(name) for name in added} != set(range(boundary)):
raise ValueError("earlier adapter placement differs from the declared expansion")
for name in added:
layer = model.encoder.model.get_submodule(name)
if not isinstance(layer, LoRALinear) or torch.count_nonzero(layer.lora_b):
raise ValueError("new adapters must initially have exactly zero output")
model.adapter_names = sorted((*previous, *added))
model.adaptation = {**model.adaptation, "layers": total}
model.condition_readout.requires_grad_(False)
model.encoder.model.gradient_checkpointing_enable(
gradient_checkpointing_kwargs={"use_reentrant": False}
)
if not all(layer.gradient_checkpointing for layer in model.encoder.model.language_model.layers):
raise ValueError("language-layer gradient checkpointing was not enabled")
return {"previous_layers": old_depth, "expanded_layers": total, "boundary": boundary}
def training_mode(model):
language = model.encoder.model.language_model
if language.config.attention_dropout != 0 or any(
isinstance(module, torch.nn.Dropout) and module.p != 0 for module in language.modules()
):
raise ValueError("matched checkpointed training requires zero language dropout")
model.train()
model.encoder.model.eval()
# Checkpointing is activated by decoder-layer training mode. The vision stack stays in eval.
language.train()
def optimizer_groups(model, earlier_lr, later_lr, readout_lr):
earlier, later = [], []
for name, parameter in model.named_parameters():
if ".lora_" in name:
(earlier if layer_number(name) < 12 else later).append(parameter)
readout = list(model.readout.parameters()) + list(model.binding_head.parameters())
groups = [
{"params": earlier, "lr": earlier_lr, "name": "earlier_adapters"},
{"params": later, "lr": later_lr, "name": "later_adapters"},
{"params": readout, "lr": readout_lr, "name": "readout"},
]
if any(not group["params"] for group in groups):
raise ValueError("an expected trainable parameter group is empty")
expected = {id(p) for p in model.parameters() if p.requires_grad}
actual = [id(p) for group in groups for p in group["params"]]
if len(actual) != len(set(actual)) or set(actual) != expected:
raise ValueError("optimizer groups do not exactly cover trainable parameters")
return groups
|