48 lines
2.7 KiB
Python
48 lines
2.7 KiB
Python
"""Read-only FP32 effective weights for an existing pure-LoRA checkpoint."""
|
|
import re
|
|
import numpy as np
|
|
from functional.model import PREFIX,StreamingLM
|
|
from pilot.tensors import load_tensors
|
|
BASE_REVISION='15852e8c16360a2fea060d615a32b45270f8a8fc'
|
|
TARGETS={'in_proj_qkv','in_proj_a','in_proj_b','out_proj'}
|
|
|
|
class EffectiveWeight:
|
|
def __init__(self,base,a,b,scale):
|
|
assert a.ndim==b.ndim==2 and (b.shape[0],a.shape[1])==base.shape and a.shape[0]==b.shape[1]
|
|
self.base,self.a,self.b,self.scale=base,a,b,scale
|
|
self.shape,self.numel,self.dtype=base.shape,base.numel,'F32'
|
|
def read(self,selector=slice(None)):
|
|
if not isinstance(selector,slice):raise ValueError('Only contiguous row chunks supported')
|
|
start,stop,step=selector.indices(self.numel);dim=self.shape[1]
|
|
if step!=1 or start%dim or stop%dim:raise ValueError('Row aligned chunks required')
|
|
w=self.base.read(slice(start,stop)).reshape(-1,dim)
|
|
return (w+self.scale*(self.b[start//dim:stop//dim]@self.a)).astype(np.float32).reshape(-1)
|
|
def evict(self):self.base.evict()
|
|
|
|
|
|
def install(lm,directory,meta,multiplier=1.):
|
|
assert meta['base_model']=='Qwen/Qwen3.5-2B' and meta['base_revision']==BASE_REVISION and meta['gate'] is None and meta['arm']=='lora'
|
|
tensors=load_tensors(directory);pairs={};pattern=re.compile(r'^base_model\.model\.model\.(layers\.(\d+)\.linear_attn\.(in_proj_qkv|in_proj_a|in_proj_b|out_proj))\.lora_([AB])\.default\.weight$')
|
|
for name,tensor in tensors.items():
|
|
match=pattern.fullmatch(name)
|
|
if not match:raise ValueError('Unexpected trainable tensor: '+name)
|
|
target=PREFIX+match[1]+'.weight';assert lm.config['layer_types'][int(match[2])]=='linear_attention'
|
|
assert target in lm.tensors
|
|
value=tensor.read().reshape(tensor.shape).copy();tensor.evict();assert np.isfinite(value).all()
|
|
pairs.setdefault(target,{})[match[4]]=value
|
|
expected={PREFIX+f'layers.{layer}.linear_attn.{role}.weight' for layer,kind in enumerate(lm.config['layer_types']) if kind=='linear_attention' for role in TARGETS}
|
|
assert set(pairs)==expected and len(tensors)==meta['n_tensors']
|
|
scale=meta['lora_alpha']/meta['lora_rank']*multiplier
|
|
for target,pair in pairs.items():
|
|
assert set(pair)=={'A','B'} and pair['A'].shape[0]==meta['lora_rank']
|
|
lm.tensors[target]=EffectiveWeight(lm.tensors[target],pair['A'],pair['B'],scale)
|
|
return pairs,scale
|
|
|
|
|
|
class ExplicitLoRA(StreamingLM):
|
|
"""Independent reference: base projection + scaled low-rank output, no W merge."""
|
|
def linear(self,x,name,rows=256):
|
|
y=super().linear(x,name,rows);pair=self.pairs.get(PREFIX+name+'.weight')
|
|
if pair is not None:y=y+self.scale*((x@pair['A'].T)@pair['B'].T)
|
|
return y.astype(np.float32)
|