File size: 3,494 Bytes
3ea3da7 | 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 | """Fuse equal-rotation packed QKV and gate/up projections without BF16 weights."""
import types
import torch
from torch import nn
from orbitquant import OrbitQuantConfig
from orbitquant.layers import OrbitQuantLinear
class PackedSlice(nn.Module):
"""Compatibility view; the owning attention/MLP registers the packed group."""
def __init__(self, group, start, length):
super().__init__()
object.__setattr__(self, '_group', group)
self.start, self.out_features, self.in_features = start, length, group.in_features
self.bias = None
def forward(self, x):
return self._group(x)[..., self.start:self.start+self.out_features]
def group(modules, config, name):
first=modules[0]
if not all(isinstance(m,OrbitQuantLinear) and m.in_features==first.in_features and m.bias is None for m in modules):
raise ValueError('Fusion requires compatible bias-free packed linears')
for m in modules[1:]:
if (m.weight_bits,m.activation_bits,m.activation_eps,m.rotation.seed,m.rotation.block_size)!=(first.weight_bits,first.activation_bits,first.activation_eps,first.rotation.seed,first.rotation.block_size):
raise ValueError('Fusion requires identical quantization and rotation')
return OrbitQuantLinear(in_features=first.in_features,out_features=sum(m.out_features for m in modules),
config=config,module_name=name,source_weight_layout='out_in',bias=None,
packed_weight_indices=torch.cat([m.packed_weight_indices for m in modules]),
row_norms=torch.cat([m.row_norms for m in modules]),debug_weight=None).eval()
def project_qkv(self,x,cos,sin):
from yue2.modeling_yue2 import _apply_rotary
b,t,_=x.shape
q,k,v=self.packed_qkv(x).split((self.num_heads*self.head_dim,self.num_kv_heads*self.head_dim,self.num_kv_heads*self.head_dim),dim=-1)
q=self.q_norm(q.reshape(b,t,self.num_heads,self.head_dim))
k=self.k_norm(k.reshape(b,t,self.num_kv_heads,self.head_dim))
return _apply_rotary(q,cos.unsqueeze(2),sin.unsqueeze(2)), _apply_rotary(k,cos.unsqueeze(2),sin.unsqueeze(2)), v.reshape(b,t,self.num_kv_heads,self.head_dim)
def mlp_forward(self,x):
gate,up=self.packed_gate_up(x).chunk(2,dim=-1)
return self.down_proj(torch.nn.functional.silu(gate)*up)
def fuse_model(model,manifest):
if getattr(model,'_orbitquant_fused',False): return model
config=OrbitQuantConfig.from_dict(manifest['config'])
for i,layer in enumerate(model.model.layers):
for attr in ('self_attn','nar_self_attn'):
attn=getattr(layer,attr)
parts=[getattr(attn,n) for n in ('q_proj','k_proj','v_proj')]
packed=group(parts,config,f'model.layers.{i}.{attr}.packed_qkv')
attn.packed_qkv=packed
offset=0
for name,part in zip(('q_proj','k_proj','v_proj'),parts):
setattr(attn,name,PackedSlice(packed,offset,part.out_features));offset+=part.out_features
attn.project_qkv=types.MethodType(project_qkv,attn)
del parts
for attr in ('mlp','nar_mlp'):
mlp=getattr(layer,attr)
n=mlp.gate_proj.out_features
packed=group([mlp.gate_proj,mlp.up_proj],config,f'model.layers.{i}.{attr}.packed_gate_up')
mlp.packed_gate_up=packed
mlp.gate_proj=PackedSlice(packed,0,n);mlp.up_proj=PackedSlice(packed,n,n)
mlp.forward=types.MethodType(mlp_forward,mlp)
model._orbitquant_fused=True
return model
|