File size: 2,919 Bytes
45cf443 3d0a2b3 57010bc 3d0a2b3 c6bc767 3d0a2b3 45cf443 | 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 | import torch
from torch import optim, nn
# LoRA for MiniMind-series backbones.
# Applies to every nn.Linear whose in_features == out_features, i.e. the
# q/k/v/o_proj and gate/up/down_proj inside Block's Attention/MLP.
# Because VLM/Omni reuse the same MiniMind Attention/MLP for their LLM
# (thinker)主干, apply_lora also works on VLM / VAM, but only
# touches the shared LLM layers -- vision/audio projectors and the speech
# talker head (in != out) are left untouched by design.
class LoRA(nn.Module):
def __init__(self, in_features, out_features, rank):
super().__init__()
self.rank = rank
self.A = nn.Linear(in_features, rank, bias=False)
self.B = nn.Linear(rank, out_features, bias=False)
self.A.weight.data.normal_(mean=0.0, std=0.02)
self.B.weight.data.zero_()
def forward(self, x):
return self.B(self.A(x))
def apply_lora(model, rank=16):
for name, module in model.named_modules():
if isinstance(module, nn.Linear) and module.in_features == module.out_features:
lora = LoRA(module.in_features, module.out_features, rank=rank).to(model.device)
setattr(module, "lora", lora)
original_forward = module.forward
def forward_with_lora(x, layer1=original_forward, layer2=lora):
return layer1(x) + layer2(x)
module.forward = forward_with_lora
def load_lora(model, path):
state_dict = torch.load(path, map_location=model.device)
state_dict = {(k[7:] if k.startswith('module.') else k): v for k, v in state_dict.items()}
for name, module in model.named_modules():
if hasattr(module, 'lora'):
lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
module.lora.load_state_dict(lora_state)
def save_lora(model, path):
raw_model = getattr(model, '_orig_mod', model)
state_dict = {}
for name, module in raw_model.named_modules():
if hasattr(module, 'lora'):
clean_name = name[7:] if name.startswith("module.") else name
lora_state = {f'{clean_name}.lora.{k}': v.cpu().half() for k, v in module.lora.state_dict().items()}
state_dict.update(lora_state)
torch.save(state_dict, path)
def merge_lora(model, lora_path, save_path):
load_lora(model, lora_path)
raw_model = getattr(model, '_orig_mod', model)
state_dict = {k: v.cpu().half() for k, v in raw_model.state_dict().items() if '.lora.' not in k}
for name, module in raw_model.named_modules():
if isinstance(module, nn.Linear) and '.lora.' not in name:
state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
if hasattr(module, 'lora'):
state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
torch.save(state_dict, save_path)
|