Reinforcement Learning from Human Feedback (RLHF)

# Clone nanochat repository
git clone https://github.com/karpathy/nanochat.git
cd nanochat
 
# Install uv package manager (fast, reliable dependency management)
curl -LsSf https://astral.sh/uv/install.sh | sh
 
# Create virtual environment and install dependencies
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
uv pip install -e .
 
# Verify installation
python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA available: {torch.cuda.is_available()}')"

# Verify SFT checkpoint exists
ls -la chatsft_checkpoints/
 
# Should see directories like: d12/, d20/, or d24/ with model_step_*.pt files
# If missing, see: /blog/nanochat-deep-dive/fine-tuning-for-chat-sft

"""
I put GRPO in quotes because we actually end up with something a lot
simpler and more similar to just REINFORCE:
 
1) Delete trust region, so there is no KL regularization to a reference model
2) We are on policy, so there's no need for PPO ratio+clip.
3) We use GAPO style normalization that is token-level, not sequence-level.
4) Instead of z-score normalization (r - mu)/sigma, only use (r - mu) as the advantage.
"""

source = "sft"                    # Start from SFT model (not base)
device_batch_size = 8            # Max samples per forward pass
examples_per_step = 16           # Training examples per gradient step
num_samples = 16                 # Samples per example (for advantage estimation)
max_new_tokens = 256             # Max response length
temperature = 1.0                # Sampling temperature
top_k = 50                       # Top-k sampling
 
# Learning rates (same structure as SFT)
unembedding_lr = 0.004
embedding_lr = 0.2
matrix_lr = 0.02
weight_decay = 0.0
init_lr_frac = 0.05              # Start at 5% of base LR
 
num_epochs = 1                   # One pass through GSM8K train set
save_every = 60                  # Checkpoint frequency
eval_every = 60                  # Evaluation frequency
eval_examples = 400              # Examples for pass@k evaluation

messages = [
    {"role": "user", "content": "Weng earns $12 an hour..."},
    {"role": "assistant", "content": [
        {"type": "text", "text": "Weng earns 12/60 = $"},
        {"type": "python", "text": "12/60"},
        {"type": "python_output", "text": "0.2"},
        {"type": "text", "text": "0.2 per minute.\nWorking 50 minutes, she earned 0.2 x 50 = $"},
        {"type": "python", "text": "0.2*50"},
        {"type": "python_output", "text": "10"},
        {"type": "text", "text": "10.\n#### 10"},
    ]}
]

@torch.no_grad()
def get_batch():
    assistant_end = tokenizer.encode_special("<|assistant_end|>")
    rank_indices = range(ddp_rank, len(train_task), ddp_world_size)
    
    for example_idx in itertools.cycle(rank_indices):
        # Get the conversation
        conversation = train_task[example_idx]
        
        # Tokenize: remove assistant message, prime for completion
        tokens = tokenizer.render_for_completion(conversation)
        prefix_length = len(tokens)
        
        # Generate num_samples completions
        model.eval()
        generated_token_sequences = []
        masks = []
        
        num_sampling_steps = num_samples // device_batch_size
        for sampling_step in range(num_sampling_steps):
            seed = hash((step, example_idx, sampling_step)) & 0x7FFFFFFF
            
            try:
                with autocast_ctx:
                    sequences_batch, masks_batch = engine.generate_batch(
                        tokens,
                        num_samples=device_batch_size,
                        max_tokens=max_new_tokens,
                        temperature=temperature,
                        top_k=top_k,
                        seed=seed,
                    )
                generated_token_sequences.extend(sequences_batch)
                masks.extend(masks_batch)
                
            except RuntimeError as e:
                if "out of memory" in str(e):
                    logging.error(f"OOM during rollout at step {step}, example {example_idx}. Clearing cache.")
                    torch.cuda.empty_cache()
                    # Reduce batch size for this rollout
                    device_batch_size = max(1, device_batch_size // 2)
                    continue
                else:
                    raise e

rewards = []
for sample_tokens in generated_token_sequences:
    # Extract generated response (after prompt)
    generated_tokens = sample_tokens[prefix_length:]
    generated_text = tokenizer.decode(generated_tokens)
    
    # Calculate reward
    reward = train_task.reward(conversation, generated_text)
    rewards.append(reward)

def reward(self, conversation, assistant_response):
    """Binary reward: 1.0 if answer is correct, 0.0 otherwise"""
    is_correct = self.evaluate(conversation, assistant_response)
    return float(is_correct)

GSM_RE = re.compile(r"#### (\-?[0-9\.\,]+)")
 
def extract_answer(completion):
    """Extract numerical answer after #### marker"""
    match = GSM_RE.search(completion)
    if match:
        match_str = match.group(1).strip().replace(",", "")
        return match_str
    return None

rewards = torch.tensor(rewards, dtype=torch.float, device=device)  # (B,)
mu = rewards.mean()
advantages = rewards - mu  # Simple baseline, no z-score normalization

max_length = max(len(seq) for seq in generated_token_sequences)
padded_sequences = [seq + [assistant_end] * (max_length - len(seq)) 
                    for seq in generated_token_sequences]
padded_masks = [mask + [0] * (max_length - len(mask)) 
                for mask in masks]
 
ids = torch.tensor(padded_sequences, dtype=torch.long, device=device)
mask_ids = torch.tensor(padded_masks, dtype=torch.long, device=device)
 
# Autoregressive setup
inputs = ids[:, :-1]
targets = ids[:, 1:].clone()
targets[mask_ids[:, 1:] == 0] = -1  # Mask prompt and tool outputs

for example_step in range(examples_per_rank):
    # Get batch for one training example
    sequences_all, inputs_all, targets_all, rewards_all, advantages_all = next(batch_iterator)
    
    model.train()
    
    # Process in device_batch_size chunks
    num_passes = inputs_all.size(0) // device_batch_size
    for pass_idx in range(num_passes):
        b0, b1 = pass_idx * device_batch_size, (pass_idx + 1) * device_batch_size
        inputs = inputs_all[b0:b1]
        targets = targets_all[b0:b1]
        advantages = advantages_all[b0:b1]
        
        try:
            # Calculate log probabilities
            with autocast_ctx:
                logp = -model(inputs, targets, loss_reduction='none').view_as(inputs)  # (B, T)
            
            # Check for NaN in loss (indicates instability)
            if torch.isnan(logp).any():
                logging.warning(f"NaN detected in policy gradient at step {step}, pass {pass_idx}. Skipping.")
                optimizer.zero_grad()
                continue
            
            # Policy gradient objective
            pg_obj = (logp * advantages.unsqueeze(-1)).sum()
            
            # Normalize by valid tokens and batch structure
            num_valid = (targets >= 0).sum().clamp(min=1)
            pg_obj = pg_obj / (num_valid * num_passes * examples_per_rank)
            
            # Minimize negative objective (maximize objective)
            loss = -pg_obj
            loss.backward()
            
        except RuntimeError as e:
            if "out of memory" in str(e):
                logging.error(f"OOM during backward pass at step {step}. Clearing cache.")
                torch.cuda.empty_cache()
                optimizer.zero_grad()
                continue
            else:
                raise e

ratio = torch.exp(logp - old_logp)
clipped_ratio = torch.clamp(ratio, 1-epsilon, 1+epsilon)
loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()

loss = -pg_obj + beta * KL(π, π_ref)

if next_token == python_start:
    state.in_python_block = True
    state.python_expr_tokens = []
elif next_token == python_end and state.in_python_block:
    state.in_python_block = False
    if state.python_expr_tokens:
        expr = self.tokenizer.decode(state.python_expr_tokens)
        result = use_calculator(expr)
        if result is not None:
            result_tokens = self.tokenizer.encode(str(result))
            state.forced_tokens.append(output_start)
            state.forced_tokens.extend(result_tokens)
            state.forced_tokens.append(output_end)
    state.python_expr_tokens = []
elif state.in_python_block:
    state.python_expr_tokens.append(next_token)

def use_calculator(expr):
    """Evaluate a math expression safely"""
    expr = expr.replace(",", "")
    # Only allow numeric chars and basic operators
    if any([x not in "0123456789*+-/.() " for x in expr]):
        return None
    # Disallow power operator (expensive)
    if "**" in expr:
        return None
    return eval_with_timeout(expr, max_time=3)

if step % eval_every == 0:
    model.eval()
    passk = torch.zeros(device_batch_size, device=device)
    
    with autocast_ctx:
        records_iter = run_gsm8k_eval(val_task, tokenizer, engine, 
                                       num_samples=device_batch_size, 
                                       max_examples=eval_examples, 
                                       temperature=1.0)
        records = list(records_iter)
    
    # Calculate pass@k for k=1..device_batch_size
    for k in range(1, device_batch_size + 1):
        passk[k-1] = sum(any(o["is_correct"] for o in r["outcomes"][:k]) 
                         for r in records)
    
    # Aggregate across ranks
    if ddp:
        dist.all_reduce(passk, op=dist.ReduceOp.SUM)
    
    passk = passk / num_records

def get_lr_multiplier(it):
    lrm = 1.0 - it / num_steps
    return lrm
 
# Apply each step
for opt in optimizers:
    for group in opt.param_groups:
        group["lr"] = group["initial_lr"] * lrm

# Check for SFT checkpoint
ls -la chatsft_checkpoints/
 
# Should see: d12/ or d24/ with model_step_*.pt

run = "gsm8k_rl_run1"
source = "sft"
num_epochs = 1
examples_per_step = 16
device_batch_size = 8
num_samples = 16
temperature = 1.0
init_lr_frac = 0.05
eval_every = 30
save_every = 30

python -m scripts.chat_rl --config rl_config.txt

torchrun --standalone --nproc_per_node=8 -m scripts.chat_rl -- --config rl_config.txt

def reward(self, conversation, assistant_response):
    # Base reward: correctness
    is_correct = self.evaluate(conversation, assistant_response)
    reward = float(is_correct)
    
    # Bonus for shorter solutions
    length = len(assistant_response)
    reward += 0.1 * max(0, 1 - length / 500)
    
    # Bonus for showing work
    if "<<" in assistant_response:  # Used calculator
        reward += 0.05
    
    # Penalty for format violations
    if "####" not in assistant_response:
        reward -= 0.2
    
    return reward

def reward(self, conversation, assistant_response):
    # Extract all calculator results
    calc_results = re.findall(r'<<(.+?)=(.+?)>>', assistant_response)
    
    # Check if intermediate steps are correct
    correct_steps = 0
    for expr, result in calc_results:
        expected = eval(expr)
        if abs(float(result) - expected) < 0.01:
            correct_steps += 1
    
    # Partial credit
    if correct_steps > 0:
        return 0.1 * correct_steps
    
    # Full credit for correct final answer
    if self.evaluate(conversation, assistant_response):
        return 1.0
    
    return 0.0

rewards_dict = {
    "correctness": 1.0 if correct else 0.0,
    "efficiency": -len(response) / 1000,  # Shorter is better
    "clarity": count_explanation_sentences(response) / 10,
}
 
# Weighted sum
reward = (
    1.0 * rewards_dict["correctness"] +
    0.1 * rewards_dict["efficiency"] +
    0.05 * rewards_dict["clarity"]
)

# BAD: Reward word count
reward += 0.01 * word_count  # Model learns to be verbose

# BAD: Contradictory rewards
reward += 0.1 if len(response) < 100 else 0  # Reward brevity
reward += 0.1 if "detailed explanation" in response else 0  # Reward detail

# BAD: Exploitable pattern
reward += 0.5 if "Therefore, the answer is" in response else 0
# Model learns to always write this phrase regardless of correctness

# During rollout generation
print(f"Reward distribution: {torch.bincount(rewards.long())}")
# Output: tensor([10, 6])  → 10 wrong, 6 correct
 
# Calculate success rate
success_rate = rewards.sum() / len(rewards)
print(f"Success rate: {success_rate:.2%}")

# After loss.backward()
total_norm = 0
for p in model.parameters():
    if p.grad is not None:
        param_norm = p.grad.data.norm(2)
        total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
print(f"Gradient norm: {total_norm:.4f}")

advantages = rewards - rewards.mean()
print(f"Advantages: min={advantages.min():.3f}, max={advantages.max():.3f}, std={advantages.std():.3f}")

sequence_lengths = [len(seq) for seq in generated_token_sequences]
print(f"Seq lengths: min={min(sequence_lengths)}, max={max(sequence_lengths)}, mean={sum(sequence_lengths)/len(sequence_lengths):.1f}")

def reward(self, conversation, assistant_response):
    # Extract code, run test cases
    code = extract_code_block(assistant_response)
    test_results = run_tests(code, self.test_cases)
    return float(all(test_results))

def reward(self, conversation, assistant_response):
    # Check if response follows constraints
    instruction = parse_instruction(conversation)
    follows_format = check_format(assistant_response, instruction.format)
    includes_keywords = check_keywords(assistant_response, instruction.keywords)
    return 0.5 * follows_format + 0.5 * includes_keywords

def reward(self, conversation, assistant_response):
    # Use a reward model (small LM trained on human preferences)
    return reward_model.score(conversation + [assistant_response])

def get_batch_multiturn():
    # Start with conversation up to turn N-1
    conversation = sample_conversation(max_turns=3)
    
    # Generate turn N
    tokens = tokenizer.render_for_completion(conversation)
    sequences = engine.generate_batch(tokens, num_samples=16)
    
    # Reward based on full conversation
    rewards = [reward_conversation(conversation + [response]) for response in sequences]
    
    yield sequences, rewards