"""Execute standard Linear LoRA as an ARTI Formula, on the reading-QA task. Prepare the reading data and trained PEFT profiles with reading_qa.py first. This check transports those exact A/B weights; it does not independently retrain a different adapter and present its score as a framework ranking. The declared profile is unmerged Linear LoRA, dropout=0, no adapter bias, FP32 A/B and BF16 Qwen host, with identical mounting and alpha/r scaling. """ from __future__ import annotations import argparse import hashlib import json from pathlib import Path import torch from torch import nn from torch.nn import functional as F from peft import PeftModel from transformers import AutoTokenizer from arti import mechanisms as api import reading_qa as qa class TorchLoRA(nn.Module): def __init__(self, original, a, b, scaling): super().__init__() self.original = original self.A = nn.Parameter(a.detach().clone().to(original.weight.device)) self.B = nn.Parameter(b.detach().clone().to(original.weight.device)) self.register_buffer('gain', torch.tensor(scaling, dtype=torch.float32, device=original.weight.device)) def forward(self, x): base = self.original(x) with torch.autocast(device_type=x.device.type, enabled=False): update = F.linear(F.linear(x.float(), self.A), self.B) * self.gain value = base.float() + update return value.to(base.dtype) class FormulaLoRA(TorchLoRA): def __init__(self, original, a, b, scaling): super().__init__(original, a, b, scaling) declared = api.build_lora_program( input_dim=original.in_features, output_dim=original.out_features, rank=a.shape[0], source_ref='arti/formula-operand-bank@1', dtype='float32', contract_accumulation_dtype='float32', pointwise_accumulation_dtype='float32') # Exercise the public program codec, then execute that reconstructed IR. self.program = api.FormulaProgram.from_dict(json.loads(json.dumps(declared.to_dict()))) assert self.program.fingerprint == declared.fingerprint self.fabric = api.FormulaFabricV2(self.program) self.bank_bindings = {binding.name: binding for binding in self.program.bindings if isinstance(binding, api.BankBinding)} def forward(self, x): base = self.original(x) with torch.autocast(device_type=x.device.type, enabled=False): value = self.fabric( inputs={'x': x.float(), 'base': base.float(), 'lora.gain': self.gain}, banks={'lora.A': self.bank_bindings['lora.A'].bind(self.A), 'lora.B': self.bank_bindings['lora.B'].bind(self.B)}) return value.values[0].to(base.dtype) def matrices(model): values = {} for index in (26, 27): layer = model.model.layers[index].self_attn.v_proj if isinstance(layer, TorchLoRA): values[f'{index}.A'], values[f'{index}.B'] = layer.A, layer.B else: values[f'{index}.A'] = layer.lora_A['default'].weight values[f'{index}.B'] = layer.lora_B['default'].weight return values def snapshot(values, gradient=False): result = {} for name, value in values.items(): tensor = value.grad if gradient else value assert tensor is not None and torch.isfinite(tensor).all() result[name] = tensor.detach().cpu().clone() return result def difference(left, right): assert left.keys() == right.keys() return max((left[key] - right[key]).abs().max().item() for key in left) def continuation(model, batches): parameters = list(matrices(model).values()) optimizer = torch.optim.AdamW(parameters, lr=.001, weight_decay=.01, fused=True) model.train() losses, gradients, values = [], [], [] for batch in batches: optimizer.zero_grad(set_to_none=True) loss = model(**batch, use_cache=False).loss loss.backward() losses.append(loss.detach().item()) gradients.append(snapshot(matrices(model), gradient=True)) torch.nn.utils.clip_grad_norm_(parameters, 1.) optimizer.step() values.append(snapshot(matrices(model))) return {'losses': losses, 'gradients': gradients, 'parameters': values} def check(work, output, base_path, language, seed, qualification_only=False): target = output / f'{language}-{seed}' if (target / 'result.json').exists(): completed = json.loads((target / 'result.json').read_text(encoding='utf8')) assert completed['identical_peft_answers'] == completed['questions'] == 256 assert all(value == 0 for checks in completed['qualification'].values() for value in checks.values()) print(f'Completed: {language}, seed {seed}', flush=True) return target.mkdir(parents=True, exist_ok=True) qa.seed_all(seed) torch.backends.cuda.matmul.allow_tf32 = False model = PeftModel.from_pretrained(qa.load_base(base_path), work / language / f'peft-{seed}' / 'adapter', is_trainable=True) host = model.get_base_model() host.config.use_cache = False tokenizer = AutoTokenizer.from_pretrained(base_path, local_files_only=True) profile = model.peft_config['default'] assert profile.r == profile.lora_alpha == 32 and profile.lora_dropout == 0 assert profile.bias == 'none' and not profile.use_dora and not profile.use_rslora initial = snapshot(matrices(host)) assert all(value.dtype == torch.float32 for value in initial.values()) assert sum(value.numel() for value in initial.values()) == 131072 source_rows = json.loads((work / language / 'train.json').read_text(encoding='utf8'))[:24] encoded = [qa.encode_final_answer(row['id'], row['messages'], tokenizer, template_kwargs=qa.TEMPLATE) for row in source_rows] batches = [qa.collate_chat(encoded[begin:begin+8], pad_token_id=tokenizer.pad_token_id, pad_to_multiple_of=8) for begin in (0, 8, 16)] batches = [{key: value.to('cuda') for key, value in batch.items()} for batch in batches] print(f'Checking the same LoRA: {language}, seed {seed}', flush=True) controls = {'peft': continuation(host, batches)} originals = {} gains = {} for index in (26, 27): layer = host.model.layers[index].self_attn.v_proj dropout = layer.lora_dropout['default'] assert isinstance(dropout, nn.Identity) or (isinstance(dropout, nn.Dropout) and dropout.p == 0) assert not layer.merged and not layer.lora_variant originals[index], gains[index] = layer.base_layer, layer.scaling['default'] for name, implementation in [('torch', TorchLoRA), ('arti', FormulaLoRA)]: for index in (26, 27): host.model.layers[index].self_attn.v_proj = implementation( originals[index], initial[f'{index}.A'], initial[f'{index}.B'], gains[index]) controls[name] = continuation(host, batches) qualified = {} for name in ('torch', 'arti'): qualified[name] = { 'loss_max_abs': max(abs(a-b) for a, b in zip(controls['peft']['losses'], controls[name]['losses'])), 'gradient_max_abs': max(difference(a,b) for a,b in zip(controls['peft']['gradients'], controls[name]['gradients'])), 'parameter_max_abs': max(difference(a,b) for a,b in zip(controls['peft']['parameters'], controls[name]['parameters'])), } result = {'language': language, 'seed': seed, 'base': qa.BASE, 'revision': qa.REVISION, 'scope': 'same trained LoRA weights expressed in PEFT, handwritten PyTorch and native ARTI Formula', 'trainable_parameters': 131072, 'continuation_steps': 3, 'training_ids': [row['id'] for row in source_rows], 'losses': {name: control['losses'] for name, control in controls.items()}, 'qualification': qualified, 'program_codec_roundtrip': True, 'source_adapter_sha256': qa.digest((work/language/f'peft-{seed}'/'adapter/adapter_model.safetensors').read_bytes()), 'source_sha256': qa.digest(Path(__file__).read_bytes())} result['finite_nonzero_gradients'] = all( any(value.abs().max().item() > 0 for value in gradients.values()) for control in controls.values() for gradients in control['gradients']) qa.write(target / 'qualification.json', result) # Exact deltas are reported and required; failed attempts remain on disk. assert all(value == 0 for checks in qualified.values() for value in checks.values()), qualified assert result['finite_nonzero_gradients'] if qualification_only: print(json.dumps({'language': language, 'seed': seed, 'qualification': qualified}), flush=True) return for index in (26, 27): layer = host.model.layers[index].self_attn.v_proj with torch.no_grad(): layer.A.copy_(initial[f'{index}.A'].to(layer.A.device)) layer.B.copy_(initial[f'{index}.B'].to(layer.B.device)) qa.write(target / f'layer-{index}-program.json', layer.program.to_dict()) rows = json.loads((work/language/'heldout.json').read_text(encoding='utf8')) answers, _ = qa.generate(host, tokenizer, rows) scorer = qa.score_module(work, language) for answer in answers: answer['em'], answer['f1'] = qa.per_answer(scorer, language, answer['references'], answer['prediction']) qa.write(target / 'predictions.json', answers) expected = json.loads((work/language/f'peft-{seed}'/'predictions.json').read_text(encoding='utf8')) assert [x['id'] for x in answers] == [x['id'] for x in expected] identical = sum(a['prediction'] == b['prediction'] for a,b in zip(answers,expected)) result.update({'questions': len(answers), 'identical_peft_answers': identical, 'em': 100*sum(x['em'] for x in answers)/len(answers), 'f1': 100*sum(x['f1'] for x in answers)/len(answers)}) qa.write(target / 'comparison.json', result) assert identical == len(answers), result qa.write(target / 'result.json', result) print(json.dumps({'language':language,'seed':seed,'identical_answers':identical, 'em':result['em'],'f1':result['f1']}),flush=True) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--work', type=Path, default=Path('reading-qa-work')) parser.add_argument('--output', type=Path, default=Path('lora-formula-check')) parser.add_argument('--language', choices=['squad','cmrc']) parser.add_argument('--seed', type=int, choices=qa.SEEDS) parser.add_argument('--qualification-only', action='store_true') args = parser.parse_args() from huggingface_hub import snapshot_download base = Path(snapshot_download(qa.BASE, revision=qa.REVISION, cache_dir=args.work/'hub-cache', local_files_only=True)) for language in [args.language] if args.language else ['squad','cmrc']: for seed in [args.seed] if args.seed else qa.SEEDS: check(args.work,args.output,base,language,seed,args.qualification_only) qa.cleanup() if __name__ == '__main__': main()