From 7467b5c313393e9c01b5d877f45e680e0c35213f Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 29 Jul 2026 04:46:31 -0500 Subject: =?UTF-8?q?RESULT=2069:=20=E9=A5=B1=E5=92=8C=E7=9B=B2=E5=8C=BA?= =?UTF-8?q?=E8=A2=AB=E8=87=AA=E5=B7=B1=E7=9A=84=E6=B2=BB=E7=96=97=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E8=AF=81=E4=BC=AA(cap30/50=E5=85=B3=E9=97=AD0.0000,?= =?UTF-8?q?=E8=AF=81=E4=B9=A6cap=3D1=E7=88=86=E6=A8=A1=E5=9E=8B);=20cap?= =?UTF-8?q?=E4=BB=B7=E6=A0=BC=E2=89=880(=E7=94=A8=E6=88=B7=E6=8B=85?= =?UTF-8?q?=E5=BF=A7=E5=AE=9A=E4=BB=B7);=20=E5=AD=98=E6=B4=BB=E7=BA=A6?= =?UTF-8?q?=E6=9D=9F=3D=E9=A1=B6=E5=8D=8A+attn+1/=CE=B2+11=E9=A1=B9?= =?UTF-8?q?=E5=85=8D=E7=96=AB;=20dgain=E4=BD=8D=E7=A7=BB=E8=87=82=E5=9C=A8?= =?UTF-8?q?=E9=A3=9E(cos0.83=E7=9A=84=E4=BB=A3=E4=BB=B7=E7=94=B1CE?= =?UTF-8?q?=E8=A3=81=E5=86=B3)+mixffn=E8=A1=A5=E5=88=86=E8=A7=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_bp_train.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) (limited to 'ep_run/casc_bp_train.py') diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 42acf16..4ecfee9 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -13,6 +13,7 @@ ap.add_argument('--B', type=int, default=24); ap.add_argument('--steps', type=in ap.add_argument('--lr', type=float, default=3e-4); ap.add_argument('--warmup', type=int, default=200) ap.add_argument('--seed', type=int, default=0) ap.add_argument('--save_every', type=int, default=500); ap.add_argument('--log', type=int, default=200) +ap.add_argument('--logit_cap', type=float, default=0.0) # >0: Gemma-2-style attn logit softcap ap.add_argument('--watch_every', type=int, default=2000) # wandb-only telemetry: weight/act RMS ap.add_argument('--wandb', default='auto') # ON BY DEFAULT; 'auto' = per-regime project (ept-fineweb-72m / ept-tinystories-42m); --wandb '' to disable ap.add_argument('--wandb_run', default='') @@ -114,7 +115,13 @@ class Olmo2Attn(nn.Module): q = self.rope(q.view(B, T, self.H, self.hd).transpose(1, 2)) k = self.rope(k.view(B, T, self.H, self.hd).transpose(1, 2)) v = v.view(B, T, self.H, self.hd).transpose(1, 2) - y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + if args.logit_cap > 0: + lg = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5) + lg = args.logit_cap * torch.tanh(lg / args.logit_cap) + cm = torch.ones(T, T, dtype=torch.bool, device=x.device).tril() + y = lg.masked_fill(~cm, float('-inf')).softmax(-1) @ v + else: + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) return self.proj(y.transpose(1, 2).contiguous().view(B, T, C)) class Olmo2Block(nn.Module): -- cgit v1.2.3