Winning without optimization / SFT, by modeling (outcome, moves) pairs

benchmark.py

@@ -6,45 +6,61 @@
 import random
 
 
-model = load_from_checkpoint()
-model.eval()
-model.to(device)
-
-with torch.no_grad():
-    counts = {"player_2": 0, "player_1": 0, "draw": 0, "invalid": 0}
-    for _ in range(1000):
-        board = np.zeros((3, 3), dtype=int)
-        player = 1
-        winner = None
-        moves = [START]
-        while winner is None and not board_full(board):
-            if player == 1:
-                x = torch.tensor(moves, dtype=torch.long, device=device)[None, ...]
-                y = model.generate(x, max_new_tokens=1, temperature=1.0, top_k=3)
-                y = y[0][-1].item()
-
-                if y not in set(range(9)) or y in moves:
-                    print(f"invalid move: {y} moves: {moves}")
-                    winner = None
-                    break
-
-                i, j = divmod(y, 3)
+
+def main(incl_winner: bool = False, target_score: int = 1):
+    """
+    incl_winner: Whether to include the winner in the input sequence just after the START token
+    target_score: Used only when incl_winner is True. 0 for player 2, 1 for player 1, 2 for draw
+    """
+
+    model = load_from_checkpoint()
+    model.eval()
+    model.to(device)
+
+    with torch.no_grad():
+        counts = {"player_2": 0, "player_1": 0, "draw": 0, "invalid": 0}
+        for _ in range(1000):
+            board = np.zeros((3, 3), dtype=int)
+            player = 1
+            winner = None
+
+            moves = []
+            if incl_winner:
+                assert target_score in [0, 1, 2]
+
+            while winner is None and not board_full(board):
+                if player == 1:
+                    x_list = [START, target_score] + moves if incl_winner else [START] + moves
+                    x = torch.tensor(x_list, dtype=torch.long, device=device)[None, ...]
+                    y = model.generate(x, max_new_tokens=1, temperature=1.0, top_k=3)
+                    y = y[0][-1].item()
+
+                    if y not in set(range(9)) or y in moves:
+                        print(f"invalid move: {y} moves: {moves}")
+                        winner = None
+                        break
+
+                    i, j = divmod(y, 3)
+                else:
+                    i, j = random.choice(get_valid_moves(board))
+
+                moves.append(i * 3 + j)
+                board[i][j] = player
+                player *= -1
+                winner = check_winner(board)
+
+            if winner == 1:
+                counts["player_1"] += 1
+            elif winner == -1:
+                counts["player_2"] += 1
+            elif board_full(board):
+                counts["draw"] += 1
             else:
-                i, j = random.choice(get_valid_moves(board))
-
-            moves.append(i * 3 + j)
-            board[i][j] = player
-            player *= -1
-            winner = check_winner(board)
-
-        if winner == 1:
-            counts["player_1"] += 1
-        elif winner == -1:
-            counts["player_2"] += 1
-        elif board_full(board):
-            counts["draw"] += 1
-        else:
-            counts["invalid"] += 1
-
-    print(counts)
-    print(counts["player_1"] / sum(counts.values()))
+                counts["invalid"] += 1
+
+        print(counts)
+        print(counts["player_1"] / sum(counts.values()))
+
+import typer
+if __name__ == "__main__":
+    typer.run(main)

generate_data.py

@@ -53,29 +53,52 @@ def seq_to_board(seq):
     return board
 
 
-def save_data(trajectories):
+w_map = {-1: 0, 1: 1, None: 2}
+
+def save_data(trajectories, incl_winner=False):
     outcomes = defaultdict(int)
     for b, s, w in trajectories:
+        #print(b, s, w)
         outcomes[w] += 1
 
     print(outcomes)
 
-    data = np.full((len(trajectories), SEQ_LENGTH), PAD, dtype=np.int16)
-    for i in range(len(trajectories)):
-        row = [START] + trajectories[i][1]
-        if i < 10:
-            print(row)
-        data[i, : len(row)] = row
+    if not incl_winner:
+        data = np.full((len(trajectories), SEQ_LENGTH), PAD, dtype=np.int16)
+        for i, (b, s, w) in enumerate(trajectories):
+            # start with the START token, then sequence
+            row = [START] + s
+            if i < 10:
+                print(row)
+            data[i, : len(row)] = row
+    else:
+        data=np.full((len(trajectories), SEQ_LENGTH+1), PAD, dtype=np.int16)
+        for i, (b, s, w) in enumerate(trajectories):
+            # start with the START token, then winner, then sequence
+            row = [START, w_map(w)] + s
+            if i < 10:
+                print(row)
+            data[i, : len(row)] = row
 
     np.random.shuffle(data)
 
     np.save("data/train.npy", data)
 
 
-if __name__ == "__main__":
+# Use typer instead
+import typer
+app = typer.Typer()
+
+def main(optimal: bool = False, incl_winner: bool = False):
     board = np.zeros((3, 3), dtype=int)
-    # trajectories = all_trajectories(board, [], 1)
+    if optimal:
+        trajectories = all_optimal_trajectories(board, [], 1)
+    else:
+        trajectories = all_trajectories(board, [], 1)
     # trajectories = all_optimal_trajectories(board, [], 1, {-1, 1})
-    trajectories = all_optimal_trajectories(board, [], 1)
+    # trajectories = all_optimal_trajectories(board, [], 1)
     print(len(trajectories))
-    save_data(trajectories)
+    save_data(trajectories, incl_winner)
+
+if __name__ == "__main__":
+    typer.run(main)

requirements.txt

@@ -0,0 +1,4 @@
+torch
+typer
+wandb
+chardet

setup.py

@@ -14,17 +14,18 @@
 torch.backends.cudnn.allow_tf32 = True
 
 
-def init_model():
+def init_model(incl_winner: bool = False):
     print("Initializing a new model from scratch")
     config = GPTConfig(
-        block_size=SEQ_LENGTH,
+        block_size=SEQ_LENGTH + 1 if incl_winner else SEQ_LENGTH,
         vocab_size=VOCAB_SIZE,
         n_layer=1,
         n_head=1,
         n_embd=14,
         dropout=0.0,
         bias=False,
     )
+    print("config:", config)
     return GPT(config)
 
 

train.py

@@ -1,5 +1,6 @@
 import os
 import time
+import typer
 
 import numpy as np
 import torch
@@ -39,54 +40,61 @@ def get_batch():
     return x, y
 
 
-iter_num = 0
 
-model = init_model()
-model.to(device)
-model.train()
 
-optimizer = model.configure_optimizers(
-    weight_decay, learning_rate, (beta1, beta2), device
-)
 
+def main(incl_winner: bool = False):
+    iter_num = 0
+    model = init_model(incl_winner=incl_winner)
+    model.to(device)
+    model.train()
 
-if wandb_log:
-    import wandb
-
-    wandb.init(project=wandb_project)
+    optimizer = model.configure_optimizers(
+        weight_decay, learning_rate, (beta1, beta2), device
+    )
 
-X, Y = get_batch()
-t0 = time.time()
-while iter_num < max_iters:
-    if iter_num > 0 and iter_num % save_interval == 0:
-        save_checkpoint(model)
 
-    logits, loss = model(X, Y)
+    if wandb_log:
+        import wandb
 
-    loss.backward()
+        wandb.init(project=wandb_project)
 
     X, Y = get_batch()
+    t0 = time.time()
+    print("AAAAAAAAAAAAAAAA")
+    while iter_num < max_iters:
+        if iter_num > 0 and iter_num % save_interval == 0:
+            save_checkpoint(model)
 
-    if grad_clip != 0.0:
-        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+        logits, loss = model(X, Y)
 
-    optimizer.step()
+        loss.backward()
 
-    optimizer.zero_grad(set_to_none=True)
+        X, Y = get_batch()
 
-    t1 = time.time()
-    dt = t1 - t0
-    t0 = t1
-    lossf = loss.item()
-    print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms")
+        if grad_clip != 0.0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
 
-    if wandb_log:
-        wandb.log(
-            {
-                "iter": iter_num,
-                "train/loss": lossf,
-                "lr": learning_rate,
-            }
-        )
-
-    iter_num += 1
+        optimizer.step()
+
+        optimizer.zero_grad(set_to_none=True)
+
+        t1 = time.time()
+        dt = t1 - t0
+        t0 = t1
+        lossf = loss.item()
+        print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms")
+
+        if wandb_log:
+            wandb.log(
+                {
+                    "iter": iter_num,
+                    "train/loss": lossf,
+                    "lr": learning_rate,
+                }
+            )
+
+        iter_num += 1
+
+if __name__ == "__main__":
+    typer.run(main)