feat: add Connect Four LLM training pipeline (SFT + GRPO)

[FloLey]

https://claude.ai/code/session_015dRfJUv1tpJrJ4AVgqLaoL

[gemini-code-assist]

Code Review

This pull request introduces a training pipeline for Connect Four LLM agents using SFT and GRPO via the Unsloth library, including a comprehensive RunPod setup guide and evaluation scripts. The review feedback identifies critical improvements regarding data leakage and reproducibility in the dataset splitting logic, as well as potential runtime errors due to division by zero in the evaluation statistics. Additionally, recommendations were made to refactor global state into a more encapsulated class structure and to clarify command-line examples in the documentation for better usability.

[gemini-code-assist]

There are a couple of issues with how the training and evaluation datasets are split, which can lead to data leakage and lack of reproducibility:

1. Data Leakage: The training data is taken from the start of the shuffled list (data[:max_rows]) in prepare_sft_dataset and prepare_grpo_dataset, while the evaluation data is taken from the end (raw_data[-10_000:]). If sft_max_rows or grpo_max_rows is large enough to overlap with the last 10,000 samples, the model will be evaluated on data it has already seen during training.
2. Reproducibility: random.shuffle(data) is called in prepare_sft_dataset and prepare_grpo_dataset without setting a random seed. This means the training set will be different on each run, making experiments difficult to reproduce. While run_eval sets a seed, the training data preparation does not.

A better approach is to perform a single, deterministic train/test split at the beginning of the script or in main.

Here's a suggested refactoring:

# In main()
raw_data = load_csv_data(config["csv_path"])
random.seed(42)
random.shuffle(raw_data)
eval_data = raw_data[-10_000:]
train_data = raw_data[:-10_000]

# Then pass `train_data` to run_sft/run_grpo and `eval_data` to run_eval
# and remove data loading and splitting from within those functions.
# Also remove `random.shuffle` from the `prepare_*_dataset` functions.

[gemini-code-assist]

If total or valid is 0, this block of code will raise a ZeroDivisionError when calculating and printing statistics. You should add checks to handle these cases before performing division.

[gemini-code-assist]

The placeholder {4b,8b,14b} in the command could be slightly confusing for users. It's clearer to provide a concrete example to show how to replace the placeholder.

Suggested change

	python connect4_train.py --model {4b,8b,14b} --stage export --csv connect4_data.csv
	```
	python connect4_train.py --model 8b --stage export --csv connect4_data.csv

[gemini-code-assist]

Using a global variable SCORE_LOOKUP that is modified in run_grpo and read by the reward functions makes the code harder to reason about and test. This creates a hidden dependency between these functions.

A better approach would be to encapsulate this state. For example, you could use a class to hold the lookup table and define the reward functions as its methods. This would make the data flow explicit.

Example refactoring:

class RewardCalculator:
    def __init__(self, data):
        self.score_lookup = build_lookup_table(data)

    def reward_move_quality(self, completions, move_sequence, **kwargs):
        # ... use self.score_lookup
        ...

# In run_grpo:
raw_data = load_csv_data(config["csv_path"])
reward_calculator = RewardCalculator(raw_data)
trainer = GRPOTrainer(
    ...
    reward_funcs=[
        reward_calculator.reward_format, 
        reward_calculator.reward_move_quality, 
        reward_calculator.reward_is_best_move
    ],
    ...
)

[FloLey]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a Connect Four LLM training pipeline utilizing SFT and GRPO with the Unsloth library, complemented by RunPod setup documentation and environment configuration. The reviewer suggested several improvements, including exposing more training parameters as command-line arguments, adhering to PEP 8 standards for imports and assignments, removing unused imports, and avoiding the use of global variables to improve code maintainability.

[gemini-code-assist]

The troubleshooting tip suggests editing the script for loss_type="dr_grpo". It would be more user-friendly and consistent with other configurable parameters if this could also be passed as a command-line argument to connect4_train.py.

[gemini-code-assist]

For better readability and adherence to PEP 8, it is recommended to have one import statement per line.

import argparse
import csv
import json
import os
import re
import random

[gemini-code-assist]

The defaultdict import is not used in this script. Please remove unused imports to keep the codebase clean.

Suggested change

	from collections import defaultdict
	import torch

[gemini-code-assist]

Modifying a global variable (WANDB_AVAILABLE) directly within a function is generally discouraged as it can lead to unexpected side effects and make code harder to reason about. Consider passing WANDB_AVAILABLE as an argument to main or encapsulating this logic within a class.

Suggested change

	if args.no_wandb: global WANDB_AVAILABLE; WANDB_AVAILABLE = False
	if args.no_wandb:
	WANDB_AVAILABLE = False

[gemini-code-assist]

Multiple assignments on a single line can sometimes reduce readability. For better clarity and adherence to PEP 8, consider splitting this into separate lines.

Suggested change

	config = get_config(args.model); config["csv_path"] = args.csv; config["hf_repo"] = args.hf_repo
	config = get_config(args.model)
	config["csv_path"] = args.csv
	config["hf_repo"] = args.hf_repo

[FloLey]

/Gemini review

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive training and evaluation pipeline for Connect Four agents using GRPO and Unsloth, accompanied by a detailed RunPod setup guide and updated ignore rules. The implementation covers board reconstruction, reward calculation, and model export functionality. The review feedback provides several actionable improvements to enhance the robustness and maintainability of the code, including better CSV file handling, ensuring deterministic evaluation through greedy decoding, preventing redundant repository naming suffixes during HuggingFace uploads, and refactoring global state management to avoid side effects.

[gemini-code-assist]

For improved robustness and to prevent platform-specific issues, it's best practice to specify encoding="utf-8" and newline="" when opening CSV files with Python's csv module. This ensures consistent behavior across different operating systems and correctly handles newlines within quoted fields.

Suggested change

	with open(csv_path, "r") as f:
	with open(csv_path, "r", encoding="utf-8", newline="") as f:

[gemini-code-assist]

The evaluation process is currently stochastic due to do_sample=True. For reproducible and deterministic evaluation results, it's recommended to use greedy decoding. This can be achieved by setting do_sample=False. The temperature parameter can then be removed as it has no effect with greedy decoding.

Suggested change

	outputs = model.generate(input_ids=inputs, max_new_tokens=3072, temperature=0.01, do_sample=True)
	outputs = model.generate(input_ids=inputs, max_new_tokens=3072, do_sample=False)

[gemini-code-assist]

The push_to_hub function constructs the GGUF repository name by appending -GGUF to the provided --hf-repo argument. If a user mistakenly provides a repo name that already ends with -GGUF, it will result in a doubly-appended name (e.g., my-repo-GGUF-GGUF). It would be more robust to add a check for this.

You could add a check after getting hf_repo from the config:

    hf_repo = config.get("hf_repo")
    if hf_repo and hf_repo.upper().endswith("-GGUF"):
        print("WARNING: --hf-repo should be the base repo name. Stripping '-GGUF' suffix.")
        hf_repo = hf_repo[:-5]

[gemini-code-assist]

Modifying a global variable (WANDB_AVAILABLE) within the main function can make the code harder to understand and maintain. A cleaner approach is to determine the desired wandb state within main and pass it as an argument to the functions that need it. This avoids side effects and makes data flow explicit.

Consider this alternative approach:

1. In main(), remove the global modification and instead create a local variable and pass it via the config dictionary:

   # In main()
   use_wandb = WANDB_AVAILABLE and not args.no_wandb
   config['use_wandb'] = use_wandb

2. Update run_grpo to use this config value:

   # In run_grpo()
   use_wandb = config.get('use_wandb', WANDB_AVAILABLE) # Fallback to global default
   if use_wandb:
       wandb.init(...)
   
   grpo_kwargs = dict(
       ...
       report_to="wandb" if use_wandb else "none",
       ...
   )

This makes the dependency on wandb explicit.