Spaces:
Sleeping
Sleeping
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| Training & Knowledge Distillation Pipelines for Q-TensorFormer (`src/training.py`). | |
| Provides: | |
| - `LanguageModelTrainer`: Deterministic causal LM trainer with AdamW, cosine LR schedule, | |
| gradient clipping, validation tracking, and full training history logging. | |
| - `DistillationTrainer`: Teacher-student knowledge distillation trainer for `scripts/distill.py`. | |
| """ | |
| import math | |
| import time | |
| from typing import Any, Dict, List, Optional | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| from .config import TrainingConfig | |
| from .metrics import evaluate_model | |
| class LanguageModelTrainer: | |
| """ | |
| Canonical causal language modeling trainer for Q-TensorFormer and DenseBaseline. | |
| """ | |
| def __init__( | |
| self, | |
| model: nn.Module, | |
| train_config: Optional[TrainingConfig] = None, | |
| device: str = "cpu", | |
| pad_token_id: int = 0, | |
| ): | |
| self.model = model.to(device) | |
| self.cfg = train_config or TrainingConfig() | |
| self.device = device | |
| self.pad_token_id = pad_token_id | |
| self.optimizer = torch.optim.AdamW( | |
| self.model.parameters(), | |
| lr=self.cfg.learning_rate, | |
| weight_decay=self.cfg.weight_decay, | |
| ) | |
| def _get_lr(self, step: int, total_steps: int) -> float: | |
| base_lr = self.cfg.learning_rate | |
| min_lr = base_lr * self.cfg.lr_min_factor | |
| warmup = min(self.cfg.warmup_steps, max(1, total_steps // 5)) | |
| if step < warmup: | |
| return base_lr * float(step + 1) / float(max(1, warmup)) | |
| if self.cfg.lr_scheduler == "cosine": | |
| progress = float(step - warmup) / float(max(1, total_steps - warmup)) | |
| progress = min(1.0, max(0.0, progress)) | |
| return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress)) | |
| if self.cfg.lr_scheduler == "linear": | |
| progress = float(step - warmup) / float(max(1, total_steps - warmup)) | |
| progress = min(1.0, max(0.0, progress)) | |
| return base_lr - progress * (base_lr - min_lr) | |
| return base_lr | |
| def train( | |
| self, | |
| train_loader: DataLoader, | |
| val_loader: Optional[DataLoader] = None, | |
| verbose: bool = True, | |
| ) -> Dict[str, Any]: | |
| torch.manual_seed(self.cfg.seed) | |
| total_steps = max(1, self.cfg.max_epochs * len(train_loader)) | |
| initial_val = ( | |
| evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id) | |
| if val_loader is not None | |
| else evaluate_model(self.model, train_loader, device=self.device, pad_token_id=self.pad_token_id) | |
| ) | |
| epoch_history: List[Dict[str, Any]] = [] | |
| global_step = 0 | |
| t_start = time.perf_counter() | |
| for epoch in range(1, self.cfg.max_epochs + 1): | |
| self.model.train() | |
| epoch_nll = 0.0 | |
| epoch_tokens = 0 | |
| t_ep = time.perf_counter() | |
| self.optimizer.zero_grad() | |
| for b_idx, batch in enumerate(train_loader): | |
| if len(batch) == 3: | |
| x, y, attn_mask = batch | |
| attn_mask = attn_mask.to(self.device) | |
| else: | |
| x, y = batch | |
| attn_mask = (x != self.pad_token_id).long().to(self.device) | |
| x = x.to(self.device) | |
| y = y.to(self.device) | |
| lr = self._get_lr(global_step, total_steps) | |
| for pg in self.optimizer.param_groups: | |
| pg["lr"] = lr | |
| try: | |
| out = self.model(x, attention_mask=attn_mask) | |
| except TypeError: | |
| out = self.model(x) | |
| logits = out[0] if isinstance(out, tuple) else (out.logits if hasattr(out, "logits") else out) | |
| vocab_size = logits.size(-1) | |
| valid_mask = y != self.pad_token_id | |
| n_valid = int(valid_mask.sum().item()) | |
| if n_valid == 0: | |
| continue | |
| loss = F.cross_entropy( | |
| logits.reshape(-1, vocab_size), | |
| y.reshape(-1), | |
| ignore_index=self.pad_token_id, | |
| reduction="mean", | |
| ) | |
| scaled_loss = loss / max(1, self.cfg.gradient_accumulation_steps) | |
| scaled_loss.backward() | |
| if (b_idx + 1) % max(1, self.cfg.gradient_accumulation_steps) == 0 or (b_idx + 1) == len(train_loader): | |
| if self.cfg.max_grad_norm > 0: | |
| torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.max_grad_norm) | |
| self.optimizer.step() | |
| self.optimizer.zero_grad() | |
| epoch_nll += float(loss.item()) * n_valid | |
| epoch_tokens += n_valid | |
| global_step += 1 | |
| train_loss = epoch_nll / max(1, epoch_tokens) | |
| train_ppl = math.exp(min(train_loss, 20.0)) | |
| val_metrics = ( | |
| evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id) | |
| if val_loader is not None | |
| else {"loss": train_loss, "perplexity": train_ppl, "accuracy": 0.0} | |
| ) | |
| ep_time = time.perf_counter() - t_ep | |
| rec = { | |
| "epoch": epoch, | |
| "train_loss": round(train_loss, 6), | |
| "train_perplexity": round(train_ppl, 4), | |
| "val_loss": val_metrics["loss"], | |
| "val_perplexity": val_metrics["perplexity"], | |
| "val_accuracy": val_metrics["accuracy"], | |
| "lr": round(self.optimizer.param_groups[0]["lr"], 8), | |
| "epoch_time_s": round(ep_time, 3), | |
| } | |
| epoch_history.append(rec) | |
| if verbose: | |
| print( | |
| f" Epoch {epoch:02d}/{self.cfg.max_epochs:02d} | " | |
| f"train_loss={rec['train_loss']:.4f} (ppl={rec['train_perplexity']:.2f}) | " | |
| f"val_loss={rec['val_loss']:.4f} (ppl={rec['val_perplexity']:.2f}, acc={rec['val_accuracy']:.3f}) | " | |
| f"{rec['epoch_time_s']:.2f}s" | |
| ) | |
| total_time = time.perf_counter() - t_start | |
| final_val = ( | |
| evaluate_model(self.model, val_loader, device=self.device, pad_token_id=self.pad_token_id) | |
| if val_loader is not None | |
| else evaluate_model(self.model, train_loader, device=self.device, pad_token_id=self.pad_token_id) | |
| ) | |
| return { | |
| "initial_val": initial_val, | |
| "final_val": final_val, | |
| "epochs": epoch_history, | |
| "total_steps": global_step, | |
| "total_training_time_s": round(total_time, 3), | |
| "loss_improved": final_val["loss"] < initial_val["loss"], | |
| "perplexity_reduction_ratio": round( | |
| initial_val["perplexity"] / max(1e-6, final_val["perplexity"]), 4 | |
| ), | |
| } | |
| class DistillationTrainer: | |
| """ | |
| Knowledge distillation trainer for `scripts/distill.py`. | |
| Minimizes `alpha * KL(student || teacher) * T^2 + (1 - alpha) * CE(student, y)`. | |
| """ | |
| def __init__( | |
| self, | |
| student: nn.Module, | |
| teacher: nn.Module, | |
| train_config: Optional[TrainingConfig] = None, | |
| config: Optional[Any] = None, | |
| train_loader: Optional[DataLoader] = None, | |
| val_loader: Optional[DataLoader] = None, | |
| test_loader: Optional[DataLoader] = None, | |
| output_dir: Optional[str] = None, | |
| temperature: float = 2.0, | |
| alpha: float = 0.7, | |
| device: str = "cpu", | |
| pad_token_id: int = 0, | |
| **kwargs, | |
| ): | |
| self.student = student.to(device) | |
| self.teacher = teacher.to(device) | |
| self.teacher.eval() | |
| self.train_loader = train_loader | |
| self.val_loader = val_loader | |
| self.test_loader = test_loader | |
| self.output_dir = output_dir | |
| if train_config is not None: | |
| self.cfg = train_config | |
| elif config is not None and hasattr(config, "training"): | |
| self.cfg = config.training | |
| elif config is not None and isinstance(config, TrainingConfig): | |
| self.cfg = config | |
| else: | |
| self.cfg = TrainingConfig() | |
| self.temperature = float(temperature) | |
| self.alpha = float(alpha) | |
| self.device = device | |
| self.pad_token_id = pad_token_id | |
| self.optimizer = torch.optim.AdamW( | |
| self.student.parameters(), | |
| lr=self.cfg.learning_rate, | |
| weight_decay=self.cfg.weight_decay, | |
| ) | |
| def train( | |
| self, | |
| train_loader: Optional[DataLoader] = None, | |
| val_loader: Optional[DataLoader] = None, | |
| ) -> Dict[str, Any]: | |
| tl = train_loader if train_loader is not None else self.train_loader | |
| vl = val_loader if val_loader is not None else self.val_loader | |
| if tl is None: | |
| raise ValueError("train_loader must be provided either in __init__ or train().") | |
| torch.manual_seed(self.cfg.seed) | |
| history: List[Dict[str, Any]] = [] | |
| for epoch in range(1, self.cfg.max_epochs + 1): | |
| self.student.train() | |
| total_loss = 0.0 | |
| n_batches = 0 | |
| for batch in tl: | |
| x, y = batch[0].to(self.device), batch[1].to(self.device) | |
| self.optimizer.zero_grad() | |
| with torch.no_grad(): | |
| t_out = self.teacher(x) | |
| t_logits = t_out[0] if isinstance(t_out, tuple) else t_out | |
| s_out = self.student(x) | |
| s_logits = s_out[0] if isinstance(s_out, tuple) else s_out | |
| vocab_size = s_logits.size(-1) | |
| ce_loss = F.cross_entropy( | |
| s_logits.reshape(-1, vocab_size), | |
| y.reshape(-1), | |
| ignore_index=self.pad_token_id, | |
| ) | |
| T = self.temperature | |
| kd_loss = F.kl_div( | |
| F.log_softmax(s_logits / T, dim=-1), | |
| F.softmax(t_logits / T, dim=-1), | |
| reduction="batchmean", | |
| ) * (T * T) | |
| loss = self.alpha * kd_loss + (1.0 - self.alpha) * ce_loss | |
| loss.backward() | |
| if self.cfg.max_grad_norm > 0: | |
| torch.nn.utils.clip_grad_norm_(self.student.parameters(), self.cfg.max_grad_norm) | |
| self.optimizer.step() | |
| total_loss += float(loss.item()) | |
| n_batches += 1 | |
| val_metrics = ( | |
| evaluate_model(self.student, vl, device=self.device, pad_token_id=self.pad_token_id) | |
| if vl is not None | |
| else {"loss": total_loss / max(1, n_batches), "perplexity": 0.0} | |
| ) | |
| history.append({ | |
| "epoch": epoch, | |
| "train_loss": round(total_loss / max(1, n_batches), 4), | |
| "val_loss": val_metrics["loss"], | |
| "val_perplexity": val_metrics["perplexity"], | |
| }) | |
| return {"epochs": history} | |