Njap-png
Add Agron neural-network zoo, optimisers, algorithms, schedulers and unified trainer
a311e94 | """Objective-level algorithms: knowledge distillation and DPO.""" | |
| from __future__ import annotations | |
| from typing import Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from .base import AlgorithmContext, BaseAlgorithm | |
| def distillation_loss( | |
| student_logits: torch.Tensor, | |
| teacher_logits: torch.Tensor, | |
| labels: Optional[torch.Tensor] = None, | |
| temperature: float = 2.0, | |
| alpha: float = 0.5, | |
| ) -> torch.Tensor: | |
| """Combined cross-entropy + KL divergence distillation loss.""" | |
| tau = max(temperature, 1e-6) | |
| student_log_probs = F.log_softmax(student_logits / tau, dim=-1) | |
| teacher_probs = F.softmax(teacher_logits / tau, dim=-1) | |
| kl = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * (tau ** 2) | |
| if labels is None: | |
| return alpha * kl | |
| ce = F.cross_entropy( | |
| student_logits[..., :-1, :].reshape(-1, student_logits.shape[-1]).float(), | |
| labels[..., 1:].reshape(-1), | |
| ) | |
| return (1 - alpha) * ce + alpha * kl | |
| def dpo_loss( | |
| policy_chosen_logps: torch.Tensor, | |
| policy_rejected_logps: torch.Tensor, | |
| ref_chosen_logps: torch.Tensor, | |
| ref_rejected_logps: torch.Tensor, | |
| beta: float = 0.1, | |
| label_smoothing: float = 0.0, | |
| ) -> torch.Tensor: | |
| """Direct Preference Optimisation loss (Rafailov et al., 2023).""" | |
| chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps) | |
| rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps) | |
| logits = chosen_rewards - rejected_rewards | |
| return -F.logsigmoid(logits - label_smoothing).mean() | |
| def sequence_log_probs( | |
| model, | |
| input_ids: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| ) -> torch.Tensor: | |
| """Sum of per-token log probabilities for each sequence in the batch.""" | |
| out = model(input_ids=input_ids, attention_mask=attention_mask) | |
| logits = out["logits"] if isinstance(out, dict) else out.logits | |
| logits = logits[:, :-1, :].float() | |
| targets = input_ids[:, 1:] | |
| log_probs = F.log_softmax(logits, dim=-1) | |
| token_log_probs = log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1) | |
| if attention_mask is not None: | |
| mask = attention_mask[:, 1:].float() | |
| token_log_probs = token_log_probs * mask | |
| return token_log_probs.sum(dim=-1) | |
| class Distillation(BaseAlgorithm): | |
| """Distil a teacher model into the student during training.""" | |
| name = "distillation" | |
| category = "objective" | |
| description = "Matches a teacher model's logits via temperature-scaled KL." | |
| def __init__(self, teacher=None, temperature: float = 2.0, alpha: float = 0.5): | |
| super().__init__(temperature=temperature, alpha=alpha) | |
| self.teacher = teacher | |
| self.temperature = temperature | |
| self.alpha = alpha | |
| def on_train_begin(self, ctx): | |
| if self.teacher is None: | |
| self.teacher = ctx.get("teacher") | |
| if self.teacher is not None: | |
| self.teacher.eval() | |
| for p in self.teacher.parameters(): | |
| p.requires_grad_(False) | |
| def train_step(self, ctx, batch): | |
| batch = self.on_step_begin(ctx, batch) | |
| outputs = ctx.model(**batch) | |
| student_logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits | |
| labels = batch.get("labels") | |
| if self.teacher is None: | |
| loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss | |
| loss.backward() | |
| return loss | |
| with torch.no_grad(): | |
| teacher_out = self.teacher(input_ids=batch["input_ids"], | |
| attention_mask=batch.get("attention_mask")) | |
| teacher_logits = teacher_out["logits"] if isinstance(teacher_out, dict) else teacher_out.logits | |
| loss = distillation_loss(student_logits, teacher_logits, labels, | |
| temperature=self.temperature, alpha=self.alpha) | |
| loss.backward() | |
| return loss | |
| class DPO(BaseAlgorithm): | |
| """Direct Preference Optimisation over (chosen, rejected) pairs. | |
| The batch must contain ``chosen_input_ids``, ``rejected_input_ids`` and the | |
| matching ``*_attention_mask`` tensors. | |
| """ | |
| name = "dpo" | |
| category = "objective" | |
| description = "Trains on preference pairs with the DPO objective." | |
| def __init__(self, beta: float = 0.1, reference_model=None, label_smoothing: float = 0.0): | |
| super().__init__(beta=beta, label_smoothing=label_smoothing) | |
| self.beta = beta | |
| self.reference_model = reference_model | |
| self.label_smoothing = label_smoothing | |
| def on_train_begin(self, ctx): | |
| if self.reference_model is None: | |
| self.reference_model = ctx.get("reference_model") | |
| def train_step(self, ctx, batch): | |
| chosen_ids = batch.get("chosen_input_ids") | |
| rejected_ids = batch.get("rejected_input_ids") | |
| if chosen_ids is None or rejected_ids is None: | |
| raise ValueError("DPO requires chosen_input_ids and rejected_input_ids in the batch") | |
| chosen_mask = batch.get("chosen_attention_mask") | |
| rejected_mask = batch.get("rejected_attention_mask") | |
| policy_chosen = sequence_log_probs(ctx.model, chosen_ids, chosen_mask) | |
| policy_rejected = sequence_log_probs(ctx.model, rejected_ids, rejected_mask) | |
| if self.reference_model is not None: | |
| with torch.no_grad(): | |
| ref_chosen = sequence_log_probs(self.reference_model, chosen_ids, chosen_mask) | |
| ref_rejected = sequence_log_probs(self.reference_model, rejected_ids, rejected_mask) | |
| else: | |
| ref_chosen = torch.zeros_like(policy_chosen) | |
| ref_rejected = torch.zeros_like(policy_rejected) | |
| loss = dpo_loss(policy_chosen, policy_rejected, ref_chosen, ref_rejected, | |
| beta=self.beta, label_smoothing=self.label_smoothing) | |
| loss.backward() | |
| return loss | |