item 31 pre-registered (Nils's design): synthetic memory tokens — burst states → per-band-layer KV prefix via KVMemoryAdapter (in-place attention wrap, bit-exact disarmed, gated silent init); composed with frozen item-29 arm-1; +29 tf+fr in-flight note (30.5, FR term hurts)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -40,11 +40,25 @@ def main():
|
||||
help="load a lora_*_e*.pt loop-only band-LoRA checkpoint")
|
||||
ap.add_argument("--inner-iters", type=int, default=0, metavar="M",
|
||||
help="M in-place band iterations at the last pause")
|
||||
ap.add_argument("--kvmem", default=None, metavar="KVMEM_PT",
|
||||
help="KVMemoryAdapter checkpoint: burst states become "
|
||||
"band-layer KV prefix entries at generation")
|
||||
args = ap.parse_args()
|
||||
|
||||
model, tok = load_model(dtype=torch.bfloat16)
|
||||
tok.padding_side = "left"
|
||||
looper = BandLooper(model)
|
||||
kvmem = None
|
||||
if args.kvmem:
|
||||
from kv_memory import install, KVMemoryAdapter
|
||||
from loop_common import BAND
|
||||
install(model)
|
||||
sd = torch.load(args.kvmem, map_location="cuda")
|
||||
code = sd["trunk.1.weight"].shape[0]
|
||||
kvmem = KVMemoryAdapter(model, band=BAND, code=code).cuda()
|
||||
kvmem.load_state_dict(sd)
|
||||
kvmem.eval()
|
||||
print(f"kv-memory loaded: {args.kvmem} (code={code})", flush=True)
|
||||
if args.bandlora:
|
||||
from lora_band import inject_band_lora
|
||||
ck = torch.load(args.bandlora, map_location="cuda")
|
||||
@@ -86,7 +100,11 @@ def main():
|
||||
enc["attention_mask"], k, p,
|
||||
max_new_tokens=args.max_new,
|
||||
feedforward=args.feedforward,
|
||||
inner_iters=args.inner_iters)
|
||||
inner_iters=args.inner_iters,
|
||||
kvmem=kvmem)
|
||||
if kvmem is not None:
|
||||
from kv_memory import arm_memory
|
||||
arm_memory(None)
|
||||
for j, it in enumerate(chunk):
|
||||
txt = tok.decode(gen[j, enc["input_ids"].shape[1]:],
|
||||
skip_special_tokens=True)
|
||||
|
||||
Reference in New Issue
Block a user