137 lines
5.3 KiB
Python
137 lines
5.3 KiB
Python
import argparse
|
|
import random
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def masked_mean(x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
|
|
return (x * mask).sum() / mask.sum().clamp(min=1)
|
|
|
|
|
|
def response_logp_positions(prompt_lens: torch.Tensor, response_len: int) -> torch.Tensor:
|
|
response_idx = torch.arange(response_len, device=prompt_lens.device).unsqueeze(0)
|
|
return prompt_lens.unsqueeze(1) - 1 + response_idx
|
|
|
|
|
|
def response_token_logps(logits: torch.Tensor, output_ids: torch.Tensor, logp_pos: torch.Tensor) -> torch.Tensor:
|
|
labels = output_ids[:, 1:]
|
|
all_token_logps = F.log_softmax(logits[:, :-1, :], dim=-1).gather(
|
|
2,
|
|
labels.unsqueeze(-1),
|
|
).squeeze(-1)
|
|
return all_token_logps.gather(1, logp_pos)
|
|
|
|
|
|
def put_sequence_reward_on_last_token(rewards: torch.Tensor, response_mask: torch.Tensor) -> torch.Tensor:
|
|
token_rewards = torch.zeros_like(response_mask)
|
|
lengths = response_mask.sum(dim=1).long().clamp(min=1)
|
|
last_idx = lengths - 1
|
|
token_rewards[torch.arange(rewards.size(0)), last_idx] = rewards
|
|
return token_rewards
|
|
|
|
|
|
def group_relative_advantages(rewards: torch.Tensor, num_generations: int, eps: float = 1e-4) -> torch.Tensor:
|
|
grouped_rewards = rewards.view(-1, num_generations)
|
|
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(num_generations)
|
|
std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(num_generations)
|
|
return (rewards - mean_r) / (std_r + eps)
|
|
|
|
|
|
def reference_kl_penalty(policy_logps: torch.Tensor, ref_logps: torch.Tensor) -> torch.Tensor:
|
|
delta = ref_logps - policy_logps
|
|
return torch.exp(delta) - delta - 1.0
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Trace MiniMind reward/logprob/KL tensor relationships.")
|
|
parser.add_argument("--vocab_size", type=int, default=64)
|
|
parser.add_argument("--prompt_len", type=int, default=3)
|
|
parser.add_argument("--response_len", type=int, default=3)
|
|
parser.add_argument("--num_generations", type=int, default=3)
|
|
args = parser.parse_args()
|
|
|
|
random.seed(7)
|
|
torch.manual_seed(7)
|
|
|
|
output_ids = torch.tensor(
|
|
[
|
|
[11, 12, 13, 21, 22, 2],
|
|
[14, 15, 16, 31, 2, 0],
|
|
],
|
|
dtype=torch.long,
|
|
)
|
|
prompt_lens = torch.full((output_ids.size(0),), args.prompt_len, dtype=torch.long)
|
|
completion_ids = output_ids[:, args.prompt_len:]
|
|
response_mask = torch.tensor(
|
|
[
|
|
[1.0, 1.0, 1.0],
|
|
[1.0, 1.0, 0.0],
|
|
]
|
|
)
|
|
logp_pos = response_logp_positions(prompt_lens, args.response_len)
|
|
|
|
logits_shape = (output_ids.size(0), output_ids.size(1), args.vocab_size)
|
|
old_logits = torch.randn(logits_shape)
|
|
policy_logits = old_logits + 0.15 * torch.randn(logits_shape)
|
|
ref_logits = old_logits + 0.10 * torch.randn(logits_shape)
|
|
|
|
old_logps = response_token_logps(old_logits, output_ids, logp_pos)
|
|
policy_logps = response_token_logps(policy_logits, output_ids, logp_pos)
|
|
ref_logps = response_token_logps(ref_logits, output_ids, logp_pos)
|
|
|
|
log_ratio = policy_logps - old_logps
|
|
ratio = torch.exp(log_ratio)
|
|
approx_kl = masked_mean(0.5 * log_ratio.pow(2), response_mask)
|
|
ref_kl = reference_kl_penalty(policy_logps, ref_logps)
|
|
ref_kl_mean = masked_mean(ref_kl, response_mask)
|
|
|
|
ppo_rewards = torch.tensor([2.0, -0.5])
|
|
token_rewards = put_sequence_reward_on_last_token(ppo_rewards, response_mask)
|
|
|
|
grpo_rewards = torch.tensor([2.0, 0.0, -1.0, 0.5, 0.5, 0.5])
|
|
grpo_advantages = group_relative_advantages(grpo_rewards, args.num_generations)
|
|
|
|
chosen_policy_seq = torch.tensor([-4.0, -2.5])
|
|
rejected_policy_seq = torch.tensor([-5.0, -4.0])
|
|
chosen_ref_seq = torch.tensor([-4.5, -3.0])
|
|
rejected_ref_seq = torch.tensor([-4.8, -3.8])
|
|
pi_logratios = chosen_policy_seq - rejected_policy_seq
|
|
ref_logratios = chosen_ref_seq - rejected_ref_seq
|
|
dpo_improvement = pi_logratios - ref_logratios
|
|
|
|
print("[Token logprobs]")
|
|
print(f"output_ids.shape={tuple(output_ids.shape)}")
|
|
print(f"completion_ids={completion_ids.tolist()}")
|
|
print(f"logp_pos={logp_pos.tolist()}")
|
|
print(f"old_logps.shape={tuple(old_logps.shape)}")
|
|
print(f"policy_logps.shape={tuple(policy_logps.shape)}")
|
|
print(f"ref_logps.shape={tuple(ref_logps.shape)}")
|
|
print(f"response_mask={response_mask.tolist()}")
|
|
print()
|
|
|
|
print("[Reward to advantage]")
|
|
print(f"ppo_sequence_rewards={ppo_rewards.tolist()}")
|
|
print(f"ppo_token_rewards={token_rewards.tolist()}")
|
|
print(f"grpo_grouped_rewards={grpo_rewards.view(-1, args.num_generations).tolist()}")
|
|
print(f"grpo_group_advantages={grpo_advantages.view(-1, args.num_generations).tolist()}")
|
|
print()
|
|
|
|
print("[KL and ratio]")
|
|
print(f"ratio.shape={tuple(ratio.shape)}")
|
|
print(f"ratio_valid_min={ratio[response_mask.bool()].min().item():.8f}")
|
|
print(f"ratio_valid_max={ratio[response_mask.bool()].max().item():.8f}")
|
|
print(f"approx_kl_policy_vs_old={approx_kl.item():.8f}")
|
|
print(f"reference_kl_penalty_mean={ref_kl_mean.item():.8f}")
|
|
print()
|
|
|
|
print("[DPO contrast]")
|
|
print(f"pi_logratios={pi_logratios.tolist()}")
|
|
print(f"ref_logratios={ref_logratios.tolist()}")
|
|
print(f"dpo_improvement={dpo_improvement.tolist()}")
|
|
print("dpo_reference_usage=sequence_logratio_not_token_kl_penalty")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|