File size: 18,942 Bytes
e70aa91 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 | """Heterogeneous LFM3 decoder stack with exact source-operation transport.
Machine cells use the compiled circuit's original sparse ReLU/head operations.
No SiLU approximation, normalization, widening or numerical conversion is added.
Language cache indices remain those of the unmodified language layers.
"""
from collections import OrderedDict
import inspect,weakref
import torch
from torch import nn
from .computer.circuit import CompiledComputer,CircuitStage
class MachineDecoderLayer(nn.Module):
decoder_role='machine'
def language_identity(self,hidden_states):return hidden_states
class ComputerControlDecoderLayer(MachineDecoderLayer):
decoder_role='computer_control'
def __init__(self,computer):
super().__init__();self.spec=computer.spec
self.embedding=computer.embedding;self.output=computer.output;self.programs=computer.programs
def forward(self,hidden_states=None,*,phase='language',args=(),kwargs=None):
if phase=='language':return hidden_states
port=self._port();kwargs=kwargs or {}
if phase=='boot':return CompiledComputer.boot(port,*args,**kwargs)
if phase=='embed':return CompiledComputer.embed(port,*args,**kwargs)
raise ValueError('Unknown control phase: '+phase)
class OrderedSparseDownProjection(nn.Module):
"""Exact w2 adapter: retain coefficient multiplication and scatter order."""
def __init__(self,data):super().__init__();self.data=data
def forward(self,values,template):
values=values*self.data.coefficient
return torch.zeros_like(template).index_add_(1,self.data.output,values)
class ExactGatedMLP(nn.Module):
"""LFM gated-MLP decomposition with exact ReLU/sparse projection adapters."""
def __init__(self,operation):
super().__init__();self.w1=operation.b;self.w3=operation.a
self.w2=OrderedSparseDownProjection(operation.data);self.activation=nn.ReLU()
def forward(self,hidden_states,active):
values=self.w3(hidden_states)*self.activation(self.w1(hidden_states))
return self.w2(values,hidden_states)*active[:,None]
class ExactRecurrentAttention(nn.Module):
"""Exact memory attention adapter; persistent writes are source operations."""
def __init__(self,operation):
super().__init__();self.spec=operation.spec;self.data=operation.data
for name in operation.spec['maps']:self.add_module(name,getattr(operation,name))
def forward(self,x,state,active):
d=self.data;memory,history=state;flat=torch.cat((x,history),dim=1)
B,V=x.shape[0],d.khead.numel()
q=self.q(flat).relu().reshape(B,-1,128)[:,d.khead]
k=self.k(flat).relu().reshape(B,-1,128)[:,d.khead]
v=self.v(flat).relu().reshape(B,V,128)
z=self.z(flat).relu().reshape(B,V,128)
a,beta=self.a(flat),self.b(flat)+d.bconst
factor=torch.where(a==0,1.,torch.where(a==110,.5,0.))
updated=memory*factor[:,:,None,None];row=k.argmax(-1)
write=(k.amax(-1)>0)&(beta==1)
selector=torch.arange(128,device=x.device)[None,None,:]==row[:,:,None]
updated=torch.where((write[:,:,None]&selector)[:,:,:,None],v[:,:,None,:],updated)
memory.copy_(torch.where(active[:,None,None,None],updated,memory))
read=((q>0)[:,:,:,None]*memory).sum(2)
reference=read[:,:,:2].square().sum(-1,keepdim=True)
denominator=(reference/128.).sqrt()
normalized=torch.where(reference>0,read/torch.where(denominator>0,denominator,1.),0.)
delta=(normalized*z*d.oweight).reshape(B,-1)
return torch.zeros_like(x).index_add_(1,d.output.flatten(),delta)*active[:,None]
def commit_history(self,x,state,active):
memory,history=state
new_history=torch.cat((x,history[:,:2*x.shape[1]]),dim=1)
history.copy_(torch.where(active[:,None],new_history,history))
class ComputerExecutionDecoderLayer(MachineDecoderLayer):
decoder_role='computer_execution'
def __init__(self,operations):
super().__init__();self.stage_specs=[operation.spec for operation in operations];self.triton_widths={}
self.operator_norm=nn.Identity();self.ffn_norm=nn.Identity();self.use_sandwich_norm=False
head=next((op for op in operations if op.kind=='head'),None);ffn=next((op for op in operations if op.kind=='ffn'),None)
self.self_attn=ExactRecurrentAttention(head) if head is not None else None
self.feed_forward=ExactGatedMLP(ffn) if ffn is not None else None
def forward(self,hidden_states,state=None,active=None,backend='torch',scratch=None,*,phase='language'):
if active is None:return hidden_states
if backend=='triton':
for spec in self.stage_specs:
view=StageAdapter(self,spec)
if view.kind not in self.triton_widths:raise RuntimeError("Call circuit.prepare('triton') before CUDA capture.")
view.widths=self.triton_widths[view.kind]
hidden_states,_=CircuitStage.triton(view,hidden_states,state if view.kind=='head' else None,active,scratch)
return hidden_states,state
# Same operator/residual then gated-MLP/residual organization as LFM3.
# Missing machine sublayers are pruned, not simulated with add-zero.
if self.self_attn is not None:
residual=hidden_states
attention_output=self.self_attn(self.operator_norm(hidden_states),state,active)
hidden_states=residual+attention_output
self.self_attn.commit_history(residual,state,active)
if self.feed_forward is not None:
residual=hidden_states
ffn_output=self.feed_forward(self.ffn_norm(hidden_states),active)
hidden_states=residual+ffn_output
return hidden_states,state
class StageAdapter:
"""Parameter-free view for stable logical state names and program boot."""
def __init__(self,cell,spec):self.cell=cell;self.spec=spec;self.kind=spec['kind'];self.widths=None
@property
def data(self):return self.cell.feed_forward.w2.data if self.kind=='ffn' else self.cell.self_attn.data
def __getattr__(self,name):
if self.kind=='ffn' and name in ('a','b'):return getattr(self.cell.feed_forward,'w3' if name=='a' else 'w1')
if self.kind=='head':return getattr(self.cell.self_attn,name)
raise AttributeError(name)
def logical_modules(self):
for name in self.spec['maps']:yield name,getattr(self,name)
yield 'data',self.data
def prepare_triton(self):
CircuitStage.prepare_triton(self);self.cell.triton_widths[self.kind]=self.widths
class MemoryDecoderLayer(MachineDecoderLayer):
decoder_role='memory'
def __init__(self,ram,files):super().__init__();self.ram=ram;self.files=files
def forward(self,hidden_states=None,*,phase='language',component='ram',method='forward',args=(),kwargs=None):
if phase=='language':return hidden_states
module=getattr(self,component)
return module(*args,**(kwargs or {})) if method=='forward' else getattr(module,method)(*args,**(kwargs or {}))
class ServiceDecoderLayer(MachineDecoderLayer):
def __init__(self,module,role):super().__init__();self.payload=module;self.decoder_role=role
def forward(self,hidden_states=None,*,phase='language',method='forward',args=(),kwargs=None):
if phase=='language':return hidden_states
return self.payload(*args,**(kwargs or {})) if method=='forward' else getattr(self.payload,method)(*args,**(kwargs or {}))
class DecoderModelMixin:
def forward(self,*args,execution_phase=None,**kwargs):
if execution_phase is None or execution_phase=='language':return super().forward(*args,**kwargs)
if args:raise ValueError('Machine execution uses explicit typed arguments')
if execution_phase=='computer':
state=kwargs.pop('computer_state')
if kwargs:raise ValueError('Unexpected computer execution fields')
# The original transition function calls the actual decoder cells
# through the stage view. No separate computer module is retained.
port=self._computer_port
active=(state.current!=port.spec['halt'])&state.valid
x=port.embed(state.current);scratch=torch.zeros_like(x) if port.backend=='triton' else None
hi=0
for cell in self.layers[1:1+self.computer_block_count]:
head=state.heads[hi] if cell.self_attn is not None else None
x,_=cell(x,head,active,port.backend,scratch)
hi+=cell.self_attn is not None
out=port.output
family=(x[:,out.family].round()*out.family_weight).sum(1).long()
payload=(x[:,out.payload].round()*out.payload_weight).sum(1).long()
f=family.clamp(0,port.embedding.start.numel()-1)
valid=((family>=0)&(family<port.embedding.start.numel())&(payload>=0)&(payload<port.embedding.size[f]))
state.valid.copy_(state.valid&(valid|~active));predicted=port.embedding.start[f]+payload
state.current.copy_(torch.where(active,predicted,state.current));return state.current
index=kwargs.pop('decoder_index');method=kwargs.pop('decoder_method','forward')
call_args=kwargs.pop('decoder_args',());call_kwargs=kwargs.pop('decoder_kwargs',{})
component=kwargs.pop('decoder_component',None)
if kwargs:raise ValueError('Unexpected decoder execution fields')
layer=self.layers[index]
if isinstance(layer,ComputerControlDecoderLayer):return layer(phase=execution_phase,args=call_args,kwargs=call_kwargs)
if isinstance(layer,MemoryDecoderLayer):return layer(phase=execution_phase,component=component,method=method,args=call_args,kwargs=call_kwargs)
return layer(phase=execution_phase,method=method,args=call_args,kwargs=call_kwargs)
class StageView:
def __init__(self,owner,specs,mapping):self.owner=weakref.ref(owner);self.specs=specs;self.mapping=mapping
def __len__(self):return len(self.specs)
def __iter__(self):return iter(tuple(self[i] for i in range(len(self))))
def __getitem__(self,index):
if isinstance(index,slice):return list(self)[index]
return StageAdapter(self.owner().layers[self.mapping[index]],self.specs[index])
class ComputerPort:
"""Parameter-free compatibility API; execution always enters decoder.forward."""
def __init__(self,owner,spec,mapping,backend):
self.owner=weakref.ref(owner);self.spec=spec;self.backend=backend;self.stages=StageView(owner,spec['stages'],mapping)
@property
def embedding(self):return self.owner().layers[0].embedding
@property
def output(self):return self.owner().layers[0].output
@property
def programs(self):return self.owner().layers[0].programs
@property
def device(self):return self.embedding.base.device
def __call__(self,state):return self.owner()(execution_phase='computer',computer_state=state)
def embed(self,tokens):return self.owner()(execution_phase='embed',decoder_index=0,decoder_args=(tokens,))
def boot(self,*args,**kwargs):return self.owner()(execution_phase='boot',decoder_index=0,decoder_args=args,decoder_kwargs=kwargs)
def prepare(self,backend='torch'):return CompiledComputer.prepare(self,backend)
def _logical_modules(self):
yield 'embedding',self.embedding;yield 'output',self.output
for i,stage in enumerate(self.stages):
for name,module in stage.logical_modules():yield 'stages.'+str(i)+'.'+name,module
yield 'programs',self.programs
def state_dict(self,destination=None,prefix='',keep_vars=False):
result=OrderedDict() if destination is None else destination
for name,module in self._logical_modules():module.state_dict(destination=result,prefix=prefix+name+'.',keep_vars=keep_vars)
return result
def named_parameters(self,prefix='',recurse=True,remove_duplicate=True):
seen=set()
for name,module in self._logical_modules():
for key,value in module.named_parameters(recurse=recurse,remove_duplicate=remove_duplicate):
if remove_duplicate and id(value) in seen:continue
seen.add(id(value));yield ((prefix+'.') if prefix else '')+name+'.'+key,value
def parameters(self,recurse=True):return (value for _,value in self.named_parameters(recurse=recurse))
class ComponentPort:
def __init__(self,owner,index,component=None):self.owner=weakref.ref(owner);self.index=index;self.component=component
def _module(self):
layer=self.owner().layers[self.index]
return getattr(layer,self.component) if self.component else layer.payload
def _call(self,method,args,kwargs):
return self.owner()(execution_phase='service',decoder_index=self.index,decoder_component=self.component,decoder_method=method,decoder_args=args,decoder_kwargs=kwargs)
def __call__(self,*args,**kwargs):return self._call('forward',args,kwargs)
def __getattr__(self,name):
value=getattr(self._module(),name)
if callable(value) and not isinstance(value,nn.Module) and name not in {'state_dict','named_parameters','parameters','named_buffers','buffers','named_modules','modules','children','named_children','get_submodule'}:
return lambda *args,**kwargs:self._call(name,args,kwargs)
return value
def install_decoder_stack(model):
backbone=model.language_model.model
if getattr(backbone,'computer_decoder_count',0):raise ValueError('Decoder stack already installed')
original=model.computer;stage_count=len(original.stages)
groups=[];i=0
while i<stage_count:
group=[i]
if original.stages[i].kind=='head' and i+1<stage_count and original.stages[i+1].kind=='ffn':group.append(i+1)
groups.append(group);i+=len(group)
mapping=[0]*stage_count
for index,group in enumerate(groups,1):
for source in group:mapping[source]=index
prefix=[ComputerControlDecoderLayer(original)]+[ComputerExecutionDecoderLayer([original.stages[j] for j in group]) for group in groups]
indices={'memory':len(prefix)};prefix.append(MemoryDecoderLayer(model.ram,model.files))
if model.ram_bridge is not None:indices['bridge']=len(prefix);prefix.append(ServiceDecoderLayer(model.ram_bridge,'bridge'))
indices['analytics']=len(prefix);prefix.append(ServiceDecoderLayer(model.analytics,'analytics'))
language_layers=list(backbone.layers);backbone.layers=nn.ModuleList(prefix+language_layers)
backbone.computer_decoder_count=len(prefix);backbone.computer_block_count=len(groups);backbone.language_decoder_count=len(language_layers)
backbone.__class__=type('LFMputerDecoderModel',(DecoderModelMixin,type(backbone)),{})
computer=ComputerPort(backbone,original.spec,mapping,original.backend)
object.__setattr__(backbone,'_computer_port',computer);object.__setattr__(prefix[0],'_port',weakref.ref(computer))
for name in ['computer','ram','files','analytics']:
model._modules.pop(name)
if model.ram_bridge is not None:model._modules.pop('ram_bridge')
object.__setattr__(model,'computer',computer)
object.__setattr__(model,'ram',ComponentPort(backbone,indices['memory'],'ram'))
object.__setattr__(model,'files',ComponentPort(backbone,indices['memory'],'files'))
object.__setattr__(model,'analytics',ComponentPort(backbone,indices['analytics']))
if 'bridge' in indices:object.__setattr__(model,'ram_bridge',ComponentPort(backbone,indices['bridge']))
model.decoder_layout={'computer_stage_count':stage_count,'computer_block_count':len(groups),'stage_groups':groups,'stage_to_decoder':mapping,'computer_stage_kinds':[s.kind for s in original.stages],'machine_layer_count':len(prefix),'language_layer_count':len(language_layers),'total_layer_count':len(backbone.layers),'memory_index':indices['memory'],'bridge_index':indices.get('bridge'),'analytics_index':indices['analytics'],'roles':[layer.decoder_role for layer in prefix]+['language']*len(language_layers),'language_cache_indices':'unchanged; machine state is explicit'}
model.decoder_layout.update(machine_hidden_size=original.spec['dims'],language_hidden_size=backbone.config.hidden_size)
if hasattr(model,'lora_names'):
model.lora_names=[logical_to_decoder_key('language_model.'+name,model.decoder_layout)[len('language_model.'):] for name in model.lora_names]
model._runners={};return model
def logical_to_decoder_key(key,layout):
base='language_model.model.layers.';p=layout['machine_layer_count']
if key.startswith(base):
i,tail=key[len(base):].split('.',1);return base+str(int(i)+p)+'.'+tail
if key.startswith('computer.stages.'):
i,tail=key[len('computer.stages.'):].split('.',1);kind=layout['computer_stage_kinds'][int(i)]
if kind=='ffn':
name,rest=tail.split('.',1);tail={'a':'feed_forward.w3','b':'feed_forward.w1','data':'feed_forward.w2.data'}[name]+'.'+rest
else:tail='self_attn.'+tail
return base+str(layout['stage_to_decoder'][int(i)])+'.'+tail
for name in ['embedding','output','programs']:
prefix='computer.'+name+'.'
if key.startswith(prefix):return base+'0.'+name+'.'+key[len(prefix):]
for prefix,index,child in [('ram.',layout['memory_index'],'ram.'),('files.',layout['memory_index'],'files.'),('ram_bridge.',layout.get('bridge_index'),'payload.'),('analytics.',layout['analytics_index'],'payload.')]:
if key.startswith(prefix):
if index is None:raise ValueError('Bridge tensor without a bridge layer')
return base+str(index)+'.'+child+key[len(prefix):]
return key
def decoder_to_logical_key(key,layout):
base='language_model.model.layers.'
if not key.startswith(base):return key
i,tail=key[len(base):].split('.',1);i=int(i);p=layout['machine_layer_count']
if i>=p:return base+str(i-p)+'.'+tail
if i==0:return 'computer.'+tail
if 1<=i<=layout['computer_block_count']:
kind='ffn' if tail.startswith('feed_forward.') else 'head'
source=next(j for j,n in enumerate(layout['stage_to_decoder']) if n==i and layout['computer_stage_kinds'][j]==kind)
if kind=='ffn':
for before,after in [('feed_forward.w3.','a.'),('feed_forward.w1.','b.'),('feed_forward.w2.data.','data.')]:
if tail.startswith(before):return 'computer.stages.'+str(source)+'.'+after+tail[len(before):]
raise ValueError('Unmapped machine FFN tensor: '+key)
assert tail.startswith('self_attn.');return 'computer.stages.'+str(source)+'.'+tail[len('self_attn.'):]
for index,child,prefix in [(layout['memory_index'],'ram.','ram.'),(layout['memory_index'],'files.','files.'),(layout.get('bridge_index'),'payload.','ram_bridge.'),(layout['analytics_index'],'payload.','analytics.')]:
if i==index and tail.startswith(child):return prefix+tail[len(child):]
raise ValueError('Unmapped decoder tensor: '+key)
|