111 lines
4.4 KiB
Python
111 lines
4.4 KiB
Python
import argparse
|
|
import os
|
|
import random
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from dataset.lm_dataset import DPODataset # noqa: E402
|
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM # noqa: E402
|
|
from trainer.train_dpo import logits_to_log_probs # noqa: E402
|
|
|
|
|
|
def shape(x: torch.Tensor) -> tuple[int, ...]:
|
|
return tuple(x.shape)
|
|
|
|
|
|
def show_active_tokens(tokenizer, y: torch.Tensor, mask: torch.Tensor, label: str, limit: int = 16) -> None:
|
|
active_positions = torch.nonzero(mask, as_tuple=False).flatten().tolist()
|
|
print(f"{label}_active_positions_first={active_positions[:limit]}")
|
|
token_pieces = []
|
|
for pos in active_positions[:limit]:
|
|
token_id = int(y[pos].item())
|
|
token_pieces.append(tokenizer.decode([token_id]))
|
|
print(f"{label}_active_label_tokens_first={token_pieces}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Trace MiniMind DPO dataset and reference-model logprob flow.")
|
|
parser.add_argument("--tokenizer_path", default=str(ROOT / "model"))
|
|
parser.add_argument("--data_path", default=str(ROOT / "course/labs/tiny_dpo.jsonl"))
|
|
parser.add_argument("--index", type=int, default=0)
|
|
parser.add_argument("--max_length", type=int, default=96)
|
|
parser.add_argument("--hidden_size", type=int, default=64)
|
|
parser.add_argument("--num_hidden_layers", type=int, default=2)
|
|
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
args = parser.parse_args()
|
|
|
|
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
|
|
random.seed(42)
|
|
torch.manual_seed(42)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
|
|
dataset = DPODataset(args.data_path, tokenizer, max_length=args.max_length)
|
|
sample = dataset[args.index]
|
|
|
|
x_chosen = sample["x_chosen"].unsqueeze(0).to(args.device)
|
|
y_chosen = sample["y_chosen"].unsqueeze(0).to(args.device)
|
|
mask_chosen = sample["mask_chosen"].unsqueeze(0).to(args.device)
|
|
x_rejected = sample["x_rejected"].unsqueeze(0).to(args.device)
|
|
y_rejected = sample["y_rejected"].unsqueeze(0).to(args.device)
|
|
mask_rejected = sample["mask_rejected"].unsqueeze(0).to(args.device)
|
|
|
|
x = torch.cat([x_chosen, x_rejected], dim=0)
|
|
y = torch.cat([y_chosen, y_rejected], dim=0)
|
|
mask = torch.cat([mask_chosen, mask_rejected], dim=0)
|
|
|
|
config = MiniMindConfig(
|
|
hidden_size=args.hidden_size,
|
|
num_hidden_layers=args.num_hidden_layers,
|
|
vocab_size=len(tokenizer),
|
|
)
|
|
policy = MiniMindForCausalLM(config).to(args.device).eval()
|
|
ref = MiniMindForCausalLM(config).to(args.device).eval()
|
|
ref.load_state_dict(policy.state_dict())
|
|
ref.requires_grad_(False)
|
|
|
|
with torch.no_grad():
|
|
ref_logits = ref(x).logits
|
|
policy_logits = policy(x).logits
|
|
|
|
ref_log_probs = logits_to_log_probs(ref_logits, y)
|
|
policy_log_probs = logits_to_log_probs(policy_logits, y)
|
|
diff = (policy_log_probs - ref_log_probs).abs().max().item()
|
|
|
|
print("[Dataset]")
|
|
print(f"num_samples={len(dataset)}")
|
|
print(f"x_chosen.shape={shape(x_chosen)}")
|
|
print(f"y_chosen.shape={shape(y_chosen)}")
|
|
print(f"mask_chosen.shape={shape(mask_chosen)}")
|
|
print(f"x_rejected.shape={shape(x_rejected)}")
|
|
print(f"y_rejected.shape={shape(y_rejected)}")
|
|
print(f"mask_rejected.shape={shape(mask_rejected)}")
|
|
print(f"mask_chosen_active={int(mask_chosen.sum().item())}")
|
|
print(f"mask_rejected_active={int(mask_rejected.sum().item())}")
|
|
show_active_tokens(tokenizer, y_chosen[0].cpu(), mask_chosen[0].cpu(), "chosen")
|
|
show_active_tokens(tokenizer, y_rejected[0].cpu(), mask_rejected[0].cpu(), "rejected")
|
|
print()
|
|
print("[Concat]")
|
|
print(f"x_cat.shape={shape(x)}")
|
|
print(f"y_cat.shape={shape(y)}")
|
|
print(f"mask_cat.shape={shape(mask)}")
|
|
print()
|
|
print("[Policy / Reference]")
|
|
print(f"policy_logits.shape={shape(policy_logits)}")
|
|
print(f"ref_logits.shape={shape(ref_logits)}")
|
|
print(f"policy_log_probs.shape={shape(policy_log_probs)}")
|
|
print(f"ref_log_probs.shape={shape(ref_log_probs)}")
|
|
print(f"initial_policy_ref_logprob_max_abs_diff={diff:.12f}")
|
|
print(f"ref_requires_grad={any(p.requires_grad for p in ref.parameters())}")
|
|
print(f"policy_requires_grad={any(p.requires_grad for p in policy.parameters())}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|