diff options
Diffstat (limited to 'src/train.py')
| -rw-r--r-- | src/train.py | 9 |
1 files changed, 8 insertions, 1 deletions
diff --git a/src/train.py b/src/train.py index f0c63e9..975598d 100644 --- a/src/train.py +++ b/src/train.py @@ -133,7 +133,14 @@ def main() -> None: loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape( logits.shape[0], -1) * batch["y_mask"].float() loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1) - lam = cfg.halt_penalty if (cfg.halting and step >= cfg.halt_warmup_steps) else 0.0 + if cfg.halting and step >= cfg.halt_warmup_steps: + if step >= cfg.halt_ramp_end_steps: + lam = cfg.halt_penalty + else: + frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps) + lam = cfg.halt_penalty * frac + else: + lam = 0.0 halt = out["halt_steps"] penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0 loss = loss_tokens + penalty |