diff options
Diffstat (limited to 'ep_run/dp_ep.py')
| -rw-r--r-- | ep_run/dp_ep.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/ep_run/dp_ep.py b/ep_run/dp_ep.py index 0d76911..5c25c0b 100644 --- a/ep_run/dp_ep.py +++ b/ep_run/dp_ep.py @@ -23,7 +23,7 @@ def main(): ap.add_argument('--hr', type=float, default=0.02) ap.add_argument('--holofast', action='store_true') ap.add_argument('--sdpa', action='store_true') - ap.add_argument('--log', type=int, default=50) + ap.add_argument('--log_every', type=int, default=50) cfg = ap.parse_args() dist.init_process_group('nccl') @@ -58,7 +58,7 @@ def main(): g /= world p.grad = g opt.step(); opt.zero_grad(set_to_none=True) - if step % cfg.log == 0 and rank == 0: + if step % cfg.log_every == 0 and rank == 0: v = L.evaluate(blk, 150, 0.1, nb=2) print(f"[dp{world}] step {step}/{cfg.steps} | val {v:.4f} | res {res:.1e} | " f"{step / (time.time() - t0):.3f} it/s(x{world}B)", flush=True) |