diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 09a0132..74aab93 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -295,7 +295,7 @@ if args.resume: with torch.no_grad(): W_out.copy_(_ck['wout'].to(dev)) if _ck.get('lnf') is not None and not isinstance(ln_f, nn.Identity): ln_f.load_state_dict(_ck['lnf']) start_step = int(_ck.get('step', 0)) - if _ck.get('bsimp') is not None: GOV['bsimp'] = float(_ck['bsimp']) + _bsimp_resume = _ck.get('bsimp') # applied after GOV exists (NameError fix: GOV defined below) print(f'[resume] loaded {args.resume} at step {start_step}', flush=True) if args.bf16: for _m in (tok, pos, blocks): @@ -458,6 +458,8 @@ def dFdtheta(zs, x, y, beta): SIG0 = None BGEN = torch.Generator().manual_seed(args.seed + 990) # separate RNG: sign flips must not shift the data stream GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0} +if args.resume and '_bsimp_resume' in dir() and _bsimp_resume is not None: + GOV['bsimp'] = float(_bsimp_resume) WSNAP = {'p': None, 'o': None} def _clone_state(sd): if torch.is_tensor(sd): return sd.clone() |