Update chess-mamba-vs-xformer/train_bygame.py

chess-mamba-vs-xformer/train_bygame.py

@@ -394,6 +394,9 @@ if init_from == 'scratch':
     print(f"saving checkpoint to {out_dir}\n")
     torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
 
+GAMES_SEEN_CHECKPOINTS = [12652800, 22275000, 11536000, 16250000, 18000000, 19690000, 22005050]
+TOKENS_SEEN_PADDED_CHECKPOINTS = [7798839804]
+
 t0 = time.time()
 while True:
     # Determine and set the learning rate for this iteration
@@ -402,7 +405,7 @@ while True:
         param_group['lr'] = lr
 
     # Evaluate the loss on train/val sets and write checkpoints
-    if master_process and ((iter_num % eval_interval == 0 and local_iter_num > 0) or abs(games_seen - 12652800) <= 151 or abs(games_seen - 22275000) <= 151 or abs(games_seen - 11536000) <= 151 or abs(games_seen - 16250000) <= 151 or abs(games_seen - 18000000) <= 151 or abs(games_seen - 19690000) <= 151 or abs(games_seen - 22005050) <= 151 or abs(tokens_seen_padded - 7798839804) <= 46238):
+    if master_process and ((iter_num % eval_interval == 0 and local_iter_num > 0) or any(abs(games_seen - checkpoint) <= 151 for checkpoint in GAMES_SEEN_CHECKPOINTS) or any(abs(tokens_seen_padded - checkpoint) <= 46238 for checkpoint in TOKENS_SEEN_PADDED_CHECKPOINTS)):
         torch.cuda.empty_cache()
         losses = estimate_loss()
         if init_from == 'anneal':
@@ -453,7 +456,7 @@ while True:
                 if losses['val'] < best_val_loss: # Temporary / only good after it's settled
                     best_val_loss = losses['val']
                     torch.save(checkpoint, os.path.join(out_dir, f'ckpt_{int(games_seen)}b.pt'))
-                elif current_nearest_multiple != last_crossed_multiple or abs(games_seen - 12652800) <= 151 or abs(games_seen - 22275000) <= 151 or abs(games_seen - 11536000) <= 151 or abs(games_seen - 16250000) <= 151 or abs(games_seen - 18000000) <= 151 or abs(games_seen - 19690000) <= 151 or abs(games_seen - 22005050) <= 151 or abs(tokens_seen_padded - 7798839804) <= 46238: # elif so we don't double up
+                elif current_nearest_multiple != last_crossed_multiple or any(abs(games_seen - checkpoint) <= 151 for checkpoint in GAMES_SEEN_CHECKPOINTS) or any(abs(tokens_seen_padded - checkpoint) <= 46238 for checkpoint in TOKENS_SEEN_PADDED_CHECKPOINTS): # elif so we don't double up
                     last_crossed_multiple = current_nearest_multiple
                     torch.save(checkpoint, os.path.join(out_dir, f'ckpt_{int(games_seen)}g_{tokens_seen_padded}t.pt'))