119 lines
4.5 KiB
Python
119 lines
4.5 KiB
Python
import argparse
|
|
import os
|
|
import random
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from dataset.lm_dataset import SFTDataset # noqa: E402
|
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM # noqa: E402
|
|
from trainer.train_distillation import distillation_loss # noqa: E402
|
|
|
|
|
|
def shape(x: torch.Tensor) -> tuple[int, ...]:
|
|
return tuple(x.shape)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Trace MiniMind distillation CE + KL loss on a tiny SFT sample.")
|
|
parser.add_argument("--tokenizer_path", default=str(ROOT / "model"))
|
|
parser.add_argument("--data_path", default=str(ROOT / "course/labs/tiny_sft.jsonl"))
|
|
parser.add_argument("--index", type=int, default=0)
|
|
parser.add_argument("--max_length", type=int, default=96)
|
|
parser.add_argument("--student_hidden_size", type=int, default=64)
|
|
parser.add_argument("--student_num_layers", type=int, default=2)
|
|
parser.add_argument("--teacher_hidden_size", type=int, default=96)
|
|
parser.add_argument("--teacher_num_layers", type=int, default=3)
|
|
parser.add_argument("--alpha", type=float, default=0.5)
|
|
parser.add_argument("--temperature", type=float, default=1.5)
|
|
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 = SFTDataset(args.data_path, tokenizer, max_length=args.max_length)
|
|
input_ids, labels = dataset[args.index]
|
|
input_ids = input_ids.unsqueeze(0).to(args.device)
|
|
labels = labels.unsqueeze(0).to(args.device)
|
|
|
|
student_config = MiniMindConfig(
|
|
hidden_size=args.student_hidden_size,
|
|
num_hidden_layers=args.student_num_layers,
|
|
vocab_size=len(tokenizer),
|
|
)
|
|
teacher_config = MiniMindConfig(
|
|
hidden_size=args.teacher_hidden_size,
|
|
num_hidden_layers=args.teacher_num_layers,
|
|
vocab_size=len(tokenizer),
|
|
)
|
|
student = MiniMindForCausalLM(student_config).to(args.device).eval()
|
|
teacher = MiniMindForCausalLM(teacher_config).to(args.device).eval()
|
|
teacher.requires_grad_(False)
|
|
|
|
with torch.no_grad():
|
|
student_outputs = student(input_ids)
|
|
teacher_outputs = teacher(input_ids)
|
|
|
|
student_logits = student_outputs.logits[..., :-1, :].contiguous()
|
|
teacher_logits = teacher_outputs.logits[..., :-1, :].contiguous()
|
|
shift_labels = labels[..., 1:].contiguous()
|
|
loss_mask = (shift_labels != -100)
|
|
loss_mask_flat = loss_mask.view(-1)
|
|
|
|
flat_student_logits = student_logits.view(-1, student_logits.size(-1))
|
|
flat_teacher_logits = teacher_logits.view(-1, teacher_logits.size(-1))
|
|
flat_labels = shift_labels.view(-1)
|
|
|
|
per_token_ce = F.cross_entropy(
|
|
flat_student_logits,
|
|
flat_labels,
|
|
ignore_index=-100,
|
|
reduction="none",
|
|
)
|
|
ce_loss_raw = torch.sum(per_token_ce * loss_mask_flat.float()) / (loss_mask_flat.float().sum() + 1e-8)
|
|
|
|
active_student_logits = flat_student_logits[loss_mask_flat]
|
|
active_teacher_logits = flat_teacher_logits[loss_mask_flat]
|
|
distill = distillation_loss(
|
|
active_student_logits,
|
|
active_teacher_logits,
|
|
temperature=args.temperature,
|
|
)
|
|
total = args.alpha * ce_loss_raw + (1 - args.alpha) * distill
|
|
|
|
print("[Config]")
|
|
print(f"student_hidden_size={args.student_hidden_size}")
|
|
print(f"teacher_hidden_size={args.teacher_hidden_size}")
|
|
print(f"alpha={args.alpha}")
|
|
print(f"temperature={args.temperature}")
|
|
print()
|
|
print("[Shapes]")
|
|
print(f"input_ids.shape={shape(input_ids)}")
|
|
print(f"labels.shape={shape(labels)}")
|
|
print(f"student_logits.shape={shape(student_logits)}")
|
|
print(f"teacher_logits.shape={shape(teacher_logits)}")
|
|
print(f"shift_labels.shape={shape(shift_labels)}")
|
|
print(f"active_student_logits.shape={shape(active_student_logits)}")
|
|
print(f"active_teacher_logits.shape={shape(active_teacher_logits)}")
|
|
print()
|
|
print("[Loss]")
|
|
print(f"active_tokens={int(loss_mask_flat.sum().item())}")
|
|
print(f"ce_loss_raw={ce_loss_raw.item():.8f}")
|
|
print(f"distill_loss={distill.item():.8f}")
|
|
print(f"total_loss={total.item():.8f}")
|
|
print(f"manual_total_formula={args.alpha:.2f}*CE + {1 - args.alpha:.2f}*Distill")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|