Upload 37 files

chess-gpt-eval/gpt_query.py

@@ -0,0 +1,265 @@
+import openai
+import tiktoken
+import json
+import os
+
+# import replicate
+
+# for hugging face inference endpoints for codellama
+import requests
+
+from typing import Optional
+
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)  # for exponential backoff
+
+# system message is used in openai_request()
+system_message = """Provide the next move in the chess game. Only provide the move, no move numbers."""
+
+# dollars per 1k tokens, per openai.com/pricing
+pricing_dict = {
+    "gpt-4": 0.03,
+    "gpt-4-0301": 0.03,
+    "gpt-4-0613": 0.03,
+    "gpt-3.5-turbo": 0.0015,
+    "gpt-3.5-turbo-0301": 0.0015,
+    "gpt-3.5-turbo-0613": 0.0015,
+    "gpt-3.5-turbo-16k": 0.003,
+    "babbage": 0.0005,
+    "gpt-3.5-turbo-instruct": 0.0015,
+}
+
+MAX_TOKENS = 10
+
+completion_models = [
+    "gpt-3.5-turbo-instruct",
+    "babbage",
+    "davinci",
+]
+
+
+# tenacity is to handle anytime a request fails
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def get_gpt_response(
+    prompt: str, model: str = "gpt-4", temperature: float = 0.0
+) -> Optional[str]:
+    try:
+        messages = []
+        # system message is used in openai_request()
+        # system_message_dict = {
+        #     "role": "system",
+        #     "content": system_message,
+        # }
+        initial_message = {"role": "user", "content": prompt}
+        messages.append(initial_message)
+
+        record_messages(messages, model)
+
+        # num_tokens = count_all_tokens(model, messages)
+        # prompt_cost = get_prompt_cost(model, num_tokens)
+        # print("prompt cost in $:", prompt_cost)
+
+        if model in completion_models:
+            response = get_completions_response(model, messages, temperature)
+        elif model.startswith("gpt"):
+            response = openai_chat_completion_request(model, messages, temperature)
+        elif model.startswith("openrouter"):
+            response = openrouter_request(model, messages, temperature)
+        elif model.startswith("huggingface"):
+            response = hugging_face_request(model, messages, temperature)
+        elif model.startswith("replicate"):
+            response = replicate_request(model, messages, temperature)
+        else:
+            raise Exception("Invalid model name")
+
+        # response_cost = get_response_cost(model, count_tokens(model, response))
+        # print("response cost in $:", response_cost)
+
+        messages.append({"role": "assistant", "content": response})
+        record_messages(messages, model)
+
+        return response
+    except Exception as e:
+        print(f"Error while getting GPT response: {e}")
+        return None
+
+
+def openai_chat_completion_request(
+    model: str, messages: list[dict], temperature: float
+) -> str:
+    system_message_dict = {
+        "role": "system",
+        "content": system_message,
+    }
+    messages.append(system_message_dict)
+    completion = openai.ChatCompletion.create(
+        model=model,
+        temperature=temperature,
+        messages=messages,
+    )
+    response = completion.choices[0].message.content
+    return response
+
+
+def openrouter_request(model: str, messages: list[dict], temperature: float) -> str:
+    if temperature == 0:
+        temperature = 0.001
+
+    with open("gpt_inputs/openrouter_api_key.txt", "r") as f:
+        openai.api_key = f.read().strip()
+
+    openai.api_base = "https://openrouter.ai/api/v1"
+    OPENROUTER_REFERRER = "https://github.com/adamkarvonen/nanoGPT"
+
+    model = model.replace("openrouter/", "")
+
+    completion = openai.ChatCompletion.create(
+        model=model,
+        headers={"HTTP-Referer": OPENROUTER_REFERRER},
+        messages=messages,
+        temperature=temperature,
+        max_tokens=MAX_TOKENS,
+    )
+    response = completion.choices[0].message.content
+    return response
+
+
+def replicate_request(model: str, messages: list[dict], temperature: float) -> str:
+    if temperature == 0:
+        temperature = 0.001
+
+    with open("gpt_inputs/replicate_api_key.txt", "r") as f:
+        api_key = f.read().strip()
+    os.environ["REPLICATE_API_TOKEN"] = api_key
+
+    model = model.replace("replicate/", "")
+
+    messages = translate_to_string_input(messages)
+
+    output = replicate.run(
+        model,
+        input={
+            "prompt": messages,
+            "max_new_tokens": MAX_TOKENS,
+            "temperature": temperature,
+        },
+    )
+
+    # The meta/llama-2-7b model can stream output as it's running.
+    response = ""
+    # The predict method returns an iterator, and you can iterate over that output.
+    for item in output:
+        # https://replicate.com/meta/llama-2-7b/versions/527827021d8756c7ab79fde0abbfaac885c37a3ed5fe23c7465093f0878d55ef/api#output-schema
+        response += item
+
+    return response
+
+
+def hugging_face_request(model: str, messages: list[dict], temperature: float) -> str:
+    def query(payload):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        return response.json()
+
+    messages = translate_to_string_input(messages)
+    API_URL = "https://xxxxxxxx.us-east-1.aws.endpoints.huggingface.cloud"
+    headers = {
+        "Authorization": "Bearer xxxxx",
+        "Content-Type": "application/json",
+    }
+
+    if temperature == 0:
+        temperature = 0.001
+
+    output = query(
+        {
+            "inputs": messages,
+            "parameters": {"temperature": temperature, "max_new_tokens": MAX_TOKENS},
+        }
+    )
+
+    return output[0]["generated_text"]
+
+
+def translate_to_string_input(
+    openai_messages: list[dict], roles_included: bool = False
+):
+    # Translate from OpenAI's dict to a single string input
+    messages = []
+    for message in openai_messages:
+        if roles_included:
+            messages.append(message["role"] + ": ")
+        messages.append(message["content"])
+    if roles_included:
+        messages.append("assistant: ")
+    return "\n".join(messages)
+
+
+# for gpt-3 models and instruct models
+def get_completions_response(
+    model: str,
+    messages: list[dict] | str,
+    temperature: float,
+    max_tokens: int = MAX_TOKENS,
+) -> str:
+    if not isinstance(messages, str):
+        prompt = translate_to_string_input(messages, roles_included=False)
+    else:
+        prompt = messages
+
+    completion = openai.Completion.create(
+        model=model, temperature=temperature, prompt=prompt, max_tokens=max_tokens
+    )
+
+    response = completion.choices[0].text
+    return response
+
+
+def count_all_tokens(model: str, messages: list[dict[str, str]]) -> int:
+    total_tokens = 0
+    for message in messages:
+        total_tokens += count_tokens(model, message["content"])
+    return total_tokens
+
+
+def count_tokens(model: str, prompt: str) -> int:
+    if "gpt" not in model:
+        model = "gpt-4"
+
+    encoding = tiktoken.encoding_for_model(model)
+    num_tokens = len(encoding.encode(prompt))
+    return num_tokens
+
+
+def get_prompt_cost(model: str, num_tokens: int) -> float:
+    # good enough for quick evals
+    if model not in pricing_dict:
+        return num_tokens * 0.001 * pricing_dict["gpt-4"]
+    return num_tokens * 0.001 * pricing_dict[model]
+
+
+def get_response_cost(model: str, num_tokens: int) -> float:
+    # good enough for quick evals
+    if model not in pricing_dict:
+        return num_tokens * 0.001 * pricing_dict["gpt-4"]
+
+    cost = num_tokens * 0.001 * pricing_dict[model]
+
+    if model == "gpt-4":
+        cost *= 2
+
+    return cost
+
+
+def record_messages(messages: list[dict], model: str):
+    # create the conversation in a human-readable format
+    conversation_text = ""
+    for message in messages:
+        conversation_text += message["content"]
+
+    # write the conversation to the next available text file
+    with open(f"gpt_outputs/transcript.txt", "w") as f:
+        f.write(model + "\n\n")
+        f.write(conversation_text)

chess-gpt-eval/llama_module.py

@@ -0,0 +1,71 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+import torch
+
+from typing import Optional
+
+
+# There are a couple non optimal parts of this code:
+# 1. It doesn't inherit the Player class in main.py, which throws type checking errors
+# 2. get_move_from_response() is duplicated from main.py
+# However, I didn't want to add clutter and major dependencies like torch, peft, and transformers
+# to those not using this class. So, this was my compromise.
+class BaseLlamaPlayer:
+    def __init__(
+        self, tokenizer: AutoTokenizer, model: AutoModelForCausalLM, model_name: str
+    ):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.model_name = model_name
+
+    def get_llama_response(self, game_state: str, temperature: float) -> Optional[str]:
+        prompt = game_state
+        tokenized_input = self.tokenizer(prompt, return_tensors="pt").to("cuda")
+        result = self.model.generate(
+            **tokenized_input, max_new_tokens=10, temperature=temperature
+        ).to("cpu")
+        input_ids_tensor = tokenized_input["input_ids"]
+        # transformers generate() returns <s> + prompt + output. This grabs only the output
+        res_sliced = result[:, input_ids_tensor.shape[1] :]
+        return self.tokenizer.batch_decode(res_sliced)[0]
+
+    def get_move_from_response(self, response: Optional[str]) -> Optional[str]:
+        if response is None:
+            return None
+
+        # Parse the response to get only the first move
+        moves = response.split()
+        first_move = moves[0] if moves else None
+
+        return first_move
+
+    def get_move(
+        self, board: str, game_state: str, temperature: float
+    ) -> Optional[str]:
+        completion = self.get_llama_response(game_state, temperature)
+        return self.get_move_from_response(completion)
+
+    def get_config(self) -> dict:
+        return {"model": self.model_name}
+
+
+class LocalLlamaPlayer(BaseLlamaPlayer):
+    def __init__(self, model_name: str):
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16, device_map=0
+        ).to("cuda")
+        super().__init__(tokenizer, model, model_name)
+
+
+class LocalLoraLlamaPlayer(BaseLlamaPlayer):
+    def __init__(self, base_model_id: str, adapter_model_path: str):
+        tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+        base_model = AutoModelForCausalLM.from_pretrained(base_model_id)
+        model = (
+            PeftModel.from_pretrained(base_model, adapter_model_path)
+            .merge_and_unload()
+            .to("cuda")
+        )
+
+        super().__init__(tokenizer, model, adapter_model_path)

chess-gpt-eval/main.py

@@ -0,0 +1,565 @@
+import openai
+import chess
+import chess.engine
+import os
+import csv
+import random
+import time
+import platform
+
+# NOTE: LLAMA AND NANOGPT ARE EXPERIMENTAL PLAYERS, if not using them, comment them out
+# from llama_module import BaseLlamaPlayer, LocalLlamaPlayer, LocalLoraLlamaPlayer
+from nanogpt.nanogpt_module import NanoGptPlayer
+from mamba_module import MambaPlayer
+import gpt_query
+from lczero.backends import Weights, Backend, GameState
+import numpy as np
+
+from typing import Optional, Tuple
+from dataclasses import dataclass
+
+
+@dataclass
+class LegalMoveResponse:
+    move_san: Optional[str] = None
+    move_uci: Optional[chess.Move] = None
+    attempts: int = 0
+    is_resignation: bool = False
+    is_illegal_move: bool = False
+
+
+# Define base Player class
+class Player:
+    def get_move(self, board: chess.Board, game_state: str, temperature: float) -> str:
+        raise NotImplementedError
+
+    def get_config(self) -> dict:
+        raise NotImplementedError
+
+
+class GPTPlayer(Player):
+    def __init__(self, model: str):
+        with open("gpt_inputs/api_key.txt", "r") as f:
+            openai.api_key = f.read().strip()
+        self.model = model
+
+    def get_move(
+        self, board: chess.Board, game_state: str, temperature: float
+    ) -> Optional[str]:
+        response = get_gpt_response(game_state, self.model, temperature)
+        return get_move_from_gpt_response(response)
+
+    def get_config(self) -> dict:
+        return {"model": self.model}
+        
+        
+class LC0PLayer(Player):
+    # "11258-32x4-se.pb.gz" = stockfish level 0- = skill 0
+    # "11258-48x5-se.pb.gz" = stockfish level 0+ = skill 1
+    # "11258-80x7-se.pb.gz" = stockfish level 1 = skill 2
+    # "11258-104x9-se.pb.gz" = stockfish level 2 = skill 3
+    # "TK-6430 aka 128x10-BPR-64M-6430000.pb.gz" = stockfish level 3 = skill 4
+    # "00af53b081e80147172e6f281c01daf5ca19ada173321438914c730370aa4267" = stockfish level 4 = skill 5
+    # "b2ec465d0fb5b5eb39d2e1e3f74041a5d2fc92d413b71aa7ea0b6fb082ccba9c" = stockfish level 5+ = skill 6
+    def __init__(self, skill):
+        self.skill = skill
+        network_paths = ["./lc0/build/release/11258-32x4-se.pb.gz", "./lc0/build/release/11258-48x5-se.pb.gz", "./lc0/build/release/11258-80x7-se.pb.gz", "./lc0/build/release/11258-104x9-se.pb.gz", "./lc0/build/release/TK-6430 aka 128x10-BPR-64M-6430000.pb.gz", "./lc0/build/release/00af53b081e80147172e6f281c01daf5ca19ada173321438914c730370aa4267", "./lc0/build/release/b2ec465d0fb5b5eb39d2e1e3f74041a5d2fc92d413b71aa7ea0b6fb082ccba9c"]
+        print(f"\n\nLoading lc0 network: {network_paths[skill]}\n\n")
+        self.weights = Weights(network_paths[skill])
+        self.backend = Backend(weights=self.weights)
+        self.gamestate = GameState()
+
+    def get_move(self, board: chess.Board, game_state: str, temperature: float):
+        self.gamestate = GameState(fen=board.fen())
+        input_planes = self.gamestate.as_input(self.backend)
+        result = self.backend.evaluate(input_planes)[0]
+        moves = self.gamestate.moves()
+        policy_indices = self.gamestate.policy_indices()
+        move_probs = np.array(result.p_softmax(*policy_indices))
+        best_move_idx = move_probs.argmax()
+        best_move = moves[best_move_idx]
+        return board.san(chess.Move.from_uci(best_move))
+        
+    def get_config(self) -> dict:
+        return {"network": self.weights, "skill_level": self.skill, "play_time": 0}
+
+
+class StockfishPlayer(Player):
+
+    @staticmethod
+    def get_stockfish_path() -> str:
+        """
+        Determines the operating system and returns the appropriate path for Stockfish.
+        
+        Returns:
+            str: Path to the Stockfish executable based on the operating system.
+        """
+        if platform.system() == 'Linux':
+            return "/usr/games/stockfish"
+        elif platform.system() == 'Darwin':  # Darwin is the system name for macOS
+            return "stockfish"
+        elif platform.system() == 'Windows':
+            return r"C:\Users\Haile\Downloads\stockfish\stockfish-windows-x86-64-avx2.exe"
+        else:
+            raise OSError("Unsupported operating system")
+    
+    def __init__(self, skill_level: int, play_time: float):
+        self._skill_level = skill_level
+        self._play_time = play_time
+        # If getting started, you need to run brew install stockfish
+        stockfish_path = StockfishPlayer.get_stockfish_path()
+        self._engine = chess.engine.SimpleEngine.popen_uci(stockfish_path)
+
+    def get_move(
+        self, board: chess.Board, game_state: str, temperature: float
+    ) -> Optional[str]:
+        if self._skill_level == -2:
+            legal_moves = list(board.legal_moves)
+            random_move = random.choice(legal_moves)
+            return board.san(random_move)
+        elif self._skill_level < 0:
+            self._engine.configure({"Skill Level": 0})
+            result = self._engine.play(
+                board, chess.engine.Limit(time=1e-8, depth=1, nodes=1)
+            )
+
+        else:
+            self._engine.configure({"Skill Level": self._skill_level})
+            result = self._engine.play(board, chess.engine.Limit(time=self._play_time))
+        if result.move is None:
+            return None
+        return board.san(result.move)
+
+    def get_config(self) -> dict:
+        return {"skill_level": self._skill_level, "play_time": self._play_time}
+
+    def close(self):
+        self._engine.quit()
+        
+        
+class HumanPlayer(Player):
+    def get_move(self, board: chess.Board, game_state: str, temperature: float) -> str:
+        # Print board for human player
+        print(board)  
+        while True:
+            move = input("Enter your move (SAN format): ")
+            try:
+                move_uci = board.parse_san(move) 
+                if move_uci in board.legal_moves:  
+                    return move
+            except:
+                print("Illegal move, try again.")
+
+    def get_config(self) -> dict:
+        return {"player": "human"}
+
+
+def get_gpt_response(game_state: str, model: str, temperature: float) -> Optional[str]:
+    # trying to prevent what I believe to be rate limit issues
+    if model == "gpt-4":
+        time.sleep(0.4)
+    response = gpt_query.get_gpt_response(game_state, model, temperature)
+    return response
+
+
+def get_move_from_gpt_response(response: Optional[str]) -> Optional[str]:
+    if response is None:
+        return None
+
+    # Parse the response to get only the first move
+    moves = response.split()
+    first_move = moves[0] if moves else None
+
+    return first_move
+
+
+def record_results(
+    board: chess.Board,
+    player_one: Player,
+    player_two: Player,
+    game_state: str,
+    player_one_illegal_moves: int,
+    player_two_illegal_moves: int,
+    player_one_legal_moves: int,
+    player_two_legal_moves: int,
+    total_time: float,
+    player_one_resignation: bool,
+    player_two_resignation: bool,
+    player_one_failed_to_find_legal_move: bool,
+    player_two_failed_to_find_legal_move: bool,
+    total_moves: int,
+    illegal_moves: int,
+):
+    unique_game_id = generate_unique_game_id()
+
+    (
+        player_one_title,
+        player_two_title,
+        player_one_time,
+        player_two_time,
+    ) = get_player_titles_and_time(player_one, player_two)
+
+    if player_one_resignation or player_one_failed_to_find_legal_move:
+        result = "0-1"
+        player_one_score = 0
+        player_two_score = 1
+    elif player_two_resignation or player_two_failed_to_find_legal_move:
+        result = "1-0"
+        player_one_score = 1
+        player_two_score = 0
+    else:
+        result = board.result()
+        # Hmmm.... debating this one. Annoying if I leave it running and it fails here for some reason, probably involving some
+        # resignation / failed move situation I didn't think of
+        # -1e10 at least ensures it doesn't fail silently
+        if "-" in result:
+            player_one_score = result.split("-")[0]
+            player_two_score = result.split("-")[1]
+        elif result == "*": # Draw due to hitting max moves
+            player_one_score = 0#1/2
+            player_two_score = 1#1/2
+        else:
+            player_one_score = -1e10
+            player_two_score = -1e10
+
+    info_dict = {
+        "game_id": unique_game_id,
+        "transcript": game_state,
+        "result": result,
+        "player_one": player_one_title,
+        "player_two": player_two_title,
+        "player_one_time": player_one_time,
+        "player_two_time": player_two_time,
+        "player_one_score": player_one_score,
+        "player_two_score": player_two_score,
+        "player_one_illegal_moves": player_one_illegal_moves,
+        "player_two_illegal_moves": player_two_illegal_moves,
+        "player_one_legal_moves": player_one_legal_moves,
+        "player_two_legal_moves": player_two_legal_moves,
+        "player_one_resignation": player_one_resignation,
+        "player_two_resignation": player_two_resignation,
+        "player_one_failed_to_find_legal_move": player_one_failed_to_find_legal_move,
+        "player_two_failed_to_find_legal_move": player_two_failed_to_find_legal_move,
+        "game_title": f"{player_one_title} vs. {player_two_title}",
+        "number_of_moves": board.fullmove_number,
+        "time_taken": total_time,
+        "total_moves": total_moves,
+        "illegal_moves": illegal_moves,
+    }
+
+    if RUN_FOR_ANALYSIS:
+        csv_file_path = f"logs/{player_one_recording_name}_vs_{player_two_recording_name}"
+        csv_file_path = csv_file_path.replace(".", "_") # Because I'm using ckpt filenames for nanogpt models 
+        csv_file_path += ".csv"
+    else:
+        csv_file_path = recording_file
+
+
+
+    # Determine if we need to write headers (in case the file doesn't exist yet)
+    write_headers = not os.path.exists(csv_file_path)
+
+    # Append the results to the CSV file
+    with open(csv_file_path, "a", newline="") as csv_file: # THIS WAS APPEND
+        writer = csv.DictWriter(csv_file, fieldnames=info_dict.keys())
+        if write_headers:
+            writer.writeheader()
+        writer.writerow(info_dict)
+
+    with open("game.txt", "w") as f:
+        f.write(game_state)
+
+
+def generate_unique_game_id() -> str:
+    timestamp = int(time.time())
+    random_num = random.randint(1000, 9999)  # 4-digit random number
+    return f"{timestamp}-{random_num}"
+
+
+def get_player_titles_and_time(
+    player_one: Player, player_two: Player
+) -> Tuple[str, str, Optional[float], Optional[float]]:
+    player_one_config = player_one.get_config()
+    player_two_config = player_two.get_config()
+
+    # For player one
+    if "model" in player_one_config:
+        player_one_title = player_one_config["model"]
+        player_one_time = None
+    else:
+        player_one_title = f"Stockfish {player_one_config['skill_level']}"
+        player_one_time = player_one_config["play_time"]
+
+    # For player two
+    if "model" in player_two_config:
+        player_two_title = player_two_config["model"]
+        player_two_time = None
+    else:
+        player_two_title = f"Stockfish {player_two_config['skill_level']}"
+        player_two_time = player_two_config["play_time"]
+
+    return (player_one_title, player_two_title, player_one_time, player_two_time)
+
+
+used_openings = []
+def initialize_game_with_opening(
+    game_state: str, board: chess.Board
+) -> Tuple[str, chess.Board]:
+    global used_openings
+    with open("openings.csv", "r") as file:
+        lines = file.readlines()[1:]  # Skip header
+    moves_string = random.choice(lines)
+    while moves_string in used_openings:
+        moves_string = random.choice(lines)
+    used_openings.append(moves_string)
+    if move_num_in_gamestate:
+        game_state = moves_string.rstrip() + " "
+    else:
+        game_state = ' '.join(['.' + m.split(".")[-1] if "." in m else m for m in moves_string.split()])
+        game_state = game_state.rstrip() + " "
+    # Splitting the moves string on spaces
+    tokens = moves_string.split()
+
+    for token in tokens:
+        # If the token contains a period, it's a move number + move combination
+        if "." in token:
+            move = token.split(".")[-1]  # Take the move part after the period
+        else:
+            move = token
+
+        board.push_san(move)
+    return game_state.rstrip(), board
+
+
+# Return is (move_san, move_uci, attempts, is_resignation, is_illegal_move)
+def get_legal_move(
+    player: Player,
+    board: chess.Board,
+    game_state: str,
+    player_one: bool,
+    max_attempts: int = 5,
+) -> LegalMoveResponse:
+    """Request a move from the player and ensure it's legal."""
+    move_san = None
+    move_uci = None
+
+    for attempt in range(max_attempts):
+        #print(f"get_legal_move: |{game_state}|")
+        move_san = player.get_move(
+            board, game_state, min(((attempt / max_attempts) * 1) + 0.001, 0.75)
+        )
+
+        # Sometimes when GPT thinks it's the end of the game, it will just output the result
+        # Like "1-0". If so, this really isn't an illegal move, so we'll add a check for that.
+        if move_san is not None:
+            if move_san == "1-0" or move_san == "0-1" or move_san == "1/2-1/2":
+                print(f"{move_san}, player has resigned")
+                return LegalMoveResponse(
+                    move_san=None,
+                    move_uci=None,
+                    attempts=attempt,
+                    is_resignation=True,
+                )
+
+        try:
+            move_uci = board.parse_san(move_san)
+        except Exception as e:
+            print(f"Error parsing move {move_san}: {e}")
+            # check if player is gpt-3.5-turbo-instruct
+            # only recording errors for gpt-3.5-turbo-instruct because it's errors are so rare
+            if player.get_config()["model"] == "gpt-3.5-turbo-instruct":
+                with open("gpt-3.5-turbo-instruct-illegal-moves.txt", "a") as f:
+                    f.write(f"{game_state}\n{move_san}\n")
+            continue
+
+        if move_uci in board.legal_moves:
+            if player_one == False:
+                if not move_san.startswith(" "):
+                    move_san = " " + move_san
+            else:
+                if move_san.startswith(" "):
+                    move_san = move_san[1:]
+            return LegalMoveResponse(move_san, move_uci, attempt)
+        print(f"Illegal move: {move_san}")
+
+    # If we reach here, the player has made illegal moves for all attempts.
+    print(f"{player} provided illegal moves for {max_attempts} attempts.")
+    return LegalMoveResponse(
+        move_san=None, move_uci=None, attempts=max_attempts, is_illegal_move=True
+    )
+
+
+def play_turn(
+    player: Player, board: chess.Board, game_state: str, player_one: bool
+) -> Tuple[str, bool, bool, int]:
+    result = get_legal_move(player, board, game_state, player_one, 5)
+    illegal_moves = result.attempts
+    move_san = result.move_san
+    move_uci = result.move_uci
+    resignation = result.is_resignation
+    failed_to_find_legal_move = result.is_illegal_move
+
+    if resignation:
+        print(f"{player} resigned with result: {board.result()}")
+    elif failed_to_find_legal_move:
+        print(f"Game over: 5 consecutive illegal moves from {player}")
+    elif move_san is None or move_uci is None:
+        print(f"Game over: {player} failed to find a legal move")
+    else:
+        board.push(move_uci)
+        game_state += move_san
+        print(move_san, end=" ")
+
+    return game_state, resignation, failed_to_find_legal_move, illegal_moves
+
+
+def play_game(
+    player_one: Player,
+    player_two: Player,
+    max_games: int = 10,
+    random_opening_seed: bool = False,
+):
+    for z in range(max_games):
+        print(f"\nGame {z} of {max_games}\n")
+    
+        with open("gpt_inputs/prompt.txt", "r") as f:
+            game_state = f.read()
+        board = chess.Board()
+
+        if random_opening_seed:
+            game_state, board = initialize_game_with_opening(game_state, board)
+            #print(f"play_gamea after init: |{game_state}|")
+        player_one_illegal_moves = 0
+        player_two_illegal_moves = 0
+        player_one_legal_moves = 0
+        player_two_legal_moves = 0
+        player_one_resignation = False
+        player_two_resignation = False
+        player_one_failed_to_find_legal_move = False
+        player_two_failed_to_find_legal_move = False
+        start_time = time.time()
+
+        total_moves = 0
+        illegal_moves = 0
+        print_for_human = isinstance(player_one, HumanPlayer) or isinstance(player_two, HumanPlayer)
+
+        while not board.is_game_over():
+            if print_for_human:
+                print(board)
+
+            with open("game.txt", "w") as f:
+                f.write(game_state)
+            current_move_num = f"{board.fullmove_number if move_num_in_gamestate else ''}."
+            total_moves += 1
+            # I increment legal moves here so player_two isn't penalized for the game ending before its turn
+            player_one_legal_moves += 1
+            player_two_legal_moves += 1
+
+            # this if statement may be overkill, just trying to get format to exactly match PGN notation
+            if board.fullmove_number != 1:
+                game_state += " "
+            game_state += current_move_num
+            #print(f"|{game_state}|")
+            #print(f"{current_move_num}", end=" ")
+
+            (
+                game_state,
+                player_one_resignation,
+                player_one_failed_to_find_legal_move,
+                illegal_moves_one,
+            ) = play_turn(player_one, board, game_state, player_one=True)
+            player_one_illegal_moves += illegal_moves_one
+            if illegal_moves_one != 0:
+                player_one_legal_moves -= 1
+            if (
+                board.is_game_over()
+                or player_one_resignation
+                or player_one_failed_to_find_legal_move
+            ):
+                break
+
+            (
+                game_state,
+                player_two_resignation,
+                player_two_failed_to_find_legal_move,
+                illegal_moves_two,
+            ) = play_turn(player_two, board, game_state, player_one=False)
+            player_two_illegal_moves += illegal_moves_two
+            if illegal_moves_two != 0:
+                player_two_legal_moves -= 1
+            if (
+                board.is_game_over()
+                or player_two_resignation
+                or player_two_failed_to_find_legal_move
+            ):
+                break
+
+            print("\n", end="")
+
+            if total_moves > MAX_MOVES:
+                break
+
+        end_time = time.time()
+        total_time = end_time - start_time
+        print(f"\nGame over. Total time: {total_time} seconds")
+        print(f"Result: {board.result()}")
+        print(board)
+        print()
+        record_results(
+            board,
+            player_one,
+            player_two,
+            game_state,
+            player_one_illegal_moves,
+            player_two_illegal_moves,
+            player_one_legal_moves,
+            player_two_legal_moves,
+            total_time,
+            player_one_resignation,
+            player_two_resignation,
+            player_one_failed_to_find_legal_move,
+            player_two_failed_to_find_legal_move,
+            total_moves,
+            illegal_moves,
+        )
+    if isinstance(player_one, StockfishPlayer):
+        player_one.close()
+    if isinstance(player_two, StockfishPlayer):
+        player_two.close()
+
+        # print(game_state)
+
+
+RUN_FOR_ANALYSIS = True
+MAX_MOVES = 999 # Due to nanogpt max input length of 1024
+recording_file = "logs/determine.csv" # default recording file. Because we are using list [player_ones], recording_file is overwritten
+# player_one_recording_name = "ckpt_8.pt"
+#player_ones = ["ckpt_iter_20000.pt","ckpt_iter_40000.pt","ckpt_iter_60000.pt","ckpt_iter_80000.pt"] #["ckpt.pt"]
+player_ones = ["Xformer/6.6M/ckpt.pt"]
+player_two_recording_name = "lc0_sweep" #"stockfish_sweep"
+move_num_in_gamestate = False
+if __name__ == "__main__":
+    for nanogpt_player in player_ones:
+        player_one_recording_name = nanogpt_player
+        for i in range(2): #range(11):
+            num_games = 265 #265 instead of 250 for duplicates (for lc0, stockfish doesn't need it)
+            # player_one = GPTPlayer(model="gpt-3.5-turbo-instruct")
+            # player_one = LocalLlamaPlayer(model_name="meta-llama/Llama-2-7b-hf")
+            # player_one = LocalLoraLlamaPlayer("meta-llama/Llama-2-7b-hf", "/workspace/axolotl/lora2-out")
+            # player_one = GPTPlayer(model="gpt-4")
+            # player_one = StockfishPlayer(skill_level=-1, play_time=0.1)
+            
+            player_one = NanoGptPlayer(model_name=player_one_recording_name, move_num_in_gamestate=False)
+           # player_one = MambaPlayer(model_name=player_one_recording_name, move_num_in_gamestate=False)
+            #player_two = StockfishPlayer(skill_level=i, play_time=0.1)
+            player_two = LC0PLayer(skill=i)
+            
+            # player_two = GPTPlayer(model="gpt-4")
+            # player_two = GPTPlayer(model="gpt-3.5-turbo-instruct")
+            
+            print(f"\n\nSTARTING GAMES AGAINST STOCKFISH LEVEL {i}\n\n")
+            #print(f"\n\nSTARTING GAMES AGAINST LC0 LEVEL {i}\n\n")
+
+            play_game(player_one, player_two, num_games, random_opening_seed=True)
+            
+print("\n\n\n********\nDONE!\n********\n\n\n")

chess-gpt-eval/mamba.py

@@ -0,0 +1,368 @@
+import math
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pscan import pscan
+
+"""
+
+This file closely follows the mamba_simple.py from the official Mamba implementation, and the mamba-minimal by @johnma2006.
+The major differences are :
+-the convolution is done with torch.nn.Conv1d
+-the selective scan is done in PyTorch
+
+A sequential version of the selective scan is also available for comparison.
+
+- A Mamba model is composed of several layers, which are ResidualBlock.
+- A ResidualBlock is composed of a MambaBlock, a normalization, and a residual connection : ResidualBlock(x) = mamba(norm(x)) + x
+- This leaves us with the MambaBlock : its input x is (B, L, D) and its outputs y is also (B, L, D) (B=batch size, L=seq len, D=model dim).
+First, we expand x into (B, L, 2*ED) (where E is usually 2) and split it into x and z, each (B, L, ED).
+Then, we apply the short 1d conv to x, followed by an activation function (silu), then the SSM.
+We then multiply it by silu(z).
+See Figure 3 of the paper (page 8) for a visual representation of a MambaBlock.
+
+"""
+
+@dataclass
+class MambaConfig:
+    d_model: int # D
+    n_layers: int
+    dt_rank: Union[int, str] = 'auto'
+    d_state: int = 16 # N in paper/comments
+    expand_factor: int = 2 # E in paper/comments
+    d_conv: int = 4
+
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dt_init: str = "random" # "random" or "constant"
+    dt_scale: float = 1.0
+    dt_init_floor = 1e-4
+
+    bias: bool = False
+    conv_bias: bool = True
+
+    pscan: bool = True # use parallel scan mode or sequential mode when training
+
+    def __post_init__(self):
+        self.d_inner = self.expand_factor * self.d_model # E*D = ED in comments
+
+        if self.dt_rank == 'auto':
+            self.dt_rank = math.ceil(self.d_model / 16)
+
+class Mamba(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.config = config
+
+        self.layers = nn.ModuleList([ResidualBlock(config) for _ in range(config.n_layers)])
+        #self.norm_f = RMSNorm(config.d_model)
+
+    def forward(self, x):
+        # x : (B, L, D)
+
+        # y : (B, L, D)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        #x = self.norm_f(x)
+
+        return x
+    
+    def step(self, x, caches):
+        # x : (B, L, D)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        # y : (B, L, D)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        for i, layer in enumerate(self.layers):
+            x, caches[i] = layer.step(x, caches[i])
+
+        return x, caches
+
+class ResidualBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.mixer = MambaBlock(config)
+        self.norm = RMSNorm(config.d_model)
+
+    def forward(self, x):
+        # x : (B, L, D)
+
+        # output : (B, L, D)
+
+        output = self.mixer(self.norm(x)) + x
+        return output
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs: (B, ED, d_conv-1)
+
+        # output : (B, D)
+        # cache : (h, inputs)
+
+        output, cache = self.mixer.step(self.norm(x), cache)
+        output = output + x
+        return output, cache
+
+class MambaBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.config = config
+
+        # projects block input from D to 2*ED (two branches)
+        self.in_proj = nn.Linear(config.d_model, 2 * config.d_inner, bias=config.bias)
+
+        self.conv1d = nn.Conv1d(in_channels=config.d_inner, out_channels=config.d_inner, 
+                              kernel_size=config.d_conv, bias=config.conv_bias, 
+                              groups=config.d_inner,
+                              padding=config.d_conv - 1)
+                              
+        nn.init.kaiming_normal_(self.conv1d.weight, mode='fan_out', nonlinearity='leaky_relu')
+        
+        # projects x to input-dependent Δ, B, C
+        self.x_proj = nn.Linear(config.d_inner, config.dt_rank + 2 * config.d_state, bias=False)
+
+        # projects Δ from dt_rank to d_inner
+        self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
+
+        # dt initialization
+        # dt weights
+        dt_init_std = config.dt_rank**-0.5 * config.dt_scale
+        if config.dt_init == "constant":
+            nn.init.constant_(self.dt_proj.weight, dt_init_std)
+        elif config.dt_init == "random":
+            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+        
+        # dt bias
+        dt = torch.exp(
+            torch.rand(config.d_inner) * (math.log(config.dt_max) - math.log(config.dt_min)) + math.log(config.dt_min)
+        ).clamp(min=config.dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt)) # inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(inv_dt)
+        #self.dt_proj.bias._no_reinit = True # initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
+        # todo : explain why removed
+
+        # S4D real initialization
+        A = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat(config.d_inner, 1)
+        self.A_log = nn.Parameter(torch.log(A)) # why store A in log ? to keep A < 0 (cf -torch.exp(...)) ? for gradient stability ?
+        self.D = nn.Parameter(torch.ones(config.d_inner))
+
+        # projects block output from ED back to D
+        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
+
+    def forward(self, x):
+        # x : (B, L, D)
+        
+        # y : (B, L, D)
+
+        _, L, _ = x.shape
+
+        xz = self.in_proj(x) # (B, L, 2*ED)
+        x, z = xz.chunk(2, dim=-1) # (B, L, ED), (B, L, ED)
+
+        # x branch
+        x = x.transpose(1, 2) # (B, ED, L)
+        x = self.conv1d(x)[:, :, :L] # depthwise convolution over time, with a short filter
+        x = x.transpose(1, 2) # (B, L, ED)
+
+        x = F.silu(x)
+        y = self.ssm(x)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, L, D)
+
+        return output
+    
+    def ssm(self, x):
+        # x : (B, L, ED)
+
+        # y : (B, L, ED)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N)
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, L, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, L, dt_rank), (B, L, N), (B, L, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, L, ED)
+
+        if self.config.pscan:
+            y = self.selective_scan(x, delta, A, B, C, D)
+        else:
+            y = self.selective_scan_seq(x, delta, A, B, C, D)
+
+        return y
+    
+    def selective_scan(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+        
+        hs = pscan(deltaA, BX)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    def selective_scan_seq(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        _, L, _ = x.shape
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+
+        h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+        hs = []
+
+        for t in range(0, L):
+            h = deltaA[:, t] * h + BX[:, t]
+            hs.append(h)
+            
+        hs = torch.stack(hs, dim=1) # (B, L, ED, N)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    # -------------------------- inference -------------------------- #
+    """
+    Concerning auto-regressive inference
+
+    The cool part of using Mamba : inference is constant wrt to sequence length
+    We just have to keep in cache, for each layer, two things :
+    - the hidden state h (which is (B, ED, N)), as you typically would when doing inference with a RNN
+    - the last d_conv-1 inputs of the layer, to be able to compute the 1D conv which is a convolution over the time dimension
+      (d_conv is fixed so this doesn't incur a growing cache as we progress on generating the sequence)
+      (and d_conv is usually very small, like 4, so we just have to "remember" the last 3 inputs)
+
+    Concretely, these two quantities are put inside a cache tuple, and are named h and inputs respectively.
+    h is (B, ED, N), and inputs is (B, ED, d_conv-1)
+    The MambaBlock.step() receives this cache, and, along with outputing the output, alos outputs the updated cache for the next call.
+
+    The cache object is initialized as follows : (None, torch.zeros()).
+    When h is None, the selective scan function detects it and start with h=0.
+    The torch.zeros() isn't a problem (it's same as just feeding the input, because the conv1d is padded)
+
+    As we need one such cache variable per layer, we store a caches object, which is simply a list of cache object. (See mamba_lm.py)
+    """
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs : (B, ED, d_conv-1)
+        
+        # y : (B, D)
+        # cache : (h, inputs)
+        
+        h, inputs = cache
+        
+        xz = self.in_proj(x) # (B, 2*ED)
+        x, z = xz.chunk(2, dim=1) # (B, ED), (B, ED)
+
+        # x branch
+        x_cache = x.unsqueeze(2)
+        x = self.conv1d(torch.cat([inputs, x_cache], dim=2))[:, :, self.config.d_conv-1] # (B, ED)
+
+        x = F.silu(x)
+        y, h = self.ssm_step(x, h)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, D)
+
+        # prepare cache for next call
+        inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) # (B, ED, d_conv-1)
+        cache = (h, inputs)
+        
+        return output, cache
+
+    def ssm_step(self, x, h):
+        # x : (B, ED)
+        # h : (B, ED, N)
+
+        # y : (B, ED)
+        # h : (B, ED, N)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N) # todo : ne pas le faire tout le temps, puisque c'est indépendant de la timestep
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, dt_rank), (B, N), (B, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(1) # (B, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, ED, N)
+
+        if h is None:
+            h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+
+        h = deltaA * h + BX # (B, ED, N)
+
+        y = (h @ C.unsqueeze(-1)).squeeze(2) # (B, ED, N) @ (B, N, 1) -> (B, ED, 1)
+
+        y = y + D * x
+
+        # todo : pq h.squeeze(1) ??
+        return y, h.squeeze(1)
+
+# taken straight from https://github.com/johnma2006/mamba-minimal/blob/master/model.py
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5):
+        super().__init__()
+
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+        return output

chess-gpt-eval/mamba/out/meta.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1121191e401988851de5744fe27fe463c3a086fc8c9a5538ef7fc12162bfb09
+size 373

chess-gpt-eval/mamba_lm.py

@@ -0,0 +1,168 @@
+from dataclasses import dataclass, fields, asdict
+import json
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from mamba import Mamba, MambaConfig, RMSNorm
+
+"""
+
+Encapsulates a Mamba model as language model. It has an embedding layer, and a LM head which maps the model output to logits.
+
+"""
+
+# TODO generate function : batch size != 1 ? (for now B=1)
+# TODO generate function : top-p sampling
+
+@dataclass
+class MambaLMConfig(MambaConfig):
+    vocab_size: int = 32000
+    pad_vocab_size_multiple: int = 8
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        #if self.vocab_size % self.pad_vocab_size_multiple != 0:
+        #    self.vocab_size += (self.pad_vocab_size_multiple - self.vocab_size % self.pad_vocab_size_multiple)
+
+    def to_mamba_config(self) -> MambaConfig:
+        mamba_config_fields = {field.name for field in fields(MambaConfig)}
+        filtered_dict = {k: v for k, v in asdict(self).items() if k in mamba_config_fields}
+        return MambaConfig(**filtered_dict)
+
+# adapted from https://github.com/johnma2006/mamba-minimal
+def from_pretrained(name: str):
+    """
+    Returns a model loaded with pretrained weights pulled from HuggingFace.
+
+    Args:
+        name: As of now, supports
+            * 'state-spaces/mamba-2.8b-slimpj'
+            * 'state-spaces/mamba-2.8b'
+            * 'state-spaces/mamba-1.4b'
+            * 'state-spaces/mamba-790m'
+            * 'state-spaces/mamba-370m'
+            * 'state-spaces/mamba-130m'
+
+    Returns:
+        model: a Mamba model configured with the proper parameters and initialized with the proper weights
+    """   
+
+    from transformers.utils import WEIGHTS_NAME, CONFIG_NAME
+    from transformers.utils.hub import cached_file
+
+    def load_config_hf(model_name):
+        resolved_archive_file = cached_file(model_name, CONFIG_NAME, _raise_exceptions_for_missing_entries=False)
+        return json.load(open(resolved_archive_file))
+                
+    def load_state_dict_hf(model_name):
+        resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
+        return torch.load(resolved_archive_file, weights_only=True, map_location='cpu', mmap=True)
+        
+    # copy config data
+    config_data = load_config_hf(name)
+    config = MambaLMConfig(d_model=config_data['d_model'], n_layers=config_data['n_layer'], vocab_size=config_data['vocab_size'])
+
+    model = MambaLM(config)
+
+    # copy weights
+    state_dict = load_state_dict_hf(name)
+
+    new_state_dict = {}
+    for key in state_dict:
+        if key == 'backbone.embedding.weight' or key == 'backbone.norm_f.weight':
+            new_key = key.replace('backbone.', '')
+        else:
+            new_key = key.replace('backbone', 'mamba')
+
+        new_state_dict[new_key] = state_dict[key]
+
+    model.load_state_dict(new_state_dict)
+
+    return model
+
+class MambaLM(nn.Module):
+    def __init__(self, lm_config: MambaLMConfig):
+        super().__init__()
+        self.lm_config = lm_config
+        self.config = lm_config.to_mamba_config()
+
+        self.embedding = nn.Embedding(self.lm_config.vocab_size, self.config.d_model)
+        self.mamba = Mamba(self.config)
+        self.norm_f = RMSNorm(self.config.d_model)
+
+        self.lm_head = nn.Linear(self.config.d_model, self.lm_config.vocab_size, bias=False)
+        self.lm_head.weight = self.embedding.weight
+
+    def forward(self, tokens):
+        # tokens : (B, L)
+
+        # logits : (B, L, vocab_size)
+
+        x = self.embedding(tokens)
+
+        x = self.mamba(x)
+        x = self.norm_f(x)
+
+        logits = self.lm_head(x)
+
+        return logits
+    
+    def step(self, token, caches):
+        # token : (B)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        # logits : (B, vocab_size)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        x = self.embedding(token)
+
+        x, caches = self.mamba.step(x, caches)
+        x = self.norm_f(x)
+
+        logits = self.lm_head(x)
+
+        return logits, caches
+    
+    # TODO temperature
+    # TODO process prompt in parallel, and pass in sequential mode when prompt is finished ?
+    def generate(self, tokenizer, prompt: str, num_tokens: int = 50, sample: bool = True, top_k: int = 40):
+        self.eval()
+
+        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(next(self.parameters()).device) # (1, num_tokens)
+
+        # caches is a list of cache, one per layer
+        # cache is composed of : the hidden state, and the last d_conv-1 inputs
+        # the hidden state because the update is like an RNN
+        # the last d_conv-1 inputs because they are used in a 1d convolution (usually d_conv=4 so this is not large)
+        caches = [(None, torch.zeros(1, self.config.d_inner, self.config.d_conv-1, device=input_ids.device)) for _ in range(self.config.n_layers)]
+
+        for i in range(input_ids.size(1) + num_tokens - 1):
+            with torch.no_grad():
+                # forward the new output, get new cache
+                next_token_logits, caches = self.step(input_ids[:, i], caches) # (1, vocab_size), caches
+
+            # sample (no sampling when the prompt is being processed)
+            if i+1 >= input_ids.size(1):
+                probs = F.softmax(next_token_logits, dim=-1) # (1, vocab_size)
+
+                if top_k is not None:
+                    values, _ = torch.topk(probs, k=top_k) # (1, k) ordered from lowest to biggest
+                    probs[probs < values[:, -1, None]] = 0
+                    probs = probs / probs.sum(axis=1, keepdims=True)
+
+                if sample:
+                    next_token = torch.multinomial(probs, num_samples=1).squeeze(1) # (1)
+                else:
+                    next_token = torch.argmax(probs, dim=-1) # (1)
+
+                input_ids = torch.cat([input_ids, next_token.unsqueeze(1)], dim=1)
+                
+        output = [tokenizer.decode(output.tolist()) for output in input_ids][0]
+
+        self.train()
+
+        return output
+    

chess-gpt-eval/mamba_module.py

@@ -0,0 +1,144 @@
+import os
+import pickle
+import torch
+from mamba_lm import MambaLM, MambaLMConfig, from_pretrained
+from contextlib import nullcontext
+
+BASE_DIR = "mamba/"
+
+class MambaPlayer:
+    def __init__(self, model_name: str, move_num_in_gamestate: bool=False):
+        self.model_name = model_name
+        self.move_num_in_gamestate = move_num_in_gamestate
+        # -----------------------------------------------------------------------------
+
+        init_from = "resume"  # either 'resume' or a Mamba variant (e.g. 'state-spaces/mamba-1.4b')
+        out_dir = "out"  # ignored if init_from is not 'resume'
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        #device = "cpu"
+        dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float32'
+        seed = 1337
+        compile = False  # set to True if using PyTorch 2.0 and Mamba supports it
+        # -----------------------------------------------------------------------------
+
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        
+        device_type = (
+            "cuda" if "cuda" in device else "cpu"
+        )  # for later use in torch.autocast
+        ptdtype = {
+            "float32": torch.float32,
+            "bfloat16": torch.bfloat16,
+            "float16": torch.float16,
+        }[dtype]
+        ctx = (
+            nullcontext()
+            if device_type == "cpu"
+            else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+        )
+
+        # Model initialization
+        if init_from == "resume":
+            #ckpt_path = os.path.join(BASE_DIR, out_dir, self.model_name)
+            ckpt_path = os.path.normpath(f"../../mamba.py/out/{self.model_name}")
+            checkpoint = torch.load(ckpt_path, map_location=device)
+            model_config = checkpoint["model_args"]
+            model = MambaLM(model_config)
+            model.load_state_dict(checkpoint['model'])
+        elif init_from.startswith('state-spaces'):
+            model = from_pretrained(init_from).to(device)
+        else:
+            raise ValueError("Invalid init_from value")
+
+        model.eval()
+        model.to(device)
+
+        if compile and hasattr(torch, 'compile'):
+            model = torch.compile(model)
+
+        # look for the meta pickle in case it is available in the dataset folder
+        meta_path = os.path.join(BASE_DIR, "out", "meta.pkl")
+        load_meta = os.path.exists(meta_path)
+        if move_num_in_gamestate and load_meta:
+            with open(meta_path, "rb") as f:
+                meta = pickle.load(f)
+            stoi, itos = meta["stoi"], meta["itos"]
+            vocab_size = meta['vocab_size']
+            encode = lambda s: [stoi[c] for c in s]
+            decode = lambda l: "".join([itos[i] for i in l])
+        else:
+            stoi = {' ': 0, '.': 1, 'a': 2, 'b': 3, 'c': 4, 'd': 5, 'e': 6, 'f': 7, 'g': 8, 'h': 9, '1': 10, '2': 11, '3': 12, '4': 13, '5': 14, '6': 15, '7': 16, '8': 17, 'B': 18, 'N': 19, 'R': 20, 'Q': 21, 'K': 22, 'O': 23, 'x': 24, '+': 25, '#': 26, '=': 27}
+            itos = {0: ' ', 1: '.', 2: 'a', 3: 'b', 4: 'c', 5: 'd', 6: 'e', 7: 'f', 8: 'g', 9: 'h', 10: '1', 11: '2', 12: '3', 13: '4', 14: '5', 15: '6', 16: '7', 17: '8', 18: 'B', 19: 'N', 20: 'R', 21: 'Q', 22: 'K', 23: 'O', 24: 'x', 25: '+', 26: '#', 27: '='}
+            for s in stoi:
+                assert itos[stoi[s]] == s
+            vocab_size = len(stoi)
+            print(f"Vocab size {vocab_size}")
+            encode = lambda s: [stoi[c] for c in s.replace('-', '')]
+            decode = lambda l: "".join([itos[i] for i in l]).replace("OOO", "O-O-O").replace("OO", "O-O")
+
+        self.vocab_size = vocab_size
+        self.encode = encode
+        self.decode = decode
+        self.model = model
+        self.ctx = ctx
+        self.device = device
+
+    def get_mamba_response(self, game_state: str, temperature: float, max_new_tokens: int, top_k: int):
+        game_state = game_state.split("\n\n")[-1].strip()
+        #game_state = ";" + game_state
+
+        # Tokenize the game state
+        encoded_prompt = self.encode(game_state)
+        input_ids = torch.tensor([encoded_prompt], dtype=torch.long, device=self.device)
+
+        self.model.eval()  # Set the model to evaluation mode
+        with torch.no_grad():
+            have_non_space = False
+            for _ in range(max_new_tokens):
+                logits = self.model(input_ids)[0, -1, :]  # Get logits for the last token
+
+                # Apply temperature scaling and optionally sample from top k tokens
+                logits = logits / temperature
+                if top_k > 0:
+                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+                    logits[indices_to_remove] = -float('Inf')
+
+                probs = torch.nn.functional.softmax(logits, dim=-1)
+                next_token_id = torch.multinomial(probs, num_samples=1)
+                if have_non_space and (next_token_id == 0 or next_token_id==4):
+                    break
+                else:
+                    have_non_space = True
+                input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)
+
+        model_response = self.decode(input_ids[0].tolist())
+        model_response = model_response[len(game_state):].split(";")[0]
+        return model_response
+
+    #def encode(self, text: str):
+        # Implement the appropriate tokenization for MambaLM
+        # This could be a simple mapping or a more complex tokenizer
+    #    return [stoi[char] for char in text]  # Example
+
+    #def decode(self, token_ids: list):
+        # Implement the appropriate decoding for MambaLM
+     #   return ''.join([itos[id] for id in token_ids])  # Example
+        
+    def get_move_from_response(self, response: str) -> str:
+        if not response:
+            return None
+        # Parse the response to get only the first move
+        moves = response.split()
+        first_move = moves[0]
+        first_move = first_move.lstrip('.') # A patch for a weird phase during training ... doesn't seem to be an issue anymore, but don't see the harm.
+
+        return first_move
+
+    def get_move(self, board: str, game_state: str, temperature: float) -> str:
+        completion = self.get_mamba_response(game_state, temperature, 8, self.vocab_size)
+        return self.get_move_from_response(completion)
+
+    def get_config(self) -> dict:
+        return {"model": self.model_name}
+

chess-gpt-eval/nanogpt/configurator.py

@@ -0,0 +1,47 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import sys
+from ast import literal_eval
+
+for arg in sys.argv[1:]:
+    if '=' not in arg:
+        # assume it's the name of a config file
+        assert not arg.startswith('--')
+        config_file = arg
+        print(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print(f.read())
+        exec(open(config_file).read())
+    else:
+        # assume it's a --key=value argument
+        assert arg.startswith('--')
+        key, val = arg.split('=')
+        key = key[2:]
+        if key in globals():
+            try:
+                # attempt to eval it it (e.g. if bool, number, or etc)
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                # if that goes wrong, just use the string
+                attempt = val
+            # ensure the types match ok
+            assert type(attempt) == type(globals()[key])
+            # cross fingers
+            print(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")

chess-gpt-eval/nanogpt/model.py

@@ -0,0 +1,330 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters
+        #print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {} # default to empty dict
+        # only dropout can be overridden see more notes below
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        config_args['bias'] = True # always True for GPT model checkpoints
+        # we can override the dropout rate, if desired
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
+        # this means that we have to transpose these weights when we import them
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

chess-gpt-eval/nanogpt/nanogpt_module.py

@@ -0,0 +1,148 @@
+"""
+Sample from a trained model
+"""
+import os
+import pickle
+from contextlib import nullcontext
+import torch
+import tiktoken
+from nanogpt.model import GPTConfig, GPT
+
+BASE_DIR = "nanogpt/"
+
+
+class NanoGptPlayer:
+    def __init__(self, model_name: str, move_num_in_gamestate: bool=False):
+        self.model_name = model_name
+        # -----------------------------------------------------------------------------
+
+        init_from = "resume"  # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl')
+        out_dir = "out"  # ignored if init_from is not 'resume'
+        input_dir = "addition"
+        test_name = "test.txt"
+        start = "12+44="  # or "<|endoftext|>" or etc. Can also specify a file, use as: "FILE:prompt.txt"
+        num_samples = 1  # number of samples to draw
+        max_new_tokens = 6  # number of tokens generated in each sample
+        temperature = 0.01  # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
+        top_k = 200  # retain only the top_k most likely tokens, clamp others to have 0 probability
+        seed = 1337
+        device = "cuda"  # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
+        #device = "cpu"
+        dtype = "float16"  # 'float32' or 'bfloat16' or 'float16'
+        compile = False  # use PyTorch 2.0 to compile the model to be faster
+        exec(
+            open(f"{BASE_DIR}configurator.py").read()
+        )  # overrides from command line or config file
+        # -----------------------------------------------------------------------------
+
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cuda.matmul.allow_tf32 = True  # allow tf32 on matmul
+        torch.backends.cudnn.allow_tf32 = True  # allow tf32 on cudnn
+        device_type = (
+            "cuda" if "cuda" in device else "cpu"
+        )  # for later use in torch.autocast
+        ptdtype = {
+            "float32": torch.float32,
+            "bfloat16": torch.bfloat16,
+            "float16": torch.float16,
+        }[dtype]
+        ctx = (
+            nullcontext()
+            if device_type == "cpu"
+            else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+        )
+
+        # model
+        if init_from == "resume":
+            # init from a model saved in a specific directory
+            #ckpt_path = os.path.join(BASE_DIR, out_dir, self.model_name)
+            ckpt_path = os.path.normpath(f"../../mamba.py/out/{self.model_name}")
+            checkpoint = torch.load(ckpt_path, map_location=device)
+            #gptconf = GPTConfig(**checkpoint["model_args"])
+            #model = GPT(gptconf)
+            model = GPT(checkpoint["model_args"])
+            state_dict = checkpoint["model"]
+            unwanted_prefix = "_orig_mod."
+            for k, v in list(state_dict.items()):
+                if k.startswith(unwanted_prefix):
+                    state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)
+            model.load_state_dict(state_dict)
+        elif init_from.startswith("gpt2"):
+            # init from a given GPT-2 model
+            model = GPT.from_pretrained(init_from, dict(dropout=0.0))
+
+        model.eval()
+        model.to(device)
+        if compile:
+            model = torch.compile(model)  # requires PyTorch 2.0 (optional)
+
+        # look for the meta pickle in case it is available in the dataset folder
+        meta_path = os.path.join(BASE_DIR, "out", "meta.pkl")
+        load_meta = os.path.exists(meta_path)
+        if move_num_in_gamestate and load_meta:
+            with open(meta_path, "rb") as f:
+                meta = pickle.load(f)
+            stoi, itos = meta["stoi"], meta["itos"]
+            vocab_size = meta['vocab_size']
+            encode = lambda s: [stoi[c] for c in s]
+            decode = lambda l: "".join([itos[i] for i in l])
+        else:
+            stoi = {' ': 0, '.': 1, 'a': 2, 'b': 3, 'c': 4, 'd': 5, 'e': 6, 'f': 7, 'g': 8, 'h': 9, '1': 10, '2': 11, '3': 12, '4': 13, '5': 14, '6': 15, '7': 16, '8': 17, 'B': 18, 'N': 19, 'R': 20, 'Q': 21, 'K': 22, 'O': 23, 'x': 24, '+': 25, '#': 26, '=': 27}
+            itos = {0: ' ', 1: '.', 2: 'a', 3: 'b', 4: 'c', 5: 'd', 6: 'e', 7: 'f', 8: 'g', 9: 'h', 10: '1', 11: '2', 12: '3', 13: '4', 14: '5', 15: '6', 16: '7', 17: '8', 18: 'B', 19: 'N', 20: 'R', 21: 'Q', 22: 'K', 23: 'O', 24: 'x', 25: '+', 26: '#', 27: '='}
+            for s in stoi:
+                assert itos[stoi[s]] == s
+            vocab_size = len(stoi)
+            print(f"Vocab size {vocab_size}")
+            encode = lambda s: [stoi[c] for c in s.replace('-', '')]
+            decode = lambda l: "".join([itos[i] for i in l]).replace("OOO", "O-O-O").replace("OO", "O-O")
+
+        self.encode = encode
+        self.decode = decode
+        self.model = model
+        self.ctx = ctx
+        self.device = device
+
+    def get_nanogpt_response(self, game_state: str, temperature: float) -> str:
+        num_samples = 1  # number of samples to draw
+        top_k = 200  # retain only the top_k most likely tokens, clamp others to have 0 probability
+        max_new_tokens = 8
+
+        # Remove ["stockfish elo xxx"]\n["stockfish elo xxx"]\n\n from game_state
+        # nanogpt was trained only on pgn transcripts
+        game_state = game_state.split("\n\n")[-1].strip()
+
+        # print("game_state", game_state)
+
+        #game_state = ";" + game_state
+
+        start_ids = self.encode(game_state)
+
+        x = torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...]
+        with torch.no_grad():
+            with self.ctx:
+                for k in range(num_samples):
+                    y = self.model.generate(
+                        x, max_new_tokens, temperature=temperature, top_k=top_k
+                    )
+
+                    model_response = self.decode(y[0].tolist())
+
+        # print("model_response", model_response)
+        # model_response includes the input string
+        model_response = model_response[len(game_state):].split(";")[0]
+        return model_response
+
+    def get_move_from_response(self, response: str) -> str:
+        # Parse the response to get only the first move
+        moves = response.split()
+        first_move = moves[0]
+
+        return first_move
+
+    def get_move(self, board: str, game_state: str, temperature: float) -> str:
+        completion = self.get_nanogpt_response(game_state, temperature)
+        return self.get_move_from_response(completion)
+
+    def get_config(self) -> dict:
+        return {"model": self.model_name}

chess-gpt-eval/nanogpt/out/meta.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1121191e401988851de5744fe27fe463c3a086fc8c9a5538ef7fc12162bfb09
+size 373

chess-gpt-eval/nanogpt/out/view_ckpt.ipynb

@@ -0,0 +1,61 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "\n",
+    "def load_checkpoint(filepath: str) -> dict:\n",
+    "    \"\"\"\n",
+    "    Load a checkpoint file.\n",
+    "\n",
+    "    Args:\n",
+    "    filepath (str): Path to the .ckpt file.\n",
+    "\n",
+    "    Returns:\n",
+    "    dict: Contents of the checkpoint file.\n",
+    "    \"\"\"\n",
+    "    checkpoint = torch.load(filepath, map_location=torch.device('cpu'))\n",
+    "    return checkpoint\n",
+    "\n",
+    "checkpoint_path = 'ckpt.pt'\n",
+    "checkpoint_data = load_checkpoint(checkpoint_path)\n",
+    "\n",
+    "# Print the keys to understand what's inside\n",
+    "print(checkpoint_data.keys())\n",
+    "\n",
+    "# If you want to view specific information, access it using the keys\n",
+    "# For example, to view the model's state_dict\n",
+    "model_state = checkpoint_data.get('state_dict', None)\n",
+    "if model_state:\n",
+    "    print(\"Model's state dict:\", model_state)\n",
+    "\n",
+    "# To view training information like current learning rate, iterations, etc.\n",
+    "training_info = checkpoint_data.get('training_info', None)\n",
+    "if training_info:\n",
+    "    print(\"Training Info:\", training_info)\n",
+    "\n",
+    "# To view config, if it's stored in the checkpoint\n",
+    "config = checkpoint_data.get('config', None)\n",
+    "if config:\n",
+    "    print(\"Configurations:\", config)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "openai",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

chess-gpt-eval/pscan.py

@@ -0,0 +1,226 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+"""
+
+An implementation of the parallel scan operation in PyTorch (Blelloch version).
+Please see docs/pscan.ipynb for a detailed explanation of what happens here.
+
+"""
+
+def npo2(len):
+    """
+    Returns the next power of 2 above len
+    """
+
+    return 2 ** math.ceil(math.log2(len))
+
+def pad_npo2(X):
+    """
+    Pads input length dim to the next power of 2
+
+    Args:
+        X : (B, L, D, N)
+
+    Returns:
+        Y : (B, npo2(L), D, N)
+    """
+
+    len_npo2 = npo2(X.size(1))
+    pad_tuple = (0, 0, 0, 0, 0, len_npo2 - X.size(1))
+    return F.pad(X, pad_tuple, "constant", 0)
+
+class PScan(torch.autograd.Function):
+    @staticmethod
+    def pscan(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # modifies X in place by doing a parallel scan.
+        # more formally, X will be populated by these values :
+        # H[t] = A[t] * H[t-1] + X[t] with H[0] = 0
+        # which are computed in parallel (2*log2(T) sequential steps (ideally), instead of T sequential steps)
+
+        # only supports L that is a power of two (mainly for a clearer code)
+        
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+            
+            Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0]))
+            Aa[:, :, :, 1].mul_(Aa[:, :, :, 0])
+
+            Aa = Aa[:, :, :, 1]
+            Xa = Xa[:, :, :, 1]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            Aa[:, :, 1].mul_(Aa[:, :, 0])
+
+            Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1])))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa = X[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1]))
+        Aa[:, :, 2].mul_(Aa[:, :, 1])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 2**k-1:L:2**k]
+            Xa = X[:, :, 2**k-1:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1]))
+            Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1])
+
+    @staticmethod
+    def pscan_rev(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # the same function as above, but in reverse
+        # (if you flip the input, call pscan, then flip the output, you get what this function outputs)
+        # it is used in the backward pass
+
+        # only supports L that is a power of two (mainly for a clearer code)
+
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+                    
+            Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1]))
+            Aa[:, :, :, 0].mul_(Aa[:, :, :, 1])
+
+            Aa = Aa[:, :, :, 0]
+            Xa = Xa[:, :, :, 0]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3]))
+            Aa[:, :, 2].mul_(Aa[:, :, 3])
+
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2]))))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 0:L:2**(num_steps-2)]
+        Xa = X[:, :, 0:L:2**(num_steps-2)]
+        Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2]))
+        Aa[:, :, 1].mul_(Aa[:, :, 2])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 0:L:2**k]
+            Xa = X[:, :, 0:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0]))
+            Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0])
+
+    @staticmethod
+    def forward(ctx, A_in, X_in):
+        """
+        Applies the parallel scan operation, as defined above. Returns a new tensor.
+        If you can, privilege sequence lengths that are powers of two.
+
+        Args:
+            A_in : (B, L, D, N)
+            X_in : (B, L, D, N)
+
+        Returns:
+            H : (B, L, D, N)
+        """
+
+        L = X_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            A = A_in.clone()
+            X = X_in.clone()
+        else:
+            # pad tensors (and clone btw)
+            A = pad_npo2(A_in) # (B, npo2(L), D, N)
+            X = pad_npo2(X_in) # (B, npo2(L), D, N)
+        
+        # prepare tensors
+        A = A.transpose(2, 1) # (B, D, npo2(L), N)
+        X = X.transpose(2, 1) # (B, D, npo2(L), N)
+
+        # parallel scan (modifies X in-place)
+        PScan.pscan(A, X)
+
+        ctx.save_for_backward(A_in, X)
+        
+        # slice [:, :L] (cut if there was padding)
+        return X.transpose(2, 1)[:, :L]
+    
+    @staticmethod
+    def backward(ctx, grad_output_in):
+        """
+        Flows the gradient from the output to the input. Returns two new tensors.
+
+        Args:
+            ctx : A_in : (B, L, D, N), X : (B, D, L, N)
+            grad_output_in : (B, L, D, N)
+
+        Returns:
+            gradA : (B, L, D, N), gradX : (B, L, D, N)
+        """
+
+        A_in, X = ctx.saved_tensors
+
+        L = grad_output_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            grad_output = grad_output_in.clone()
+            # the next padding will clone A_in
+        else:
+            grad_output = pad_npo2(grad_output_in) # (B, npo2(L), D, N)
+            A_in = pad_npo2(A_in) # (B, npo2(L), D, N)
+
+        # prepare tensors
+        grad_output = grad_output.transpose(2, 1)
+        A_in = A_in.transpose(2, 1) # (B, D, npo2(L), N)
+        A = torch.nn.functional.pad(A_in[:, :, 1:], (0, 0, 0, 1)) # (B, D, npo2(L), N) shift 1 to the left (see hand derivation)
+
+        # reverse parallel scan (modifies grad_output in-place)
+        PScan.pscan_rev(A, grad_output)
+
+        Q = torch.zeros_like(X)
+        Q[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:])
+
+        return Q.transpose(2, 1)[:, :L], grad_output.transpose(2, 1)[:, :L]
+    
+pscan = PScan.apply

chess-gpt-eval/requirements.txt

@@ -0,0 +1,6 @@
+openai==0.28.0
+tiktoken==0.4.0
+tenacity==8.2.3
+python-chess==1.999
+matplotlib==3.8.0
+pandas==2.1.1

chess-gpt-eval/xformer.py

@@ -0,0 +1,330 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters
+        #print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {} # default to empty dict
+        # only dropout can be overridden see more notes below
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        config_args['bias'] = True # always True for GPT model checkpoints
+        # we can override the dropout rate, if desired
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
+        # this means that we have to transpose these weights when we import them
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx

chess-mamba-vs-xformer/config/Mamba/11M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 256 
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 200
+eval_iters = 33
+log_interval = 33
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 500 # not super necessary potentially
+learning_rate = 1e-3 
+min_lr = 6.6667e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 400000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Mamba/11M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-mamba-v2'
+wandb_run_name = 'Mamba-11M'
+
+dataset = 'stable'
+
+# 11M param
+model_type = 'mamba'
+n_layer = 20
+d_model = 288
+d_state = 16
+dt_rank = 'auto'  #ceil(d_model/16) ... 18 here
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Mamba/250M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 256 
+
+batch_size = 10
+gradient_accumulation_steps = 10
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 300
+eval_iters = 33
+log_interval = 75
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 500 # not super necessary potentially
+learning_rate = 2.0e-3 # tested 1.5e-3 from 112k-156k, before that 3.5e-3  #8e-3 
+min_lr = 1.3333e-4
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 400000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Mamba/250M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True
+wandb_project = 'chess-mamba-v2'
+wandb_run_name = 'Mamba-250M'
+
+dataset = 'stable'
+
+# 251M param
+model_type = 'mamba'
+n_layer = 96
+d_model = 578
+d_state = 56
+dt_rank = 176
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Mamba/29M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 256 
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 250
+eval_iters = 33
+log_interval = 50
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 500 # not super necessary potentially
+learning_rate = 1.25e-3 
+min_lr = 8.3333e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 400000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Mamba/29M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-mamba-v2'
+wandb_run_name = 'Mamba-29M'
+
+dataset = 'stable'
+
+# 29.3M param
+model_type = 'mamba'
+n_layer = 33
+d_model = 360
+d_state = 24
+dt_rank = 36
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Mamba/50M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 256 
+
+batch_size = 50
+gradient_accumulation_steps = 2
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 250
+eval_iters = 33
+log_interval = 50
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 500 # not super necessary potentially
+learning_rate = 1.5e-3 # tested 1.5e-3 from 112k-156k, before that 3.5e-3  #8e-3 
+min_lr = 1.0e-4 # was planning 8.5e-5 w/ /6.75 anneal  #... before 2e-4  # 4.75e-4
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 400000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Mamba/50M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True
+wandb_project = 'chess-mamba-v2'
+wandb_run_name = 'Mamba-50M'
+
+dataset = 'stable'
+
+# 50.4M param
+model_type = 'mamba'
+n_layer = 48
+d_model = 384
+d_state = 32
+dt_rank = 56
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Mamba/6.6M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 256 
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 200
+eval_iters = 33
+log_interval = 33
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 500 # not super necessary potentially
+learning_rate = 8.16667e-4 
+min_lr = 5.4444e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 400000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Mamba/6.6M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-mamba-v2'
+wandb_run_name = 'Mamba-6.6M'
+
+dataset = 'stable'
+
+# 6.6M param
+model_type = 'mamba'
+n_layer = 15
+d_model = 256
+d_state = 16
+dt_rank = 'auto'  #ceil(d_model/16) ... 16 here
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Xformer/11M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 100 
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 600
+eval_iters = 100
+log_interval = 100
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 1280 # not super necessary potentially
+learning_rate = 2e-4
+min_lr = 1.33333e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 1024000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Xformer/11M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-xformer'
+wandb_run_name = 'Xformer-11M'
+
+dataset = 'stable'
+
+# 11.2M param
+model_type = 'xformer'
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.0
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Xformer/250M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 100
+
+batch_size = 10
+gradient_accumulation_steps = 10
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 900
+eval_iters = 100
+log_interval = 225
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 1280 # not super necessary potentially
+learning_rate = 4e-4
+min_lr = 2.6667e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 1024000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Xformer/250M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True
+wandb_project = 'chess-xformer'
+wandb_run_name = 'Xformer-250M'
+
+dataset = 'stable'
+
+# 251.2M param
+model_type = 'xformer'
+n_layer = 51
+n_head = 16
+n_embd = 640
+dropout = 0.0
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Xformer/29M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 100
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 750
+eval_iters = 100
+log_interval = 150
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 1280 # not super necessary potentially
+learning_rate = 2.5e-4 
+min_lr = 1.6667e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 1024000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Xformer/29M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-xformer'
+wandb_run_name = 'Xformer-29M'
+
+dataset = 'stable'
+
+# 29.1M param
+model_type = 'xformer'
+n_layer = 9
+n_head = 8
+n_embd = 512
+dropout = 0.0
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Xformer/50M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 100
+
+batch_size = 50
+gradient_accumulation_steps = 2
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 750
+eval_iters = 100
+log_interval = 150
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 1280 # not super necessary potentially
+learning_rate = 3e-4 # Mamba is 9.375e-4 (adjusting for different base_batch_size)
+min_lr = 2e-5 # Same ratio min/max as w/ Mamba. It's lower than 1/10 because doing slightly long anneal.
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 1024000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Xformer/50M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True
+wandb_project = 'chess-xformer'
+wandb_run_name = 'Xformer-50M'
+
+dataset = 'stable'
+
+# 50.8M param
+model_type = 'xformer'
+n_layer = 16
+n_head = 8
+n_embd = 512
+dropout = 0.0
+move_num_in_gamestate = False
+
+init_from = 'scratch'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/config/Xformer/6.6M.py

@@ -0,0 +1,70 @@
+import numpy as np
+import math
+
+beta1 = 0.9
+beta2 = 0.95
+weight_decay = 4.5e-3
+grad_clip = 0.5
+auto_clip = True
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 400
+grad_clip_percentile = 10  #7.5 (try it at 10, tested @7.75)
+max_seq_len = 1536
+
+# batch size below values are based on this. When actual batch size adjusted, the below are adjusted automatically
+base_batch_size = 100
+
+batch_size = 100
+gradient_accumulation_steps = 1
+effective_batch_size = batch_size * gradient_accumulation_steps
+
+always_save_checkpoint = True
+eval_interval = 600
+eval_iters = 100
+log_interval = 100
+train_file_update_interval = 10 # 23 was original ... 7 definitely crashes (maybe try 10 on Lambda)
+
+warmup_iters = 1280 # not super necessary potentially
+learning_rate = 1.633333e-4 
+min_lr = 1.08889e-5
+# max_iters is for auto-stopping end of stable phase. Reported %complete progress is wrt this (that is, % complete doesn't include anneal).
+max_iters = 1024000  #~=102M games
+
+# # # # #
+
+warmup_iters = int(warmup_iters * (base_batch_size / effective_batch_size))
+learning_rate = learning_rate * np.sqrt(effective_batch_size / base_batch_size) # with baby networks can afford to go a bit higher
+max_iters = int(max_iters * (base_batch_size / effective_batch_size))
+min_lr = min_lr * np.sqrt(effective_batch_size / base_batch_size) # learning_rate / 10 usually
+
+out_dir = 'out/Xformer/6.6M'
+eval_interval = int(eval_interval * (base_batch_size / effective_batch_size)) # keep frequent because we'll overfit
+eval_iters = int(eval_iters * (base_batch_size / batch_size)) # intentionally scaled by batch_size instead of effective_batch_size
+log_interval = int(math.ceil(log_interval * (base_batch_size / effective_batch_size))) # don't print too too often
+
+print(f'warmup iters: {warmup_iters}')
+print(f'Max iters: {max_iters} ({max_iters * effective_batch_size} games)')
+print(f'Eval iters: {eval_iters}')
+print(f'Eval interval: {eval_interval}')
+print(f'Log interval: {log_interval}')
+
+wandb_log = True # override via command line if you like
+wandb_project = 'chess-xformer'
+wandb_run_name = 'Xformer-6.6M'
+
+dataset = 'stable'
+
+# 6.6M param
+model_type = 'xformer'
+n_layer = 5
+n_head = 5
+n_embd = 320
+dropout = 0.0
+move_num_in_gamestate = False
+
+init_from = 'resume'
+
+device = 'cuda'  # run on cpu only
+compile = False # do not torch compile the model

chess-mamba-vs-xformer/configurator.py

@@ -0,0 +1,47 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import sys
+from ast import literal_eval
+
+for arg in sys.argv[1:]:
+    if '=' not in arg:
+        # assume it's the name of a config file
+        assert not arg.startswith('--')
+        config_file = arg
+        print(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print(f.read())
+        exec(open(config_file).read())
+    else:
+        # assume it's a --key=value argument
+        assert arg.startswith('--')
+        key, val = arg.split('=')
+        key = key[2:]
+        if key in globals():
+            try:
+                # attempt to eval it it (e.g. if bool, number, or etc)
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                # if that goes wrong, just use the string
+                attempt = val
+            # ensure the types match ok
+            assert type(attempt) == type(globals()[key])
+            # cross fingers
+            print(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")

chess-mamba-vs-xformer/data/anneal/anneal.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37569f7bafca0eb2a7361e4ae29ef1b9fed64dbeb061d2653215c348586f7a7e
+size 679959998

chess-mamba-vs-xformer/mamba.py

@@ -0,0 +1,368 @@
+import math
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pscan import pscan
+
+"""
+
+This file closely follows the mamba_simple.py from the official Mamba implementation, and the mamba-minimal by @johnma2006.
+The major differences are :
+-the convolution is done with torch.nn.Conv1d
+-the selective scan is done in PyTorch
+
+A sequential version of the selective scan is also available for comparison.
+
+- A Mamba model is composed of several layers, which are ResidualBlock.
+- A ResidualBlock is composed of a MambaBlock, a normalization, and a residual connection : ResidualBlock(x) = mamba(norm(x)) + x
+- This leaves us with the MambaBlock : its input x is (B, L, D) and its outputs y is also (B, L, D) (B=batch size, L=seq len, D=model dim).
+First, we expand x into (B, L, 2*ED) (where E is usually 2) and split it into x and z, each (B, L, ED).
+Then, we apply the short 1d conv to x, followed by an activation function (silu), then the SSM.
+We then multiply it by silu(z).
+See Figure 3 of the paper (page 8) for a visual representation of a MambaBlock.
+
+"""
+
+@dataclass
+class MambaConfig:
+    d_model: int # D
+    n_layers: int
+    dt_rank: Union[int, str] = 'auto'
+    d_state: int = 16 # N in paper/comments
+    expand_factor: int = 2 # E in paper/comments
+    d_conv: int = 4
+
+    dt_min: float = 0.001
+    dt_max: float = 0.1
+    dt_init: str = "random" # "random" or "constant"
+    dt_scale: float = 1.0
+    dt_init_floor = 1e-4
+
+    bias: bool = False
+    conv_bias: bool = True
+
+    pscan: bool = True # use parallel scan mode or sequential mode when training
+
+    def __post_init__(self):
+        self.d_inner = self.expand_factor * self.d_model # E*D = ED in comments
+
+        if self.dt_rank == 'auto':
+            self.dt_rank = math.ceil(self.d_model / 16)
+
+class Mamba(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.config = config
+
+        self.layers = nn.ModuleList([ResidualBlock(config) for _ in range(config.n_layers)])
+        #self.norm_f = RMSNorm(config.d_model)
+
+    def forward(self, x):
+        # x : (B, L, D)
+
+        # y : (B, L, D)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        #x = self.norm_f(x)
+
+        return x
+    
+    def step(self, x, caches):
+        # x : (B, L, D)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        # y : (B, L, D)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        for i, layer in enumerate(self.layers):
+            x, caches[i] = layer.step(x, caches[i])
+
+        return x, caches
+
+class ResidualBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.mixer = MambaBlock(config)
+        self.norm = RMSNorm(config.d_model)
+
+    def forward(self, x):
+        # x : (B, L, D)
+
+        # output : (B, L, D)
+
+        output = self.mixer(self.norm(x)) + x
+        return output
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs: (B, ED, d_conv-1)
+
+        # output : (B, D)
+        # cache : (h, inputs)
+
+        output, cache = self.mixer.step(self.norm(x), cache)
+        output = output + x
+        return output, cache
+
+class MambaBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        super().__init__()
+
+        self.config = config
+
+        # projects block input from D to 2*ED (two branches)
+        self.in_proj = nn.Linear(config.d_model, 2 * config.d_inner, bias=config.bias)
+
+        self.conv1d = nn.Conv1d(in_channels=config.d_inner, out_channels=config.d_inner, 
+                              kernel_size=config.d_conv, bias=config.conv_bias, 
+                              groups=config.d_inner,
+                              padding=config.d_conv - 1)
+                              
+        nn.init.kaiming_normal_(self.conv1d.weight, mode='fan_out', nonlinearity='leaky_relu')
+        
+        # projects x to input-dependent Δ, B, C
+        self.x_proj = nn.Linear(config.d_inner, config.dt_rank + 2 * config.d_state, bias=False)
+
+        # projects Δ from dt_rank to d_inner
+        self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
+
+        # dt initialization
+        # dt weights
+        dt_init_std = config.dt_rank**-0.5 * config.dt_scale
+        if config.dt_init == "constant":
+            nn.init.constant_(self.dt_proj.weight, dt_init_std)
+        elif config.dt_init == "random":
+            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
+        else:
+            raise NotImplementedError
+        
+        # dt bias
+        dt = torch.exp(
+            torch.rand(config.d_inner) * (math.log(config.dt_max) - math.log(config.dt_min)) + math.log(config.dt_min)
+        ).clamp(min=config.dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt)) # inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+        with torch.no_grad():
+            self.dt_proj.bias.copy_(inv_dt)
+        #self.dt_proj.bias._no_reinit = True # initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
+        # todo : explain why removed
+
+        # S4D real initialization
+        A = torch.arange(1, config.d_state + 1, dtype=torch.float32).repeat(config.d_inner, 1)
+        self.A_log = nn.Parameter(torch.log(A)) # why store A in log ? to keep A < 0 (cf -torch.exp(...)) ? for gradient stability ?
+        self.D = nn.Parameter(torch.ones(config.d_inner))
+
+        # projects block output from ED back to D
+        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
+
+    def forward(self, x):
+        # x : (B, L, D)
+        
+        # y : (B, L, D)
+
+        _, L, _ = x.shape
+
+        xz = self.in_proj(x) # (B, L, 2*ED)
+        x, z = xz.chunk(2, dim=-1) # (B, L, ED), (B, L, ED)
+
+        # x branch
+        x = x.transpose(1, 2) # (B, ED, L)
+        x = self.conv1d(x)[:, :, :L] # depthwise convolution over time, with a short filter
+        x = x.transpose(1, 2) # (B, L, ED)
+
+        x = F.silu(x)
+        y = self.ssm(x)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, L, D)
+
+        return output
+    
+    def ssm(self, x):
+        # x : (B, L, ED)
+
+        # y : (B, L, ED)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N)
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, L, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, L, dt_rank), (B, L, N), (B, L, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, L, ED)
+
+        if self.config.pscan:
+            y = self.selective_scan(x, delta, A, B, C, D)
+        else:
+            y = self.selective_scan_seq(x, delta, A, B, C, D)
+
+        return y
+    
+    def selective_scan(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+        
+        hs = pscan(deltaA, BX)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    def selective_scan_seq(self, x, delta, A, B, C, D):
+        # x : (B, L, ED)
+        # Δ : (B, L, ED)
+        # A : (ED, N)
+        # B : (B, L, N)
+        # C : (B, L, N)
+        # D : (ED)
+
+        # y : (B, L, ED)
+
+        _, L, _ = x.shape
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, L, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2) # (B, L, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, L, ED, N)
+
+        h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+        hs = []
+
+        for t in range(0, L):
+            h = deltaA[:, t] * h + BX[:, t]
+            hs.append(h)
+            
+        hs = torch.stack(hs, dim=1) # (B, L, ED, N)
+
+        y = (hs @ C.unsqueeze(-1)).squeeze(3) # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)
+
+        y = y + D * x
+
+        return y
+    
+    # -------------------------- inference -------------------------- #
+    """
+    Concerning auto-regressive inference
+
+    The cool part of using Mamba : inference is constant wrt to sequence length
+    We just have to keep in cache, for each layer, two things :
+    - the hidden state h (which is (B, ED, N)), as you typically would when doing inference with a RNN
+    - the last d_conv-1 inputs of the layer, to be able to compute the 1D conv which is a convolution over the time dimension
+      (d_conv is fixed so this doesn't incur a growing cache as we progress on generating the sequence)
+      (and d_conv is usually very small, like 4, so we just have to "remember" the last 3 inputs)
+
+    Concretely, these two quantities are put inside a cache tuple, and are named h and inputs respectively.
+    h is (B, ED, N), and inputs is (B, ED, d_conv-1)
+    The MambaBlock.step() receives this cache, and, along with outputing the output, alos outputs the updated cache for the next call.
+
+    The cache object is initialized as follows : (None, torch.zeros()).
+    When h is None, the selective scan function detects it and start with h=0.
+    The torch.zeros() isn't a problem (it's same as just feeding the input, because the conv1d is padded)
+
+    As we need one such cache variable per layer, we store a caches object, which is simply a list of cache object. (See mamba_lm.py)
+    """
+    
+    def step(self, x, cache):
+        # x : (B, D)
+        # cache : (h, inputs)
+                # h : (B, ED, N)
+                # inputs : (B, ED, d_conv-1)
+        
+        # y : (B, D)
+        # cache : (h, inputs)
+        
+        h, inputs = cache
+        
+        xz = self.in_proj(x) # (B, 2*ED)
+        x, z = xz.chunk(2, dim=1) # (B, ED), (B, ED)
+
+        # x branch
+        x_cache = x.unsqueeze(2)
+        x = self.conv1d(torch.cat([inputs, x_cache], dim=2))[:, :, self.config.d_conv-1] # (B, ED)
+
+        x = F.silu(x)
+        y, h = self.ssm_step(x, h)
+
+        # z branch
+        z = F.silu(z)
+
+        output = y * z
+        output = self.out_proj(output) # (B, D)
+
+        # prepare cache for next call
+        inputs = torch.cat([inputs[:, :, 1:], x_cache], dim=2) # (B, ED, d_conv-1)
+        cache = (h, inputs)
+        
+        return output, cache
+
+    def ssm_step(self, x, h):
+        # x : (B, ED)
+        # h : (B, ED, N)
+
+        # y : (B, ED)
+        # h : (B, ED, N)
+
+        A = -torch.exp(self.A_log.float()) # (ED, N) # todo : ne pas le faire tout le temps, puisque c'est indépendant de la timestep
+        D = self.D.float()
+        # TODO remove .float()
+
+        deltaBC = self.x_proj(x) # (B, dt_rank+2*N)
+
+        delta, B, C = torch.split(deltaBC, [self.config.dt_rank, self.config.d_state, self.config.d_state], dim=-1) # (B, dt_rank), (B, N), (B, N)
+        delta = F.softplus(self.dt_proj(delta)) # (B, ED)
+
+        deltaA = torch.exp(delta.unsqueeze(-1) * A) # (B, ED, N)
+        deltaB = delta.unsqueeze(-1) * B.unsqueeze(1) # (B, ED, N)
+
+        BX = deltaB * (x.unsqueeze(-1)) # (B, ED, N)
+
+        if h is None:
+            h = torch.zeros(x.size(0), self.config.d_inner, self.config.d_state, device=deltaA.device) # (B, ED, N)
+
+        h = deltaA * h + BX # (B, ED, N)
+
+        y = (h @ C.unsqueeze(-1)).squeeze(2) # (B, ED, N) @ (B, N, 1) -> (B, ED, 1)
+
+        y = y + D * x
+
+        # todo : pq h.squeeze(1) ??
+        return y, h.squeeze(1)
+
+# taken straight from https://github.com/johnma2006/mamba-minimal/blob/master/model.py
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5):
+        super().__init__()
+
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+        return output

chess-mamba-vs-xformer/mamba_lm.py

@@ -0,0 +1,168 @@
+from dataclasses import dataclass, fields, asdict
+import json
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from mamba import Mamba, MambaConfig, RMSNorm
+
+"""
+
+Encapsulates a Mamba model as language model. It has an embedding layer, and a LM head which maps the model output to logits.
+
+"""
+
+# TODO generate function : batch size != 1 ? (for now B=1)
+# TODO generate function : top-p sampling
+
+@dataclass
+class MambaLMConfig(MambaConfig):
+    vocab_size: int = 32000
+    pad_vocab_size_multiple: int = 8
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        #if self.vocab_size % self.pad_vocab_size_multiple != 0:
+        #    self.vocab_size += (self.pad_vocab_size_multiple - self.vocab_size % self.pad_vocab_size_multiple)
+
+    def to_mamba_config(self) -> MambaConfig:
+        mamba_config_fields = {field.name for field in fields(MambaConfig)}
+        filtered_dict = {k: v for k, v in asdict(self).items() if k in mamba_config_fields}
+        return MambaConfig(**filtered_dict)
+
+# adapted from https://github.com/johnma2006/mamba-minimal
+def from_pretrained(name: str):
+    """
+    Returns a model loaded with pretrained weights pulled from HuggingFace.
+
+    Args:
+        name: As of now, supports
+            * 'state-spaces/mamba-2.8b-slimpj'
+            * 'state-spaces/mamba-2.8b'
+            * 'state-spaces/mamba-1.4b'
+            * 'state-spaces/mamba-790m'
+            * 'state-spaces/mamba-370m'
+            * 'state-spaces/mamba-130m'
+
+    Returns:
+        model: a Mamba model configured with the proper parameters and initialized with the proper weights
+    """   
+
+    from transformers.utils import WEIGHTS_NAME, CONFIG_NAME
+    from transformers.utils.hub import cached_file
+
+    def load_config_hf(model_name):
+        resolved_archive_file = cached_file(model_name, CONFIG_NAME, _raise_exceptions_for_missing_entries=False)
+        return json.load(open(resolved_archive_file))
+                
+    def load_state_dict_hf(model_name):
+        resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False)
+        return torch.load(resolved_archive_file, weights_only=True, map_location='cpu', mmap=True)
+        
+    # copy config data
+    config_data = load_config_hf(name)
+    config = MambaLMConfig(d_model=config_data['d_model'], n_layers=config_data['n_layer'], vocab_size=config_data['vocab_size'])
+
+    model = MambaLM(config)
+
+    # copy weights
+    state_dict = load_state_dict_hf(name)
+
+    new_state_dict = {}
+    for key in state_dict:
+        if key == 'backbone.embedding.weight' or key == 'backbone.norm_f.weight':
+            new_key = key.replace('backbone.', '')
+        else:
+            new_key = key.replace('backbone', 'mamba')
+
+        new_state_dict[new_key] = state_dict[key]
+
+    model.load_state_dict(new_state_dict)
+
+    return model
+
+class MambaLM(nn.Module):
+    def __init__(self, lm_config: MambaLMConfig):
+        super().__init__()
+        self.lm_config = lm_config
+        self.config = lm_config.to_mamba_config()
+
+        self.embedding = nn.Embedding(self.lm_config.vocab_size, self.config.d_model)
+        self.mamba = Mamba(self.config)
+        self.norm_f = RMSNorm(self.config.d_model)
+
+        self.lm_head = nn.Linear(self.config.d_model, self.lm_config.vocab_size, bias=False)
+        self.lm_head.weight = self.embedding.weight
+
+    def forward(self, tokens):
+        # tokens : (B, L)
+
+        # logits : (B, L, vocab_size)
+
+        x = self.embedding(tokens)
+
+        x = self.mamba(x)
+        x = self.norm_f(x)
+
+        logits = self.lm_head(x)
+
+        return logits
+    
+    def step(self, token, caches):
+        # token : (B)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        # logits : (B, vocab_size)
+        # caches : [cache(layer) for all layers], cache : (h, inputs)
+
+        x = self.embedding(token)
+
+        x, caches = self.mamba.step(x, caches)
+        x = self.norm_f(x)
+
+        logits = self.lm_head(x)
+
+        return logits, caches
+    
+    # TODO temperature
+    # TODO process prompt in parallel, and pass in sequential mode when prompt is finished ?
+    def generate(self, tokenizer, prompt: str, num_tokens: int = 50, sample: bool = True, top_k: int = 40):
+        self.eval()
+
+        input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(next(self.parameters()).device) # (1, num_tokens)
+
+        # caches is a list of cache, one per layer
+        # cache is composed of : the hidden state, and the last d_conv-1 inputs
+        # the hidden state because the update is like an RNN
+        # the last d_conv-1 inputs because they are used in a 1d convolution (usually d_conv=4 so this is not large)
+        caches = [(None, torch.zeros(1, self.config.d_inner, self.config.d_conv-1, device=input_ids.device)) for _ in range(self.config.n_layers)]
+
+        for i in range(input_ids.size(1) + num_tokens - 1):
+            with torch.no_grad():
+                # forward the new output, get new cache
+                next_token_logits, caches = self.step(input_ids[:, i], caches) # (1, vocab_size), caches
+
+            # sample (no sampling when the prompt is being processed)
+            if i+1 >= input_ids.size(1):
+                probs = F.softmax(next_token_logits, dim=-1) # (1, vocab_size)
+
+                if top_k is not None:
+                    values, _ = torch.topk(probs, k=top_k) # (1, k) ordered from lowest to biggest
+                    probs[probs < values[:, -1, None]] = 0
+                    probs = probs / probs.sum(axis=1, keepdims=True)
+
+                if sample:
+                    next_token = torch.multinomial(probs, num_samples=1).squeeze(1) # (1)
+                else:
+                    next_token = torch.argmax(probs, dim=-1) # (1)
+
+                input_ids = torch.cat([input_ids, next_token.unsqueeze(1)], dim=1)
+                
+        output = [tokenizer.decode(output.tolist()) for output in input_ids][0]
+
+        self.train()
+
+        return output
+    

chess-mamba-vs-xformer/pscan.py

@@ -0,0 +1,226 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+"""
+
+An implementation of the parallel scan operation in PyTorch (Blelloch version).
+Please see docs/pscan.ipynb for a detailed explanation of what happens here.
+
+"""
+
+def npo2(len):
+    """
+    Returns the next power of 2 above len
+    """
+
+    return 2 ** math.ceil(math.log2(len))
+
+def pad_npo2(X):
+    """
+    Pads input length dim to the next power of 2
+
+    Args:
+        X : (B, L, D, N)
+
+    Returns:
+        Y : (B, npo2(L), D, N)
+    """
+
+    len_npo2 = npo2(X.size(1))
+    pad_tuple = (0, 0, 0, 0, 0, len_npo2 - X.size(1))
+    return F.pad(X, pad_tuple, "constant", 0)
+
+class PScan(torch.autograd.Function):
+    @staticmethod
+    def pscan(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # modifies X in place by doing a parallel scan.
+        # more formally, X will be populated by these values :
+        # H[t] = A[t] * H[t-1] + X[t] with H[0] = 0
+        # which are computed in parallel (2*log2(T) sequential steps (ideally), instead of T sequential steps)
+
+        # only supports L that is a power of two (mainly for a clearer code)
+        
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+            
+            Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0]))
+            Aa[:, :, :, 1].mul_(Aa[:, :, :, 0])
+
+            Aa = Aa[:, :, :, 1]
+            Xa = Xa[:, :, :, 1]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            Aa[:, :, 1].mul_(Aa[:, :, 0])
+
+            Xa[:, :, 3].add_(Aa[:, :, 3].mul(Xa[:, :, 2] + Aa[:, :, 2].mul(Xa[:, :, 1])))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 0]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa = X[:, :, 2**(num_steps-2)-1:L:2**(num_steps-2)]
+        Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 1]))
+        Aa[:, :, 2].mul_(Aa[:, :, 1])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 2**k-1:L:2**k]
+            Xa = X[:, :, 2**k-1:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, 1:, 0].add_(Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1]))
+            Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1])
+
+    @staticmethod
+    def pscan_rev(A, X):
+        # A : (B, D, L, N)
+        # X : (B, D, L, N)
+
+        # the same function as above, but in reverse
+        # (if you flip the input, call pscan, then flip the output, you get what this function outputs)
+        # it is used in the backward pass
+
+        # only supports L that is a power of two (mainly for a clearer code)
+
+        B, D, L, _ = A.size()
+        num_steps = int(math.log2(L))
+
+        # up sweep (last 2 steps unfolded)
+        Aa = A
+        Xa = X
+        for _ in range(num_steps-2):
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+                    
+            Xa[:, :, :, 0].add_(Aa[:, :, :, 0].mul(Xa[:, :, :, 1]))
+            Aa[:, :, :, 0].mul_(Aa[:, :, :, 1])
+
+            Aa = Aa[:, :, :, 0]
+            Xa = Xa[:, :, :, 0]
+
+        # we have only 4, 2 or 1 nodes left
+        if Xa.size(2) == 4:
+            Xa[:, :, 2].add_(Aa[:, :, 2].mul(Xa[:, :, 3]))
+            Aa[:, :, 2].mul_(Aa[:, :, 3])
+
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1].add(Aa[:, :, 1].mul(Xa[:, :, 2]))))
+        elif Xa.size(2) == 2:
+            Xa[:, :, 0].add_(Aa[:, :, 0].mul(Xa[:, :, 1]))
+            return
+        else:
+            return
+
+        # down sweep (first 2 steps unfolded)
+        Aa = A[:, :, 0:L:2**(num_steps-2)]
+        Xa = X[:, :, 0:L:2**(num_steps-2)]
+        Xa[:, :, 1].add_(Aa[:, :, 1].mul(Xa[:, :, 2]))
+        Aa[:, :, 1].mul_(Aa[:, :, 2])
+
+        for k in range(num_steps-3, -1, -1):
+            Aa = A[:, :, 0:L:2**k]
+            Xa = X[:, :, 0:L:2**k]
+
+            T = Xa.size(2)
+            Aa = Aa.view(B, D, T//2, 2, -1)
+            Xa = Xa.view(B, D, T//2, 2, -1)
+
+            Xa[:, :, :-1, 1].add_(Aa[:, :, :-1, 1].mul(Xa[:, :, 1:, 0]))
+            Aa[:, :, :-1, 1].mul_(Aa[:, :, 1:, 0])
+
+    @staticmethod
+    def forward(ctx, A_in, X_in):
+        """
+        Applies the parallel scan operation, as defined above. Returns a new tensor.
+        If you can, privilege sequence lengths that are powers of two.
+
+        Args:
+            A_in : (B, L, D, N)
+            X_in : (B, L, D, N)
+
+        Returns:
+            H : (B, L, D, N)
+        """
+
+        L = X_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            A = A_in.clone()
+            X = X_in.clone()
+        else:
+            # pad tensors (and clone btw)
+            A = pad_npo2(A_in) # (B, npo2(L), D, N)
+            X = pad_npo2(X_in) # (B, npo2(L), D, N)
+        
+        # prepare tensors
+        A = A.transpose(2, 1) # (B, D, npo2(L), N)
+        X = X.transpose(2, 1) # (B, D, npo2(L), N)
+
+        # parallel scan (modifies X in-place)
+        PScan.pscan(A, X)
+
+        ctx.save_for_backward(A_in, X)
+        
+        # slice [:, :L] (cut if there was padding)
+        return X.transpose(2, 1)[:, :L]
+    
+    @staticmethod
+    def backward(ctx, grad_output_in):
+        """
+        Flows the gradient from the output to the input. Returns two new tensors.
+
+        Args:
+            ctx : A_in : (B, L, D, N), X : (B, D, L, N)
+            grad_output_in : (B, L, D, N)
+
+        Returns:
+            gradA : (B, L, D, N), gradX : (B, L, D, N)
+        """
+
+        A_in, X = ctx.saved_tensors
+
+        L = grad_output_in.size(1)
+
+        # cloning is requiered because of the in-place ops
+        if L == npo2(L):
+            grad_output = grad_output_in.clone()
+            # the next padding will clone A_in
+        else:
+            grad_output = pad_npo2(grad_output_in) # (B, npo2(L), D, N)
+            A_in = pad_npo2(A_in) # (B, npo2(L), D, N)
+
+        # prepare tensors
+        grad_output = grad_output.transpose(2, 1)
+        A_in = A_in.transpose(2, 1) # (B, D, npo2(L), N)
+        A = torch.nn.functional.pad(A_in[:, :, 1:], (0, 0, 0, 1)) # (B, D, npo2(L), N) shift 1 to the left (see hand derivation)
+
+        # reverse parallel scan (modifies grad_output in-place)
+        PScan.pscan_rev(A, grad_output)
+
+        Q = torch.zeros_like(X)
+        Q[:, :, 1:].add_(X[:, :, :-1] * grad_output[:, :, 1:])
+
+        return Q.transpose(2, 1)[:, :L], grad_output.transpose(2, 1)[:, :L]
+    
+pscan = PScan.apply

chess-mamba-vs-xformer/train_bygame.py

@@ -0,0 +1,541 @@
+import os
+import time
+import math
+import pickle
+from contextlib import nullcontext
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed import init_process_group, destroy_process_group
+import pyarrow.parquet as pq
+import random
+from torch.utils.data import Dataset, DataLoader
+import glob
+
+# -----------------------------------------------------------------------------
+# default config values designed for Mamba model training
+# I/O
+out_dir = 'out'
+eval_interval = 2000
+log_interval = 1
+eval_iters = 5
+eval_only = False
+always_save_checkpoint = True
+init_from = 'resume'  # 'scratch', 'resume', 'anneal', or Mamba model name
+# wandb logging
+wandb_log = False
+wandb_project = 'mamba'
+wandb_run_name = 'mamba_run'  # modify as needed
+# data
+dataset = 'chess'  # specify your dataset
+gradient_accumulation_steps = 5 * 8
+batch_size = 12
+base_batch_size = batch_size
+effective_batch_size = batch_size
+max_seq_len = 1024 # For xformer, this is the block size
+train_file_update_interval = 7
+
+# model
+model_type = 'mamba'
+# TODO: add 'xformer' type / model paramers. move model imports to after exec() (when these values finalized)
+n_layer = 12
+d_model = 768
+dt_rank = 'auto'
+d_state = 16
+expand_factor = 2
+bias = False
+conv_bias = True
+pscan = True
+vocab_size = 32
+move_num_in_gamestate = True
+# xformer-specific params. Note that n_layer, vocab_size, move_num_in_gamestate, and bias are shared by both model types
+n_head = 12
+n_embd = 768
+dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
+
+# optimizer settings
+learning_rate = 6e-4
+max_iters = 600000  # max_iters is for auto-stopping end of stable phase
+weight_decay = 1e-1
+beta1 = 0.9
+beta2 = 0.95
+grad_clip = 0.5
+auto_clip = False
+auto_clip_max = 0.5
+auto_clip_min = 3.333e-3
+grad_clip_start_size = 100
+grad_clip_max_size = 500
+grad_clip_percentile = 10
+# learning rate decay settings
+decay_lr = True
+warmup_iters = 2000
+min_lr = 6e-5
+# DDP settings
+backend = 'nccl'
+# system
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+dtype = 'bfloat16' if torch.cuda.is_bf16_supported() else 'float32'
+compile = False  # set to True if using PyTorch 2.0
+# -----------------------------------------------------------------------------
+
+config_keys = [k for k, v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open('configurator.py').read())  # overrides from command line or config file
+config = {k: globals()[k] for k in config_keys}  # will be useful for logging
+# -----------------------------------------------------------------------------
+
+
+anneal_checkpoint = 'anneal/ckpt.pt'
+anneal_dir = os.path.join(out_dir, 'anneal/')
+anneal_start_iters = None # Set at init
+anneal_decay_iters = None # Set at init
+
+if model_type == 'mamba':
+    from mamba_lm import MambaLM, MambaLMConfig
+    model_config = MambaLMConfig(
+        d_model=d_model,
+        n_layers=n_layer,
+        dt_rank=dt_rank,
+        d_state=d_state,
+        expand_factor=expand_factor,
+        bias=bias,
+        conv_bias=conv_bias,
+        pscan=pscan,
+        vocab_size=vocab_size
+    )
+elif model_type == 'xformer':
+    from xformer import GPTConfig, GPT
+    model_config = GPTConfig(
+        n_layer=n_layer,
+        n_head=n_head,
+        n_embd=n_embd,
+        block_size=max_seq_len,
+        bias=bias,
+        vocab_size=vocab_size,
+        dropout=dropout)
+else:
+    print(f"Unknown model_type {model_type}.")
+    exit()
+
+# DDP and other initializations
+ddp = int(os.environ.get('RANK', -1)) != -1
+if ddp:
+    init_process_group(backend=backend)
+    ddp_rank = int(os.environ['RANK'])
+    ddp_local_rank = int(os.environ['LOCAL_RANK'])
+    ddp_world_size = int(os.environ['WORLD_SIZE'])
+    device = f'cuda:{ddp_local_rank}'
+    torch.cuda.set_device(device)
+    master_process = ddp_rank == 0
+    seed_offset = ddp_rank
+    assert gradient_accumulation_steps % ddp_world_size == 0
+    gradient_accumulation_steps //= ddp_world_size
+else:
+    master_process = True
+    seed_offset = 0
+    ddp_world_size = 1
+
+if master_process:
+    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(anneal_dir, exist_ok=True)
+torch.manual_seed(1337 + seed_offset)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+device_type = 'cuda' if 'cuda' in device else 'cpu'
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16}[dtype]
+ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+# poor man's data loader
+data_dir = os.path.join('data', dataset)
+current_train_file_index = 0
+train_files = glob.glob(os.path.join(data_dir, 'train*.parquet'))
+train_datasets = []
+for f in train_files:
+    dataset = pq.read_table(f).to_pandas()
+    dataset = dataset[dataset['tokenized'].apply(len) >= 8]
+    train_datasets.append(dataset)
+#val_data = pq.read_table(os.path.join(data_dir, 'val.parquet')).to_pandas()
+#val_data = val_data[val_data['tokenized'].apply(len) >= 8]
+truncated_games_count = 0
+total_games_count = 0
+games_seen = 0
+tokens_seen = 0
+tokens_seen_padded = 0
+def get_batch(split):
+    global truncated_games_count, total_games_count, current_train_file_index, tokens_seen, tokens_seen_padded
+
+    # Randomly select batch_size games
+    dataset = train_datasets[current_train_file_index] if split == 'train' else None # else val_data  # Use the correct DataFrame based on the split
+    sample_df = dataset.sample(batch_size)
+    games = sample_df['tokenized'].tolist()
+
+    # Prepare sequences tensor for the batch
+    max_length_in_batch = min(max(len(game) for game in games), max_seq_len)
+    pad_to = max_length_in_batch #if model_type == 'mamba' else max_seq_len
+    sequences = torch.zeros((batch_size, pad_to), dtype=torch.int64)
+
+    for i, game in enumerate(games):
+        total_games_count += 1
+        game_len = min(len(game), pad_to)
+        tokens_seen += game_len
+        tokens_seen_padded += pad_to
+        sequences[i, :game_len] = torch.tensor(game[:game_len], dtype=torch.int64)
+            
+    if (total_games_count // batch_size) % train_file_update_interval == 0:
+        current_train_file_index = random.randint(0, len(train_files) - 1)
+        # print(f"Switched to file: {train_files[current_train_file_index]}")
+            
+    if device_type == 'cuda':
+        sequences = sequences.pin_memory().to(device, non_blocking=True)
+    else:
+        sequences = sequences.to(device)
+
+    return sequences, max_length_in_batch
+    
+# init these up here, can override if init_from='resume' (i.e. from a checkpoint)
+iter_num = 0
+best_val_loss = 1e9
+
+# attempt to derive vocab_size from the dataset
+meta_path = os.path.join(data_dir, 'meta.pkl')
+meta_vocab_size = None
+if not move_num_in_gamestate:
+    meta_vocab_size = 28
+elif os.path.exists(meta_path):
+    with open(meta_path, 'rb') as f:
+        meta = pickle.load(f)
+    meta_vocab_size = meta['vocab_size']
+    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")
+    
+# Model initialization
+if init_from == 'scratch':
+    print(f"Initializing a new {model_type} model from scratch")
+    if meta_vocab_size is None:
+        print(f"defaulting to vocab_size of {vocab_size}")
+    else:
+        model_config.vocab_size = meta_vocab_size
+    if model_type == 'mamba':
+        model = MambaLM(model_config)
+    else:
+        model = GPT(model_config)
+    if auto_clip:
+        grad_clip = 0
+        config['grad_clip'] = 0
+        grad_norm_history = []
+elif init_from == 'resume' or init_from == 'anneal':
+    print(f"Resuming training from {out_dir}")
+    if init_from == 'anneal':
+        ckpt_path = os.path.join(out_dir, anneal_checkpoint)
+    else:
+        ckpt_path = os.path.join(out_dir, 'ckpt.pt')
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    model_config = checkpoint['model_args']
+    if model_type == 'mamba':
+        model = MambaLM(model_config)
+    else:
+        model = GPT(model_config)
+    state_dict = checkpoint['model']
+    # fix the keys of the state dictionary :(
+    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
+    unwanted_prefix = '_orig_mod.'
+    for k,v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    model.load_state_dict(state_dict)
+    if 'effective_batch_size' not in checkpoint['config']:
+        print("Checkpoint was saved without `effective_batch_size`, assuming current value (will save with next checkpoint). This is used for correcting `iter_num` when the effetive batch size is changed.")
+        checkpoint['config']['effective_batch_size'] = effective_batch_size
+    iter_num = int(round(checkpoint['iter_num'] * (checkpoint['config']['effective_batch_size'] / effective_batch_size)))
+    if 'games_seen' in checkpoint:
+        games_seen = checkpoint['games_seen']
+    else:
+        games_seen = checkpoint['config']['effective_batch_size'] * checkpoint['iter_num']
+        checkpoint['games_seen'] = games_seen
+        print(f"Checkpoint was saved without `games_seen`, assuming checkpoint's effective batch size * iters (will save with next checkpoint). {games_seen}")
+    tokens_seen = checkpoint.get('tokens_seen', 0)
+    tokens_seen_padded = checkpoint.get('tokens_seen_padded', 0)
+    best_val_loss = checkpoint['best_val_loss']
+    print(f"Best val loss: {best_val_loss}")
+    if auto_clip:
+        grad_clip = checkpoint['config']['grad_clip']
+        config['grad_clip'] = grad_clip
+        #grad_norm_history = [t.item() if torch.is_tensor(t) else t for t in checkpoint.get('grad_norm_history', [])]
+        grad_norm_history = checkpoint.get('grad_norm_history', [])
+    if init_from == 'anneal':
+        print(f"\n\nANNEAL STARTING/RESUMING FROM ITERNUM: {iter_num} ({games_seen} games)\n\n")
+        anneal_start_iters = iter_num if 'anneal_start_iters' not in checkpoint else checkpoint['anneal_start_iters']
+        anneal_decay_iters = iter_num / 8 if 'anneal_decay_iters' not in checkpoint else checkpoint['anneal_decay_iters'] # / 9 is og, but going deeper on lr too (can always take earlier ckpt during anneal if it doesn't keep improving)... have used 6.75
+        print(anneal_start_iters)
+        print(anneal_decay_iters)
+        if 'anneal_start_iters' not in checkpoint:
+            grad_clip = 0
+            config['grad_clip'] = 0
+            grad_norm_history = []
+            print(f"Starting anneal. Resumed from {anneal_checkpoint}, will now decay learning rate for {anneal_decay_iters} / until iter_num {anneal_start_iters + anneal_decay_iters}.")
+        out_dir = anneal_dir
+        weight_decay = weight_decay / 12.5  # / 17.0
+        beta2 = np.sqrt(beta2) * beta2
+        auto_clip = True
+        grad_clip_percentile = 6.75
+elif init_from.startswith('state-spaces'):
+    print(f"Initializing from Mamba pre-trained weights: {init_from}")
+    model = from_pretrained(init_from)
+    model_config = model.config
+else:
+    raise ValueError("Invalid init_from value")
+
+model.to(device)
+
+print(f'Model with {sum([p.numel() for p in model.parameters()])} parameters loaded.')
+
+# Optimizer and GradScaler
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2))
+scaler = torch.cuda.amp.GradScaler(enabled=dtype == 'float16')
+if init_from == 'resume':
+    optimizer.load_state_dict(checkpoint['optimizer'])
+checkpoint = None
+
+# Compile the model if using PyTorch 2.0
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    model = torch.compile(model)
+
+# Wrap model in DDP container if necessary
+if ddp:
+    model = DDP(model, device_ids=[ddp_local_rank])
+            
+            
+def batch_to_loss(sequences, max_length_in_batch):
+    if model_type == 'mamba':
+        logits = model(sequences[:, :-1])  # Forward pass, exclude last token for input
+        # Compute loss (assuming next token prediction task)
+        targets = sequences[:, 1:].reshape(-1)  # Shifted by one for next token prediction
+        return F.cross_entropy(logits.view(-1, logits.size(-1)), targets)
+    else:
+        inputs = sequences[:, :-1]
+        targets = sequences[:, 1:].reshape(-1)
+        _, loss = model(inputs, targets)
+        return loss
+    
+
+@torch.no_grad()
+def estimate_loss():
+    global tokens_seen, tokens_seen_padded
+    out = {}
+    model.eval()
+    tokens_seen_b4 = tokens_seen
+    tokens_seen_padded_b4 = tokens_seen_padded
+    for split in ['train']: #['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            loss = batch_to_loss(*get_batch(split))
+            losses[k] = loss.item()
+
+        split = 'val' # Temporary hack
+        out[split] = losses.mean()
+    tokens_seen = tokens_seen_b4
+    tokens_seen_padded = tokens_seen_padded_b4
+    model.train()
+    return out
+
+
+# WSD scheduler
+def get_lr(it):
+    if init_from == 'anneal':
+        # Linear decay from max LR to min LR over (anneal_start_iters / 9) iters
+        decay_ratio = min(it - anneal_start_iters, anneal_decay_iters) / anneal_decay_iters
+        return learning_rate - decay_ratio * (learning_rate - min_lr)
+
+    if it < warmup_iters:
+        # Warmup
+        return learning_rate * it / warmup_iters
+    
+    # Stable max LR
+    return learning_rate
+
+# Logging setup
+if wandb_log and master_process:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+# Training loop
+local_iter_num = 0  # Number of iterations in the lifetime of this process
+last_crossed_multiple = 0 
+save_every_n_games = 150000
+raw_model = model.module if ddp else model  # Unwrap DDP container if needed
+
+# initial save
+if init_from == 'scratch':
+    checkpoint = {
+        'model': raw_model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'model_args': model_config,
+        'iter_num': 0,
+        "games_seen": 0,
+        "tokens_seen": 0,
+        "tokens_seen_padded": 0,
+        'best_val_loss': best_val_loss,
+        'config': config,
+    }
+    checkpoint['grad_norm_history'] = grad_norm_history
+    print(f"saving checkpoint to {out_dir}\n")
+    torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+
+t0 = time.time()
+while True:
+    # Determine and set the learning rate for this iteration
+    lr = get_lr(iter_num) if decay_lr else learning_rate
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    # Evaluate the loss on train/val sets and write checkpoints
+    if iter_num % eval_interval == 0 and master_process and local_iter_num > 0:
+        torch.cuda.empty_cache()
+        losses = estimate_loss()
+        print(f"\ngame {games_seen} ({iter_num}, {(iter_num / max_iters)*100.0:.3f}%): 'val' loss {losses['val']:.4f}")
+        if auto_clip and len(grad_norm_history) >= grad_clip_start_size:
+            grad_clip_prev = grad_clip
+            grad_clip = np.percentile(grad_norm_history, grad_clip_percentile)
+            grad_clip = max(min(grad_clip, auto_clip_max), auto_clip_min)
+            # Transition between grad_clips smoothly, weighed to new value
+            grad_clip = (grad_clip*9.0 + grad_clip_prev*4.0) / 13.0
+            grad_clip = max(min(grad_clip, auto_clip_max), auto_clip_min)  # should never actually clip here
+            config['grad_clip'] = grad_clip
+            print(f"Auto adjusted grad_clip to {grad_clip}")
+        torch.cuda.empty_cache()
+        if wandb_log:
+            wandb.log({
+                "etc/iter": iter_num,
+                "etc/games": games_seen,
+                "etc/tokens_seen": tokens_seen,
+                "etc/tokens_seen_padded": tokens_seen_padded,
+                "etc/grad_clip": grad_clip,
+                "etc/lr": lr,
+                "val/loss": losses['val'],
+                
+            })
+        if losses['val'] < best_val_loss or always_save_checkpoint:
+            if iter_num > 0:
+                checkpoint = {
+                    'model': raw_model.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'model_args': model_config,
+                    'iter_num': iter_num,
+                    "games_seen": games_seen,
+                    "tokens_seen": tokens_seen,
+                    "tokens_seen_padded": tokens_seen_padded,
+                    'best_val_loss': min(best_val_loss, losses['val']),
+                    'config': config,
+                }
+                checkpoint['grad_norm_history'] = grad_norm_history
+                if init_from == 'anneal':
+                    checkpoint['anneal_start_iters'] = anneal_start_iters
+                    checkpoint['anneal_decay_iters'] = anneal_decay_iters
+                print(f"saving checkpoint to {out_dir}\n")
+                torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+                current_nearest_multiple = (games_seen // save_every_n_games) * save_every_n_games
+                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: # 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)}.pt'))
+                
+    if iter_num == 0 and eval_only:
+        break
+
+    # Forward and backward pass
+    for micro_step in range(gradient_accumulation_steps):
+        if ddp:
+            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
+
+        sequences, max_length_in_batch = get_batch('train')  # Fetch the training data
+        with ctx:
+            loss = batch_to_loss(sequences, max_length_in_batch)
+            loss = loss / gradient_accumulation_steps
+
+        scaler.scale(loss).backward()
+        #print('.', end='')
+
+    # clip the gradient
+    if grad_clip != 0.0 or auto_clip:
+        scaler.unscale_(optimizer)
+        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip if grad_clip != 0.0 else 999.9) # The 0 check is for auto_clip enabled but not enough history
+        grad_norm_history.append(total_norm.item())
+        grad_norm_history = grad_norm_history[-grad_clip_max_size:]
+        
+    # step the optimizer and scaler if training in fp16
+    scaler.step(optimizer)
+    scaler.update()
+    # flush the gradients as soon as we can, no need for this memory anymore
+    optimizer.zero_grad(set_to_none=True)
+    torch.cuda.empty_cache()
+    
+    # timing and logging
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        # get loss as float. note: this is a CPU-GPU sync point
+        # scale up to undo the division above, approximating the true total loss (exact would have been a sum)
+        lossf = loss.item() * gradient_accumulation_steps
+        print(f"game {games_seen} ({iter_num}, {(iter_num / max_iters)*100.0:.3f}%): loss {lossf:.4f}, time {dt*1000:.2f}ms")
+        if wandb_log:
+            wandb.log({
+                "etc/iter": iter_num,
+                "etc/games": games_seen,
+                "etc/tokens_seen": tokens_seen,
+                "etc/tokens_seen_padded": tokens_seen_padded,
+                "etc/grad_norm": grad_norm_history[-1] if grad_norm_history else 0,
+                "etc/lr": lr,
+                "train/loss": lossf,
+            })
+    iter_num += 1
+    local_iter_num += 1
+    games_seen += effective_batch_size
+
+    # termination conditions
+    if iter_num > max_iters and not init_from == 'anneal':  # max iters is for auto-stopping end of stable phase
+        checkpoint = {
+            'model': raw_model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'model_args': model_config,
+            'iter_num': iter_num,
+            "games_seen": games_seen,
+            "tokens_seen": tokens_seen,
+            "tokens_seen_padded": tokens_seen_padded,
+            'best_val_loss': best_val_loss,
+            'config': config,
+        }
+        checkpoint['grad_norm_history'] = grad_norm_history
+        if init_from == 'anneal':
+            checkpoint['anneal_start_iters'] = anneal_start_iters
+            checkpoint['anneal_decay_iters'] = anneal_decay_iters
+        print(f"Max_iters reached. Saving pre-anneal checkpoint to {anneal_checkpoint}")
+        torch.save(checkpoint, os.path.join(out_dir, anneal_checkpoint))
+        break
+    if init_from == 'anneal' and iter_num >= anneal_start_iters + anneal_decay_iters:
+        checkpoint = {
+            'model': raw_model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'model_args': model_config,
+            'iter_num': iter_num,
+            "games_seen": games_seen,
+            "tokens_seen": tokens_seen,
+            "tokens_seen_padded": tokens_seen_padded,
+            'best_val_loss': best_val_loss,
+            'config': config,
+        }
+        checkpoint['grad_norm_history'] = grad_norm_history
+        if init_from == 'anneal':
+            checkpoint['anneal_start_iters'] = anneal_start_iters
+            checkpoint['anneal_decay_iters'] = anneal_decay_iters
+        print(f"Anneal complete. Saving checkpoint to {out_dir}")
+        torch.save(checkpoint, os.path.join(out_dir, 'anneal_complete.pt'))
+        break
+
+    
+    
+if ddp:
+    destroy_process_group()
+

chess-mamba-vs-xformer/xformer.py

@@ -0,0 +1,330 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters
+        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {} # default to empty dict
+        # only dropout can be overridden see more notes below
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        config_args['bias'] = True # always True for GPT model checkpoints
+        # we can override the dropout rate, if desired
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
+        # this means that we have to transpose these weights when we import them
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx