The Reference Model: Your AI’s Anchor to Reality

import torch
import copy
from transformers import AutoModelForCausalLM

def create_reference_model(policy_model):
    """
    Create a frozen reference model from the policy.
    
    This is literally just a deep copy with all the parameters frozen.
    We're making a snapshot of the model at this exact moment in time.
    """
    # Deep copy the entire model (this duplicates all weights in memory)
    reference_model = copy.deepcopy(policy_model)
    
    # Freeze every single parameter - no gradients, no updates
    for param in reference_model.parameters():
        param.requires_grad = False
    
    # Set to evaluation mode (disables dropout, etc.)
    reference_model.eval()
    
    return reference_model

# Let's create a reference model from GPT-2
print("Loading GPT-2 as our policy model...")
policy_model = AutoModelForCausalLM.from_pretrained("gpt2")

print("Creating frozen reference model...")
reference_model = create_reference_model(policy_model)

# Verify that the policy can be trained but the reference cannot
policy_trainable = sum(p.numel() for p in policy_model.parameters() if p.requires_grad)
ref_trainable = sum(p.numel() for p in reference_model.parameters() if p.requires_grad)

print(f"\nPolicy model: {policy_trainable:,} trainable parameters")
print(f"Reference model: {ref_trainable:,} trainable parameters")
print(f"\nThe reference is completely frozen - exactly what we want!")

# Option 1: Keep reference in half precision
def create_reference_model_fp16(policy_model):
    """
    Create reference model in half precision to save memory.
    
    FP16 (float16) uses 16 bits per parameter instead of 32.
    Cuts memory usage in half. The reference doesn't need full precision
    since we're just using it for KL divergence calculations.
    """
    reference_model = copy.deepcopy(policy_model)
    reference_model = reference_model.half()  # Convert to FP16
    
    for param in reference_model.parameters():
        param.requires_grad = False
    
    reference_model.eval()
    return reference_model

# Option 2: Move reference to CPU (slower but saves GPU memory)
def create_reference_model_cpu(policy_model):
    """
    Create reference model on CPU to save GPU memory.
    
    This moves the reference entirely to CPU. Frees up GPU memory
    for the policy, but you'll pay a speed penalty when computing
    KL divergence (have to move data between CPU and GPU).
    """
    reference_model = copy.deepcopy(policy_model)
    reference_model = reference_model.cpu()
    
    for param in reference_model.parameters():
        param.requires_grad = False
    
    reference_model.eval()
    return reference_model

# Let's try the FP16 version
print("Creating FP16 reference model...")
reference_model_fp16 = create_reference_model_fp16(policy_model)

# Check memory savings (approximate)
fp32_params = sum(p.numel() for p in reference_model.parameters())
fp16_params = sum(p.numel() for p in reference_model_fp16.parameters())

fp32_memory_mb = (fp32_params * 4) / (1024 * 1024)  # 4 bytes per FP32
fp16_memory_mb = (fp16_params * 2) / (1024 * 1024)  # 2 bytes per FP16

print(f"\nMemory usage comparison:")
print(f"  FP32 reference: ~{fp32_memory_mb:.1f} MB")
print(f"  FP16 reference: ~{fp16_memory_mb:.1f} MB")
print(f"  Savings: ~{fp32_memory_mb - fp16_memory_mb:.1f} MB ({((fp32_memory_mb - fp16_memory_mb) / fp32_memory_mb * 100):.0f}% reduction)")
print(f"\nFor a 7B parameter model, that's saving ~14GB of memory. Not bad!")

import torch.nn.functional as F

def get_log_probs(model, input_ids, attention_mask):
    """
    Get log probabilities for tokens under a model.
    
    This runs the model forward and extracts the log probability
    it assigned to each token that actually appeared in the sequence.
    """
    with torch.no_grad():  # No gradients for the reference
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits  # Raw scores for each token in vocab
    
    # Shift for next-token prediction
    # The model predicts token i+1 from tokens 0...i
    # So logits[:, 0] predicts token 1, logits[:, 1] predicts token 2, etc.
    shift_logits = logits[:, :-1, :]  # Drop last position (no next token)
    shift_labels = input_ids[:, 1:]    # Drop first token (not predicted)
    
    # Convert logits to log probabilities via log softmax
    log_probs = F.log_softmax(shift_logits, dim=-1)
    
    # Extract just the log probs for the actual tokens that appeared
    # This gathers log_probs[i, shift_labels[i]] for each position i
    token_log_probs = torch.gather(
        log_probs,
        dim=-1,
        index=shift_labels.unsqueeze(-1)
    ).squeeze(-1)
    
    return token_log_probs

# Let's test this with a real sentence
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

text = "Hello, how are you doing today?"
inputs = tokenizer(text, return_tensors="pt")

print(f"Input text: '{text}'")
print(f"Tokens: {inputs['input_ids'].tolist()[0]}")

# Get log probs from both models
policy_logprobs = get_log_probs(policy_model, inputs['input_ids'], inputs['attention_mask'])
ref_logprobs = get_log_probs(reference_model, inputs['input_ids'], inputs['attention_mask'])

print(f"\nLog probability shapes: {policy_logprobs.shape}")
print(f"(One log prob for each of the {policy_logprobs.shape[1]} tokens we're predicting)")

# Compute per-token KL divergence
kl_per_token = policy_logprobs - ref_logprobs
kl_mean = kl_per_token.mean()

print(f"\nPer-token KL divergences: {kl_per_token[0].tolist()}")
print(f"Mean KL divergence: {kl_mean.item():.4f}")

print(f"\nSince we haven't trained the policy yet, it's identical to the reference.")
print(f"That's why KL divergence is basically zero. After training, this would be > 0.")

def verify_reference_frozen(policy_model, reference_model):
    """
    Verify that reference model is properly frozen.
    
    Returns True if everything looks good, False if something's wrong.
    """
    # Check 1: No parameters should require gradients
    ref_requires_grad = any(p.requires_grad for p in reference_model.parameters())
    
    # Check 2: Weights should be identical to policy initially
    # (They'll diverge after policy training)
    first_policy_param = next(policy_model.parameters())
    first_ref_param = next(reference_model.parameters())
    weights_equal = torch.allclose(first_policy_param, first_ref_param, atol=1e-6)
    
    print("Reference Model Verification:")
    print(f"  Any parameter requires grad? {ref_requires_grad}")
    print(f"    ✓ Should be False (frozen)" if not ref_requires_grad else "    ✗ PROBLEM: Should be False!")
    
    print(f"\n  Weights equal to policy? {weights_equal}")
    print(f"    ✓ Should be True initially (same starting point)" if weights_equal else "    ✓ Different after training")
    
    print(f"\n  In eval mode? {not reference_model.training}")
    print(f"    ✓ Should be True (no dropout, etc.)" if not reference_model.training else "    ✗ PROBLEM: Should be True!")
    
    all_good = (not ref_requires_grad) and (not reference_model.training)
    
    if all_good:
        print("\n✓ Reference model is properly frozen!")
    else:
        print("\n✗ Something's wrong with the reference model setup!")
    
    return all_good

# Test it
verify_reference_frozen(policy_model, reference_model)

def compute_weight_divergence(policy_model, reference_model):
    """
    Compute how far policy weights have diverged from reference.
    
    This gives you a single number you can track during training.
    Useful for detecting if the policy is drifting too far too fast.
    """
    total_diff = 0.0
    total_norm = 0.0
    
    # Sum up the norm of differences across all parameters
    for (name, p_param), (_, r_param) in zip(
        policy_model.named_parameters(),
        reference_model.named_parameters()
    ):
        diff = (p_param - r_param).norm().item()  # L2 norm of difference
        norm = r_param.norm().item()               # L2 norm of reference
        total_diff += diff
        total_norm += norm
    
    # Relative divergence normalizes by reference magnitude
    relative_divergence = total_diff / (total_norm + 1e-8)
    
    return {
        'absolute_divergence': total_diff,
        'relative_divergence': relative_divergence
    }

# Check initial divergence (should be ~0 since we haven't trained yet)
divergence = compute_weight_divergence(policy_model, reference_model)

print(f"Weight divergence from reference:")
print(f"  Absolute: {divergence['absolute_divergence']:.6f}")
print(f"  Relative: {divergence['relative_divergence']:.6f}")
print(f"\nShould be zero initially. After training, this will increase.")
print(f"If it increases too much, you're drifting far from the reference!")
print(f"\nTypical values during RLHF:")
print(f"  Early training: 0.01 - 0.05")
print(f"  Mid training: 0.05 - 0.15")
print(f"  Late training: 0.15 - 0.30")
print(f"  Too much drift: > 0.50 (might be reward hacking!)")