"""Measure every attention head/MLP under ablation, donor and matched controls.""" import hashlib import importlib.util import json import os from pathlib import Path import re import time import torch from peft import PeftModel,get_peft_model_state_dict from safetensors.torch import load_file,save_file from transformers import AutoModelForCausalLM,AutoTokenizer from huggingface_hub import HfApi,hf_hub_download CACHE='itsalissonsilva/portfolio-experiment-cache' PREFIX='component_circuits_v1' OUT=Path('/tmp/component_circuits');OUT.mkdir(exist_ok=True) def emit(**value):print(json.dumps(value),flush=True) def summarize(rows,outputs): result={} for role in ['private','public','utility']: selected=[o for r,o in zip(rows,outputs,strict=True) if r['analysis_role']==role] result[role]=dict(n=len(selected),passed=sum(o['passed'] for o in selected), leaks=sum(o['leaked'] for o in selected),valid=sum(o['valid'] for o in selected)) return result def main(): began=time.monotonic();api=HfApi();assert api.dataset_info(CACHE).private revision=os.environ['COMPONENT_INPUT_REVISION'] def get(file):return Path(hf_hub_download(CACHE,PREFIX+'/inputs/'+file,repo_type='dataset',revision=revision)) manifest=json.loads(get('manifest.json').read_text());file=get('rows.json') assert hashlib.sha256(file.read_bytes()).hexdigest()==manifest['rows_sha256'];rows=json.loads(file.read_text()) assert all(not r['private_canary'] for r in rows['public']) assert all(r['kind']=='unrelated' for r in rows['utility']) helper_file=Path(hf_hub_download(manifest['adapter'],'train_restriction_pilot.py',revision=manifest['adapter_revision'])) spec=importlib.util.spec_from_file_location('restriction_helper',helper_file);helper=importlib.util.module_from_spec(spec);spec.loader.exec_module(helper) assert torch.cuda.is_available() and torch.cuda.is_bf16_supported();torch.set_num_threads(4) torch.manual_seed(manifest['seed']) tokenizer=AutoTokenizer.from_pretrained(manifest['model'],revision=manifest['model_revision']);tokenizer.padding_side='left' base=AutoModelForCausalLM.from_pretrained(manifest['model'],revision=manifest['model_revision'],torch_dtype=torch.bfloat16,attn_implementation='eager').cuda().eval() model=PeftModel.from_pretrained(base,manifest['adapter'],subfolder=manifest['adapter_subfolder'],revision=manifest['adapter_revision']).eval() for parameter in model.parameters():parameter.requires_grad_(False) saved=load_file(hf_hub_download(manifest['adapter'],manifest['adapter_subfolder']+'/adapter_model.safetensors',revision=manifest['adapter_revision'])) actual=get_peft_model_state_dict(model);assert set(actual)==set(saved) assert all(torch.equal(actual[k].float().cpu(),saved[k].float()) for k in saved),'Adapter changed before intervention' layers=[module for name,module in model.named_modules() if re.search(r'\.layers\.\d+$',name)];assert len(layers)==26 config=getattr(base.config,'text_config',base.config);heads=config.num_attention_heads;head_dim=config.head_dim assert heads*head_dim==layers[0].self_attn.o_proj.in_features combined=rows['private']+rows['public']+rows['utility'] def inputs(chosen): result=tokenizer([helper.render_prompt(tokenizer,r) for r in chosen],padding=True,return_tensors='pt').to('cuda') assert result['input_ids'].shape[1]<4096 return result @torch.inference_mode() def plain(chosen): outputs=[] for start in range(0,len(chosen),8): batch=chosen[start:start+8];encoded=inputs(batch) result=model.generate(**encoded,max_new_tokens=80,do_sample=False,use_cache=True,pad_token_id=tokenizer.pad_token_id) raw=tokenizer.batch_decode(result[:,encoded['input_ids'].shape[1]:],skip_special_tokens=True) outputs.extend(helper.grading(r,s) for r,s in zip(batch,raw,strict=True)) return outputs @torch.inference_mode() def capture(chosen): values={};handles=[] for index,layer in enumerate(layers): def head_hook(module,args,index=index): if args[0].shape[1]>1:values['head.'+str(index)]=args[0][:,-1].detach().cpu().clone() def mlp_hook(module,args,output,index=index): assert isinstance(output,torch.Tensor) if output.shape[1]>1:values['mlp.'+str(index)]=output[:,-1].detach().cpu().clone() handles.append(layer.self_attn.o_proj.register_forward_pre_hook(head_hook)) handles.append(layer.mlp.register_forward_hook(mlp_hook)) # Match the full generation configuration, including cache allocation; # capture only prefill tensors, without changing any activation. try:model.generate(**inputs(chosen),max_new_tokens=80,do_sample=False,use_cache=True,pad_token_id=tokenizer.pad_token_id) finally: for handle in handles:handle.remove() assert len(values)==52 and all(v.shape[0]==8 for v in values.values()) return values donors={role:capture(rows[role]) for role in ['private','public']} baseline=plain(combined);private_baseline=baseline[:8] report=dict(manifest=manifest,source_revision=revision,architecture=dict(layers=26,heads=heads,head_dim=head_dim,hidden_size=config.hidden_size), baseline=dict(summary=summarize(combined,baseline),outputs=baseline),components=[], self_restoration_exact=0,hook_removal_exact=0,all_controls_passed=False,adapter_tensors_exact=len(saved)) def persist(): path=OUT/'results.json';path.write_text(json.dumps(report,indent=2,allow_nan=False)+'\n') api.upload_file(repo_id=CACHE,repo_type='dataset',path_or_fileobj=str(path),path_in_repo=PREFIX+'/results.json',commit_message='Preserve measured synthetic head/MLP intervention controls privately; no complete-circuit claim') persist() @torch.inference_mode() def intervene(chosen,layer_index,component,head,mode): results=[];norm_errors=[] module=layers[layer_index].self_attn.o_proj if component=='head' else layers[layer_index].mlp start=head*head_dim if component=='head' else 0 stop=start+head_dim if component=='head' else config.hidden_size for offset in range(0,len(chosen),8): batch=chosen[offset:offset+8] def change(hidden): if hidden.shape[1]<=1:return hidden old=hidden[:,-1,start:stop];desired=torch.zeros_like(old) if mode in ['public','random_public','self']: assert len(chosen)==8 and offset==0 role='private' if mode=='self' else 'public' desired=donors[role][component+'.'+str(layer_index)][:,start:stop].to(old.device,old.dtype) delta=desired.float()-old.float() if mode.startswith('random'): noise=[] for i in range(len(batch)): generator=torch.Generator().manual_seed(manifest['seed']+layer_index*10000+(head+1)*100+offset+i) noise.append(torch.randn(old.shape[-1],generator=generator)) noise=torch.stack(noise).to(old.device) wanted=delta.norm(dim=-1,keepdim=True) noise=noise/noise.norm(dim=-1,keepdim=True).clamp_min(1e-20)*wanted error=((noise.norm(dim=-1)-wanted[:,0]).abs()/wanted[:,0].clamp_min(1e-20)).max().item() assert error<1e-5;norm_errors.append(error);desired=(old.float()+noise).to(old.dtype) changed=hidden.clone();changed[:,-1,start:stop]=desired return changed if component=='head': def hook(module,args):return (change(args[0]),*args[1:]) handle=module.register_forward_pre_hook(hook) else: def hook(module,args,output):return change(output) handle=module.register_forward_hook(hook) try:outputs=plain(batch) finally:handle.remove() results.extend(outputs) return dict(summary=summarize(chosen,results),outputs=results,max_float32_norm_match_error=max(norm_errors,default=0)) for layer in range(26): for component,head in [('head',h) for h in range(heads)]+[('mlp',-1)]: entry=dict(layer=layer,component=component,head=head,conditions={}) for mode in ['ablate','random_ablate','public','random_public','self']: chosen=combined if mode in ['ablate','random_ablate'] else rows['private'] result=intervene(chosen,layer,component,head,mode) if mode=='self': matches=sum(a['raw']==b['raw'] for a,b in zip(result['outputs'],private_baseline,strict=True)) assert matches==8,'Component self-restoration failed; stop causal claims' report['self_restoration_exact']+=matches;result['raw_matches_baseline']=matches entry['conditions'][mode]=result report['components'].append(entry) emit(stage='component_measured',layer=layer,component=component,head=head, ablation=entry['conditions']['ablate']['summary'],donor=entry['conditions']['public']['summary']) restored=plain(rows['private']);assert [o['raw'] for o in restored]==[o['raw'] for o in private_baseline] report['hook_removal_exact']+=8;persist() assert report['self_restoration_exact']==26*(heads+1)*8 and report['hook_removal_exact']==26*8 report['all_controls_passed']=True;report['elapsed_seconds']=time.monotonic()-began;persist() save_file({role+'.'+key:value.contiguous() for role,values in donors.items() for key,value in values.items()},str(OUT/'captured_donors.safetensors')) api.upload_file(repo_id=CACHE,repo_type='dataset',path_or_fileobj=str(OUT/'captured_donors.safetensors'),path_in_repo=PREFIX+'/captured_donors.safetensors',commit_message='Preserve actual generation-prefill component donors for reproducible synthetic intervention checks') emit(stage='complete',components=len(report['components']),self_restoration=report['self_restoration_exact'],hook_removal=report['hook_removal_exact']) if __name__=='__main__':main()