File size: 1,796 Bytes
b9fc164 | 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 | """Validate and gather adapter weights without materializing the frozen backbone."""
from collections.abc import Mapping
import torch
import torch.distributed as dist
from torch.distributed.tensor import DTensor
def validate_lora_state_dict(model, weights, adapter_name: str) -> None:
from peft.utils import get_peft_model_state_dict
if not isinstance(weights, Mapping):
raise ValueError(f"LoRA adapter {adapter_name!r} must contain a state dictionary.")
expected = get_peft_model_state_dict(model, adapter_name=adapter_name, save_embedding_layers=False)
missing = sorted(expected.keys() - weights.keys())
unexpected = sorted(weights.keys() - expected.keys())
mismatched = [
name for name in expected.keys() & weights.keys()
if not isinstance(weights[name], torch.Tensor) or weights[name].shape != expected[name].shape
]
if not expected or missing or unexpected or mismatched:
raise ValueError(
f"LoRA adapter {adapter_name!r} does not match the model: "
f"missing={missing[:20]}, unexpected={unexpected[:20]}, "
f"shape_mismatch={sorted(mismatched)[:20]}."
)
def gather_lora_state_dict(model, adapter_name: str) -> dict[str, torch.Tensor]:
from peft.utils import get_peft_model_state_dict
# FSDP2 state_dict() returns DTensors without gathering the full model.
adapter_state = get_peft_model_state_dict(model, adapter_name=adapter_name, save_embedding_layers=False)
is_writer = not dist.is_initialized() or dist.get_rank() == 0
gathered = {}
for name, value in adapter_state.items():
if isinstance(value, DTensor):
value = value.full_tensor()
if is_writer:
gathered[name] = value.detach().cpu()
return gathered
|