diff options
Diffstat (limited to 'ep_run/bp_lm.py')
| -rw-r--r-- | ep_run/bp_lm.py | 17 |
1 files changed, 15 insertions, 2 deletions
diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py index 7fe0bdd..c01cd4f 100644 --- a/ep_run/bp_lm.py +++ b/ep_run/bp_lm.py @@ -33,14 +33,27 @@ def main(): ap.add_argument('--wd', type=float, default=1e-4) ap.add_argument('--pema', type=float, default=0.999) ap.add_argument('--qknorm', action='store_true') + 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('--log', type=int, default=200) ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt') cfg = ap.parse_args() torch.manual_seed(0) blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick') blk.qknorm = cfg.qknorm - opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd) - sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05) + if cfg.stdinit: # GPT-style: N(0,0.02), scaled residual projections + with torch.no_grad(): + for W in (blk.WQ, blk.WK, blk.WV, blk.fc, blk.Wh, blk.tok): + W.normal_(0, 0.02) + for W in (blk.WO, blk.pj): + W.normal_(0, 0.02 / (2 ** 0.5)) + blk.pos.normal_(0, 0.01) + 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) + else: + sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: 1.0) ema = [p.detach().clone() for p in blk.allp] best, t0 = float('inf'), time.time() for step in range(1, cfg.steps + 1): |