Upload 8 files

anneal_complete.pt

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

mamba_module.py

@@ -0,0 +1,143 @@
+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):
+        self.model_name = model_name
+        # -----------------------------------------------------------------------------
+
+        init_from = "resume"  # either 'resume' or a Mamba variant (e.g. 'state-spaces/mamba-1.4b')
+        move_num_in_gamestate = True
+        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.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, 32)
+        return self.get_move_from_response(completion)
+
+    def get_config(self) -> dict:
+        return {"model": self.model_name}
+

train_bygame.py

@@ -0,0 +1,486 @@
+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
+from mamba_lm import MambaLM, MambaLMConfig
+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 # A trianing-only parameter for controlling VRAM
+train_file_update_interval = 7
+
+# model
+n_layer = 12
+d_model = 768
+dt_rank = 'auto'
+d_state = 16
+expand_factor = 2
+bias = False
+conv_bias = True
+pscan = True
+vocab_size = 32000
+move_num_in_gamestate = True
+
+# optimizer settings
+learning_rate = 6e-4
+max_iters = 600000
+weight_decay = 1e-1
+beta1 = 0.9
+beta2 = 0.95
+grad_clip = 1.0
+auto_clip = False    
+grad_clip_start_size = 100
+grad_clip_max_size = 500
+grad_clip_percentile = 10
+# learning rate decay settings
+decay_lr = True
+warmup_iters = 2000
+lr_decay_iters = 600000
+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_me.pt'  
+anneal_dir = os.path.join(out_dir, 'anneal/')
+anneal_start_iters = None # Set at init
+anneal_decay_iters = None # Set at init
+
+mamba_config = MambaLMConfig(
+    d_model=d_model,  # adjust as needed
+    n_layers=n_layer,  # adjust as needed
+    dt_rank=dt_rank,
+    d_state=d_state,
+    expand_factor=expand_factor,
+    bias=bias,
+    conv_bias=conv_bias,
+    pscan=pscan,
+    vocab_size=vocab_size # adjust based on your dataset
+)
+
+# 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
+tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * max_seq_len
+
+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
+def get_batch(split):
+    global truncated_games_count, total_games_count, current_train_file_index
+
+    # 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)
+    sequences = torch.zeros((batch_size, max_length_in_batch), dtype=torch.int64)
+
+    for i, game in enumerate(games):
+        total_games_count += 1
+
+        if len(game) > max_seq_len:
+            truncated_games_count += 1
+            # Randomly decide truncation strategy
+            truncation_choice = random.choice(['beginning', 'end', 'end2', 'random'])
+            if truncation_choice == 'beginning':
+                # Truncatethe beginning (use from the end backward)
+                truncated_game = game[-max_seq_len:]
+            elif truncation_choice.startswith('end'):
+                # Truncatethe end (use from the beginning forward)
+                truncated_game = game[:max_seq_len]
+            else:
+                # Random start index (truncate beginning and end)
+                start_idx = random.randint(0, len(game) - max_seq_len)
+                truncated_game = game[start_idx:start_idx + max_seq_len]
+            sequences[i, :len(truncated_game)] = torch.tensor(truncated_game, dtype=torch.int64)
+            # Report the percentage of truncated games
+            if truncated_games_count > 0 and truncated_games_count % 50 == 0:
+                truncated_percentage = (truncated_games_count / total_games_count) * 100
+                print(f"Percentage of truncated games: {truncated_percentage:.2f}%\t\t({truncated_games_count}/{total_games_count})")
+        else:
+            sequences[i, :len(game)] = torch.tensor(game, 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
+    
+# 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("Initializing a new Mamba model from scratch")
+    if meta_vocab_size is None:
+        print(f"defaulting to vocab_size of {vocab_size}")
+    else:
+        mamba_config.vocab_size = meta_vocab_size
+    model = MambaLM(mamba_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)
+    mamba_config = checkpoint['model_args']
+    model = MambaLM(mamba_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}")
+    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 / 7.0 if 'anneal_decay_iters' not in checkpoint else checkpoint['anneal_decay_iters'] # / 9 is og
+        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_me.pt, 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 / 10.0  # / 17.0
+        beta2 = np.sqrt(beta2) * beta2
+        auto_clip = True
+        grad_clip_percentile = 6.3333  # 6.75
+elif init_from.startswith('state-spaces'):
+    print(f"Initializing from Mamba pre-trained weights: {init_from}")
+    model = from_pretrained(init_from)
+    mamba_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])
+    
+
+@torch.no_grad()
+def estimate_loss():
+    out = {}
+    model.eval()
+    for split in ['train']: #['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            tokens = get_batch(split)  # Fetch tokens in the correct format
+            logits = model(tokens[:, :-1])  # Predict next tokens (ignore last token)
+
+            # The targets are the tokens shifted by one position
+            targets = tokens[:, 1:].reshape(-1)  # Flatten targets for cross-entropy
+
+            # Compute cross-entropy loss between logits and targets
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets)
+            losses[k] = loss.item()
+
+        split = 'val' # Temporary hack
+        out[split] = losses.mean()
+    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
+
+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:
+        losses = estimate_loss()
+        print(f"\ngame {games_seen} ({iter_num}, {(iter_num / max_iters)*100.0:.3f}%): 'val' loss {losses['val']:.4f}") # Temporary hack
+        #print(f"game {games_seen} ({iter_num}): train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+        if auto_clip and len(grad_norm_history) >= grad_clip_start_size:
+            grad_clip = np.percentile(grad_norm_history, grad_clip_percentile)
+            config['grad_clip'] = grad_clip
+            print(f"Auto adjusted grad_clip to {grad_clip}")
+        if wandb_log:
+            wandb.log({
+                "iter": iter_num,
+                "games": games_seen,
+                #"train/loss": losses['train'], # Temporary hack
+                "grad_clip": grad_clip,
+                "val/loss": losses['val'],
+                "lr": lr,
+            })
+        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': mamba_config,
+                    'iter_num': iter_num,
+                    "games_seen": games_seen,
+                    '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 = get_batch('train')  # Fetch the training data
+        with ctx:
+            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
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets)
+            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)
+    
+    # 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({
+                "iter": iter_num,
+                "games": games_seen,
+                "grad_norm": grad_norm_history[-1] if grad_norm_history else 0,
+                "train/loss": lossf,
+                "lr": lr,
+            })
+    iter_num += 1
+    local_iter_num += 1
+    games_seen += effective_batch_size
+
+    # termination conditions
+    if iter_num > max_iters:
+        checkpoint = {
+            'model': raw_model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'model_args': mamba_config,
+            'iter_num': iter_num,
+            "games_seen": games_seen,
+            '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 checkpoint to {out_dir}")
+        torch.save(checkpoint, os.path.join(out_dir, 'ckpt_final.pt'))
+        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': mamba_config,
+            'iter_num': iter_num,
+            "games_seen": games_seen,
+            '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()
+

train_rl.py

@@ -0,0 +1,537 @@
+import math
+import os
+import time
+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
+from mamba_lm import MambaLM, MambaLMConfig
+import random
+import chess
+from lczero.backends import Weights, Backend, GameState
+
+# Default config values
+out_dir = 'out/play'
+save_interval = 50  
+wandb_project = 'chess-training'
+wandb_run_name = 'lc0-training'
+init_from = 'resume'  # 'scratch', 'resume', 'anneal', or Mamba model name
+
+# Model parameters
+n_layer = 15
+d_model = 256
+dt_rank = 'auto'
+d_state = 16
+vocab_size = 28
+move_num_in_gamestate = False
+
+
+# wandb logging
+wandb_log = True
+wandb_project = 'mamba-rl'
+wandb_run_name = 'mamba_run'
+
+# Load openings file
+with open("openings.csv", "r") as file:
+    lines = file.readlines()[1:]  # Skip header
+    opening_lines = lines
+
+# Optimizer settings
+learning_rate = 1e-7 #7.25e-7
+min_lr = 1e-8 # 1.75e-8
+warmup_iters = 600
+lr_decay_iters = len(opening_lines)
+weight_decay = 1e-2 #5e-3
+beta1 = 0.905 #0.915
+beta2 = 0.965 #0.95
+grad_clip = 0.5 #0.25
+min_grad_clip = 1e-3 #1e-3
+max_grad_clip = 0.45 #0.45
+auto_clip = True
+grad_clip_start_size = 150
+grad_clip_max_size = 600
+grad_clip_percentile = 9
+
+# Game play / loss calculation settings
+top_k = 2 # 2
+top_k_adj_moves = 40 #999 #35
+max_illegal_moves = 8 #2
+max_moves = 87
+update_freq = 3 #1  # How often to do a backward pass
+flush_every = 1
+move_reward_scale_factor = 4.0  # 2.125  # scales down the move reward so it's not so dramatic / so that illegal moves (reward -1) are more dramatic by comparison to bad moves
+decrease_factor = 0.75 # Bonus for winning (1/x is penalty for losing)
+window_size = 300
+
+
+# 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
+
+# Initialize lc0 engines
+lc0_weights_opponent = Weights("./lc0/build/release/11258-32x4-se.pb.gz")
+lc0_backend_opponent = Backend(weights=lc0_weights_opponent)
+
+lc0_weights_evaluator = Weights("./lc0/build/release/11258-48x5-se.pb.gz")
+lc0_backend_evaluator = lc0_backend_opponent #Backend(weights=lc0_weights_evaluator)
+
+# Load tokenizer and decode function
+if move_num_in_gamestate:
+    meta_path = os.path.join(os.path.join('data', 'chess'), 'meta.pkl')
+    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")
+
+
+
+# Initialize Mamba model
+mamba_config = MambaLMConfig(
+    d_model=d_model,
+    n_layers=n_layer,
+    dt_rank=dt_rank,
+    d_state=d_state,
+    vocab_size=vocab_size  # Adjust based on your dataset
+)
+
+model = MambaLM(mamba_config)
+model.to(device)
+
+# 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')
+
+# Compile the model if using PyTorch 2.0
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    model = torch.compile(model)
+
+ddp = int(os.environ.get('RANK', -1)) != -1
+# Wrap model in DDP container if necessary
+if ddp:
+    model = DDP(model, device_ids=[ddp_local_rank])
+    
+win_rate_window = []
+win_only_rate_window = []
+# Load checkpoint if resuming training
+if init_from == 'resume':
+    print(f"Resuming training from {out_dir}")
+    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    mamba_config = checkpoint['model_args']
+    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)
+    optimizer.load_state_dict(checkpoint['optimizer'])
+    iter_num = checkpoint['iter_num']
+    games_played = checkpoint['games_seen']
+    opening_line_index = checkpoint.get('opening_line_index', 0)
+    win_rate_window = checkpoint.get('win_rate_window', [])
+    win_only_rate_window = checkpoint.get('win_only_rate_window', [])
+    best_wr = checkpoint.get('best_wr', 0.0)
+    best_wor = checkpoint.get('best_wor', 0.0)
+    if auto_clip:
+        grad_clip = checkpoint['config']['grad_clip']
+        config['grad_clip'] = grad_clip
+        grad_norm_history = checkpoint.get('grad_norm_history', [])
+    else:
+        grad_norm_history = []
+else:
+    best_wr = 0.0
+    best_wor = 0.0
+    grad_norm_history = []
+    games_played = 0
+    iter_num = 0
+    opening_line_index = 0
+    if auto_clip:
+        grad_clip = 0
+        config['grad_clip'] = 0 
+
+
+def get_model_move(game_state, top_k):
+    model.train()  # Ensure the model is in training mode
+    encoded_prompt = encode(game_state)
+    input_ids = torch.tensor([encoded_prompt], dtype=torch.long, device=device)
+    
+    have_non_space = False
+    logits_list = []  # Collect logits for analysis and potential loss calculation
+    for _ in range(8):
+        logits = model(input_ids)[0, -1, :]  # Logits for the last predicted token
+
+        # We're using top-k more as a VRAM control, not a decision enhacing tool
+        if top_k is not None and top_k < logits.size(-1):
+            logits, indices = torch.topk(logits, top_k)
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_token_id = indices[torch.multinomial(probs, 1)]
+        else:
+            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)
+        logits_list.append(logits)
+        del logits, probs
+
+    # Decode the sequence to extract the move
+    model_response = decode(input_ids.squeeze(0).tolist())
+    try:
+        move = model_response[len(game_state):].split(";")[0].split()[0]  # Extract the first move
+    except IndexError:
+        move = None
+
+    return move, torch.stack(logits_list) if len(logits_list) > 0 else None
+
+def get_lc0_move(board, backend):
+    gamestate = GameState(fen=board.fen())
+    input_planes = gamestate.as_input(backend)
+    result = backend.evaluate(input_planes)[0]
+    moves = gamestate.moves()
+    policy_indices = gamestate.policy_indices()
+    move_probs = np.array(result.p_softmax(*policy_indices))
+    try:
+        best_move_idx = move_probs.argmax()
+    except:
+        return None
+    best_move = moves[best_move_idx]
+    return chess.Move.from_uci(best_move)
+
+def evaluate_position(fen, backend):
+    gamestate = GameState(fen=fen)
+    result = backend.evaluate(gamestate.as_input(backend))[0]
+    return result.q()
+    
+def reward_from_eval(before_eval, after_eval):
+    diff = after_eval - before_eval
+    return diff / (move_reward_scale_factor + abs(diff))
+    
+def backward_pass(loss):
+    global grad_norm_history
+
+    # Backward pass
+    scaler.scale(loss).backward()
+    
+    # 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 0.1)  # 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:]
+    
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+
+def play_game():
+    global top_k
+    
+    optimizer.zero_grad(set_to_none=True)
+    torch.cuda.empty_cache()
+    board = chess.Board()
+    total_loss = 0
+    illegal_moves = 0
+    move_count = 0
+    moves_since_backward = 0
+    tot_reward = 0
+    
+    # Load opening from openings.csv
+    tokens = [m.split(".")[-1] if "." in m else m for m in opening_line.split()]
+    [board.push_san(m) for m in tokens]
+    if move_num_in_gamestate:
+        game_state = opening_line.rstrip() + " "
+    else:
+        game_state = ' '.join(['.' + m.split(".")[-1] if "." in m else m for m in opening_line.split()])
+    fail = False
+    
+    while not board.is_game_over():
+        before_eval = evaluate_position(board.fen(), lc0_backend_evaluator)
+        game_state += f"{board.fullmove_number if move_num_in_gamestate else ''}."
+        model_move, logits = get_model_move(game_state, top_k)
+        move_reward = -1
+        
+        if model_move is None or logits is None:
+            illegal_moves += 1
+            pinch_hit_move = get_lc0_move(board, lc0_backend_opponent)
+            if pinch_hit_move is None:
+                print("Failed game (lc0 couldn't find pinch-hit move)")
+                fail = True
+                tot_reward += move_reward
+                move_count += 1
+                break
+            game_state += f"{board.san(pinch_hit_move)} "
+            board.push(pinch_hit_move)
+        else:
+                try:
+                    #print(model_move)
+                    board.push(board.parse_san(model_move))
+                    game_state += f"{model_move} "
+                except:
+                    illegal_moves += 1
+                    pinch_hit_move = get_lc0_move(board, lc0_backend_opponent)
+                    if pinch_hit_move is None:
+                        print("Failed game (lc0 couldn't find pinch-hit move)")
+                        fail = True
+                        tot_reward += move_reward
+                        move_count += 1
+                        break
+                    game_state += f"{board.san(pinch_hit_move)} "
+                    board.push(pinch_hit_move)
+                else:
+                    if not board.is_valid():
+                        board.pop()
+                        illegal_moves += 1
+                        pinch_hit_move = get_lc0_move(board, lc0_backend_opponent)
+                        if pinch_hit_move is None:
+                            print("Failed game (lc0 couldn't find pinch-hit move)")
+                            fail = True
+                            tot_reward += move_reward
+                            move_count += 1
+                            break
+                        game_state += f"{board.san(pinch_hit_move)} "
+                        board.push(pinch_hit_move)
+                    else:
+                        after_eval = -evaluate_position(board.fen(), lc0_backend_evaluator)
+                        move_reward = reward_from_eval(before_eval, after_eval)
+        
+        tot_reward += move_reward
+        if not board.is_game_over():
+            black_move = get_lc0_move(board, lc0_backend_opponent)
+            if black_move is None:
+                print("Failed game (lc0 couldn't find black move)")
+                fail = True
+                move_count += 1
+                break
+            game_state += f"{board.san(black_move)} "
+            board.push(black_move)
+        
+        if logits is not None:
+            total_loss += torch.sum(torch.nn.functional.log_softmax(logits, dim=-1) * move_reward)
+        logits_none = logits is None
+        del logits
+        moves_since_backward += 1
+        if move_count % update_freq == 0 and not board.is_game_over() and not logits_none:
+            backward_pass(total_loss / moves_since_backward)
+            total_loss = 0.0  # Reset cumulative loss after update
+            moves_since_backward = 0
+        move_count += 1
+        if move_count == top_k_adj_moves:
+            top_k = top_k - 1
+        if move_count >= max_moves:
+            break
+        if move_count % flush_every == 0:
+            torch.cuda.empty_cache()
+
+    if move_count >= top_k_adj_moves:
+        top_k = top_k + 1
+    # Scale loss based on game result and illegal moves
+    avg_reward = tot_reward / move_count
+    #print(f'Avg reward {avg_reward} = {tot_reward} / {move_count}')
+    scale_factor = torch.tensor([1.0], device=device)
+    if move_count >= max_moves:
+        result = "1/2-1/2"
+    elif fail:
+        result = "*"
+    else:
+        result = board.result()
+        total_loss = total_loss / moves_since_backward 
+        if result == "0-1":  # Black wins
+            # Increase the loss for a loss, if the reward is negative (if the loss is positive)
+            scale_factor = torch.tensor([1.0 / decrease_factor], device=device) if avg_reward < 0 and illegal_moves <= max_illegal_moves else scale_factor
+            #print(f'Black win, scale factor adjusted to {scale_factor} (avg award<0 and illegal less max {avg_reward < 0 and illegal_moves <= max_illegal_moves}), illegal vs max {illegal_moves} vs {max_illegal_moves}')
+        elif result == "1-0":  # White wins
+            wdf = decrease_factor / 2.0 if avg_reward <= 0 else 1.0 / decrease_factor
+            #print(f'White win - adjusted decrease factor {wdf}')
+            # Don't update as much for (real) wins. Also change the result so our win_rate isn't inflated.
+            if illegal_moves == 0:
+                scale_factor = torch.tensor([wdf], device=device)
+                #print(f'White win, scale factor adjusted to {scale_factor} (0 illegal moves)')
+            elif illegal_moves <= max_illegal_moves:
+                scale_factor = torch.tensor([(1 + wdf) / 2], device=device)
+                #print(f'White win, scale factor adjusted to {scale_factor} ({0 < illegal_moves <= max_illegal_moves}), illegal vs max {illegal_moves} vs {max_illegal_moves}')
+                result = "1/2-1/2"
+            else:
+                result = "0-1"
+                # No adjustment to scale_factor
+        
+        if total_loss.numel():
+            try:
+                backward_pass(total_loss * scale_factor)
+            except:
+                print("Failed game (final backward pass, result not effected)")
+            total_loss = 0.0
+    
+    #print(f'Scale factor {scale_factor.item()}')
+    return avg_reward / scale_factor.item(), result, illegal_moves, move_count
+    
+    
+def get_lr(it):
+    # 1) linear warmup for warmup_iters steps
+    if it < warmup_iters:
+        return learning_rate * it / warmup_iters
+    # 2) if it > lr_decay_iters, return min learning rate
+    if it > lr_decay_iters:
+        return min_lr
+    # 3) in between, use cosine decay down to min learning rate
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    assert 0 <= decay_ratio <= 1
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
+    return min_lr + coeff * (learning_rate - min_lr)
+
+# Training loop
+if wandb_log:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+while True:
+    t0 = time.time()
+    lr = get_lr(iter_num)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+    opening_line = opening_lines[opening_line_index]
+    
+    if iter_num > 0 and iter_num % save_interval == 0:
+        if auto_clip and len(grad_norm_history) >= grad_clip_start_size:
+            grad_clip = max(min(np.percentile(grad_norm_history, grad_clip_percentile), max_grad_clip), min_grad_clip)
+            config['grad_clip'] = grad_clip
+            print(f"Auto adjusted grad_clip to {grad_clip}")
+    
+        #print(f"Game {games_played}: Loss {game_reward:.4f}, Illegal moves {illegal_moves}, Win rate {win_rate:.3f}")
+        if wandb_log:
+            wandb.log({
+                "etc/iter": iter_num,
+                "etc/lr": lr,
+                "etc/grad_clip": grad_clip,
+                "etc/games_played": games_played,
+            })
+        
+        # Save checkpoint
+        raw_model = model.module if ddp else model
+        checkpoint = {
+            'model': raw_model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'model_args': mamba_config,
+            'iter_num': iter_num,
+            "games_seen": games_played,
+            'config': config,
+            'opening_line_index': opening_line_index,
+            'grad_norm_history': grad_norm_history,
+            'win_rate_window': win_rate_window,
+            'win_only_rate_window': win_only_rate_window,
+            'best_wr': best_wr,
+            'best_wor': best_wor
+        }
+        print(f"saving checkpoint to {out_dir}\n")
+        torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+    
+    # Play a game against lc0 engine
+    game_reward, result, illegal_moves, move_count = play_game()
+    games_played += 1
+    
+    # Backward passes happen in play_game
+    
+    # Log game result and update win rate window
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    score = 0.5
+    if result == "1-0":
+        score = 1
+    elif result == "0-1":
+        score = 0
+    if result != "*":
+        win_rate_window.append(score)
+        win_rate_window = win_rate_window[-window_size:]
+        win_rate = sum(win_rate_window) / len(win_rate_window)
+        win_only_rate_window.append(int(score)) #int to discard draws
+        win_only_rate_window = win_only_rate_window[-window_size:]
+        win_only_rate = float(sum(win_only_rate_window)) / len(win_only_rate_window)
+        if win_rate > best_wr:
+            best_wr = win_rate
+            raw_model = model.module if ddp else model
+            checkpoint = {
+                'model': raw_model.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'model_args': mamba_config,
+                'iter_num': iter_num,
+                "games_seen": games_played,
+                'config': config,
+                'opening_line_index': opening_line_index,
+                'grad_norm_history': grad_norm_history,
+                'win_rate_window': win_rate_window,
+                'best_wr': best_wr,
+                'best_wor': best_wor
+            }
+            print(f"saving checkpoint to {out_dir}\n")
+            torch.save(checkpoint, os.path.join(out_dir, f'ckpt_{games_played}g_wr{best_wr}.pt'))
+        elif win_only_rate > best_wor:
+            best_wor = win_only_rate
+            raw_model = model.module if ddp else model
+            checkpoint = {
+                'model': raw_model.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'model_args': mamba_config,
+                'iter_num': iter_num,
+                "games_seen": games_played,
+                'config': config,
+                'opening_line_index': opening_line_index,
+                'grad_norm_history': grad_norm_history,
+                'win_rate_window': win_rate_window,
+                'best_wr': best_wr,
+                'best_wor': best_wor
+            }
+            print(f"saving checkpoint to {out_dir}\n")
+            torch.save(checkpoint, os.path.join(out_dir, f'ckpt_{games_played}g_wor{best_wor}.pt'))
+        best_wor = max(best_wor, win_only_rate)
+    print(f"Game {games_played} ({iter_num}, {(iter_num / len(opening_lines)) * 100.0:.3f}%): Score {score}, Reward {game_reward:.4f}, Illegal moves {illegal_moves} ({illegal_moves / move_count:.3%}), Total moves {move_count}, Win rate {win_rate:.3f}, Win only rate {win_only_rate:.3f}, time {dt * 1000:.2f}ms")
+    if wandb_log:
+        wandb.log({
+            "etc/iter": iter_num,
+            "etc/lr": lr,
+            "etc/grad_norm_mean": np.mean(grad_norm_history) if grad_norm_history else -1,
+            "etc/grad_zero_pct": float(np.count_nonzero(grad_norm_history==0))/len(grad_norm_history) if grad_norm_history else -1,
+            "etc/games_played": games_played,
+            "eval/game_reward": game_reward,
+            "eval/illegal_move_pct": illegal_moves / move_count,
+            "eval/move_ct": move_count,
+            "eval/win_rate": win_rate,
+            "eval/win_only_rate": win_only_rate,
+        })
+    
+    iter_num += 1
+    opening_line_index += 1
+    
+    # Termination condition
+    if opening_line_index >= len(opening_lines):
+        break
+
+if ddp:
+    destroy_process_group()