from typing import Tuple import os import tqdm import torch import torch.nn as nn from torch.utils.data import DataLoader from transformers.models.bert import ( BertForPreTraining, BertTokenizer, BertConfig ) from muon import SingleDeviceMuonWithAuxAdam class Dataset(): def __init__(self, file_path, tokenizer, min_length=32, max_length=512): self.sequences = [] self.tokenizer = tokenizer self.min_length = min_length self.max_length = max_length self._load_data(file_path) def _load_data(self, file_path): with open(file_path, "rb") as f: n_lines = (sum(1 for _ in f)) with open(file_path, "r") as f: all_headers = [] all_sequences = [] with tqdm.tqdm(total=n_lines) as pbar: for line in f: line = line.rstrip("\n") if line.startswith(">"): all_headers.append(line.lstrip(">")) all_sequences.append("") else: all_sequences[-1] += line pbar.update(1) self.sequences = [s for s in all_sequences if self.min_length - 2 <= len(s) <= self.max_length - 2] def __len__(self): return len(self.sequences) def __getitem__(self, idx): sequence = self.sequences[idx] tokens = self.tokenizer( " ".join(list(sequence)), max_length=self.max_length, padding="max_length", return_tensors="pt" ) return tokens class DistillationLoss(nn.Module): def __init__(self, alpha: float, temperature: float, num_labels: int, ignore_index: int): super().__init__() self.alpha = alpha self.temperature = temperature self.num_labels = num_labels self.ignore_index = ignore_index self.soft_loss_fn = nn.KLDivLoss(reduction="batchmean", log_target=True) self.hard_loss_fn = nn.CrossEntropyLoss(ignore_index=ignore_index) def forward(self, student_logits, teacher_logits, labels): """ Compute the distillation loss. Args: student_logits (torch.Tensor): Logits from the student model. teacher_logits (torch.Tensor): Logits from the teacher model. labels (torch.Tensor): Ground truth labels. Returns: torch.Tensor: The computed distillation loss. """ # Soft loss soft_loss = self.soft_loss_fn( nn.LogSoftmax(dim=-1)(student_logits / self.temperature), nn.LogSoftmax(dim=-1)(teacher_logits / self.temperature) ) # Hard loss hard_loss = self.hard_loss_fn( student_logits.view(-1, self.num_labels), labels.view(-1) ) return self.alpha * hard_loss + (1 - self.alpha) * soft_loss def main( model_name: str, student_hidden_size: int, student_intermediate_size: int, student_num_attention_heads: int, student_num_hidden_layers: int, train_data_path: str, batch_size: int, epochs: int, lr: float, default_lr: float, teacher_use_bf16: bool=True, teacher_use_sdpa: bool=True, min_length: int=32, max_length: int=512, alpha: float=0.1, temperature: float=10.0, use_muon: bool=True, weight_decay: float=0.01, betas: Tuple[float, float]=(0.9,0.95), default_weight_decay: float=0.01, default_betas: Tuple[float, float]=(0.9, 0.95), device: torch.device=torch.device("cpu"), num_workers: int=256, wandb_entity: str="giottonini-axel-unibe" ): import wandb wandb_project = f"{model_name.replace('/', "_")}-distilled" wandb_run = wandb.init( entity=wandb_entity, project=wandb_project, config=dict( hidden_size=student_hidden_size, intermediate_size=student_intermediate_size, num_attention_heads=student_num_attention_heads, num_hidden_layers=student_num_hidden_layers, alpha=alpha, temperature=temperature, use_muon=use_muon, lr=lr, weight_decay=weight_decay, betas=betas, default_lr=default_lr, default_weight_decay=default_weight_decay, default_betas=default_betas ) ) # Initialize tokenizer, teacher model and student model tokenizer = BertTokenizer.from_pretrained(model_name) teacher_model_kwargs = dict() if teacher_use_bf16: teacher_model_kwargs["torch_dtype"] = torch.bfloat16 if teacher_use_sdpa: teacher_model_kwargs["attn_implementation"] = "sdpa" teacher_model = BertForPreTraining.from_pretrained(model_name, **teacher_model_kwargs) teacher_model_compiled = torch.compile(teacher_model, mode="max-autotune", fullgraph=True) student_config = BertConfig.from_pretrained( "Rostlab/prot_bert", hidden_size=student_hidden_size, intermediate_size=student_intermediate_size, num_attention_heads=student_num_attention_heads, num_hidden_layers=student_num_hidden_layers ) student_model = BertForPreTraining(student_config) teacher_model_compiled.to(device) # type: ignore student_model.to(device) # type: ignore # Load dataset dataset = Dataset(train_data_path, tokenizer, min_length=min_length, max_length=max_length) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True) # type: ignore # Loss function loss_fn = DistillationLoss(alpha, temperature, tokenizer.vocab_size, tokenizer.pad_token_type_id) # Initialize optimizer hidden_weights = [p for p in student_model.bert.encoder.parameters() if p.ndim >= 2] hidden_biases = [p for p in student_model.bert.encoder.parameters() if p.ndim < 2] nonhidden_params = [ *student_model.bert.embeddings.parameters(), *student_model.cls.parameters() ] hidden_param_group = dict( params=hidden_weights, use_muon=use_muon, lr=lr, weight_decay=weight_decay ) if not use_muon: hidden_param_group["betas"] = betas # type: ignore default_param_group = dict( params=hidden_biases + nonhidden_params, use_muon=False, lr=default_lr, betas=default_betas, weight_decay=default_weight_decay ) optimizer = SingleDeviceMuonWithAuxAdam([hidden_param_group, default_param_group]) # Training loop wandb_run.watch(student_model) step = 0 for epoch in range(epochs): with tqdm.tqdm(dataloader) as pbar: for batch in pbar: # Clear optimizer and model gradients optimizer.zero_grad() student_model.zero_grad() # Send the data to the device batch = {k : v.squeeze(1).to(device) for k, v in batch.items()} # Compute teacher logits with torch.no_grad(): teacher_logits = teacher_model( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"] ).prediction_logits # Compute student logits student_logits = student_model( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"] ).prediction_logits # Loss backpropagation and optimization step loss = loss_fn(student_logits, teacher_logits, batch["input_ids"]) loss.backward() optimizer.step() step += 1 pbar.set_description(f"Epoch {epoch} | Loss: {loss.item():.4f}") wandb_run.log(dict(loss=loss.item())) # Save checkpoint if step % 1000 == 0: checkpoint = dict( state_dict=student_model.state_dict(), optimizer=optimizer.state_dict(), ) os.makedirs(os.path.join(wandb_run.project, wandb_run.name, str(step))) # type: ignore torch.save( checkpoint, os.path.join(wandb_run.project, wandb_run.name, str(step), "checkpoint.pt") # type: ignore ) if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--model_name", type=str, default="Rostlab/prot_bert", help="Name of the teacher model") parser.add_argument("--student_hidden_size", type=int, default=16, help="Hidden size of the student model") parser.add_argument("--student_intermediate_size", type=int, default=64, help="Intermediate size of the studen model") parser.add_argument("--student_num_attention_heads", type=int, default=4, help="Number of attention heads of the student model") parser.add_argument("--student_num_hidden_layers", type=int, default=12, help="Number of hidden layers of the student model") parser.add_argument("--train_data_path", type=str, help="Path to the training data (fasta file)") parser.add_argument("--batch_size", type=int, default=1024, help="Batch size for training") parser.add_argument("--epochs", type=int, default=3, help="Number of epochs for training") parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate for the hidden parameters") parser.add_argument("--default_lr", type=float, default=1e-4, help="Learning rate for the non-hidden parameters and biases") parser.add_argument("--device", type=int, default=-1, help="GPU device to use (-1 for CPU)") args = vars(parser.parse_args()) device = torch.device(f"cuda:{int(args['device'])}" if torch.cuda.is_available() and int(args['device']) >= 0 else "cpu") main( model_name=args["model_name"], student_hidden_size=args["student_hidden_size"], student_intermediate_size=args["student_intermediate_size"], student_num_attention_heads=args["student_num_attention_heads"], student_num_hidden_layers=args["student_num_hidden_layers"], train_data_path=args["train_data_path"], batch_size=args["batch_size"], epochs=args["epochs"], lr=args["lr"], default_lr=args["default_lr"], device=device )