diff options
Diffstat (limited to 'ep_run/bp_lm.py')
| -rw-r--r-- | ep_run/bp_lm.py | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py index c01cd4f..1bd12f8 100644 --- a/ep_run/bp_lm.py +++ b/ep_run/bp_lm.py @@ -13,7 +13,7 @@ def fwd(blk, idx): h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b) h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b) h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb) - return h @ blk.Wh + return h @ (blk.tok.t() if getattr(blk, 'tie', False) else blk.Wh) def evaluate(blk, nb=8, B=32): @@ -36,6 +36,7 @@ def main(): ap.add_argument('--stdinit', action='store_true') # standard transformer init (EQBlock's is tuned for relaxation) ap.add_argument('--beta2', type=float, default=0.999) ap.add_argument('--sched', choices=['cos', 'const'], default='cos') + ap.add_argument('--tie', action='store_true') # tok/Wh weight tying (standard small-LM trick) ap.add_argument('--log', type=int, default=200) ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt') cfg = ap.parse_args() @@ -49,6 +50,13 @@ def main(): for W in (blk.WO, blk.pj): W.normal_(0, 0.02 / (2 ** 0.5)) blk.pos.normal_(0, 0.01) + if cfg.tie: # tie head to embedding: Wh := tok^T, single parameter + with torch.no_grad(): + blk.tok.copy_(0.5 * (blk.tok + blk.Wh.t())) + blk.Wh = None # fwd() will use tok.t() when tie is on + blk.tie = True + blk.allp = [p for p in blk.allp if p is not blk.Wh] + blk.allp = blk.block + [] # block already contains tok; Wh dropped opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd, betas=(0.9, cfg.beta2)) if cfg.sched == 'cos': sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05) |