fix: correct discriminative LR and early-exit training

bert_squeeze/assistants/configs/train_berxit.yaml

@@ -15,13 +15,9 @@ train:
   accumulation_steps: 1
   auto_lr: false
   discriminative_learning: true
-  # Training stage for BERxiT:
-  # - "backbone": train encoder + ramps + final classifier (no gate loss unless train_gates=true)
-  # - "gates": freeze backbone/ramps/classifier and train only gates
-  # You can also keep "backbone" here and use `switch_step` to switch to gate training mid-run.
+  # "backbone" trains the model; "gates" calibrates only the exit gate
   train_stage: "backbone"
-  # Optional global-step at which to switch from backbone to gate training within a single run.
-  # If null, no automatic switch is performed.
+  # Set a step to switch from backbone training to gate calibration.
   switch_step:
   dropout: 0.2
   layer_lr_decay: 0.95
@@ -37,12 +33,12 @@ train:
   warmup_steps: true
   weight_decay: 0.01
 
+  # Alternate between final-exit and all-exit objectives.
   train_highway: true
   early_exit_entropy: -1
-  # BERxiT-specific options
+  # Train the shared learning-to-exit gate.
   train_gates: true
-  gate_hidden_dim: 32
-  # Either a single float applied to all layers or a list of floats per layer
+  # Use one threshold for every layer or provide one value per layer.
   gate_thresholds: 0.5
 
 model:

bert_squeeze/assistants/configs/train_deebert.yaml

@@ -29,6 +29,7 @@ train:
   warmup_steps: true
   weight_decay: 0.01
 
+  # Set false for the backbone and final exit, or true for frozen-backbone exits.
   train_highway: true
   early_exit_entropy: -1
 
@@ -51,4 +52,4 @@ data:
     label_col: label
     truncate_mode: head
   tokenizer_name: ${model.pretrained_model}
-  max_length: 256
+  max_length: 256

bert_squeeze/distillation/base_distiller.py

@@ -1,13 +1,17 @@
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import lightning.pytorch as pl
 import numpy as np
 import torch
-from omegaconf import DictConfig, ListConfig
+from omegaconf import DictConfig
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from ..utils.experiment_logging import ExperimentLogger
-from ..utils.optimizers import BertAdam
+from ..utils.optimizers import (
+    BertAdam,
+    OptimizerParameterGroup,
+    build_optimizer_parameter_groups,
+)
 from ..utils.types import DistillationLoss
 
 
@@ -31,7 +35,7 @@ def __init__(
         teacher: Union["pl.LightningModule", "torch.nn.Module"],
         student: Union[pl.LightningModule, torch.nn.Module],
         training_config: DictConfig,
-        teacher_checkpoint: str = None,
+        teacher_checkpoint: Optional[str] = None,
         **kwargs,
     ):
         super().__init__()
@@ -52,107 +56,14 @@ def _set_scorers(self) -> None:
         """"""
         raise NotImplementedError()
 
-    def _get_student_parameters(self) -> List[Dict]:
-        """
-        Method that defines the student's parameters to optimize.
-
-        Returns:
-            List[Dict]: group of parameters to optimize
-        """
-        no_decay = ['bias', 'gamma', 'beta', 'LayerNorm.weight', 'layer_norm.weight']
-
-        if self.params.discriminative_learning:
-            if (
-                isinstance(self.params.learning_rates, ListConfig)
-                and len(self.params.learning_rates) > 1
-            ):
-                groups = [
-                    (f'layer.{i}.', self.params.learning_rates[i]) for i in range(12)
-                ]
-            else:
-                lr = (
-                    self.params.learning_rates[0]
-                    if isinstance(self.params.learning_rates, ListConfig)
-                    else self.params.learning_rates
-                )
-                groups = [
-                    (f'layer.{i}.', lr * pow(self.params.layer_lr_decay, 11 - i))
-                    for i in range(12)
-                ]
-
-            group_all = [f'layer.{i}.' for i in range(12)]
-            no_decay_optimizer_parameters, decay_optimizer_parameters = [], []
-            for g, l in groups:
-                no_decay_optimizer_parameters.append(
-                    {
-                        'params': [
-                            p
-                            for n, p in self.student.named_parameters()
-                            if not any(nd in n for nd in no_decay)
-                            and any(nd in n for nd in [g])
-                        ],
-                        'weight_decay': self.params.weight_decay,
-                        'lr': l,
-                    }
-                )
-                decay_optimizer_parameters.append(
-                    {
-                        'params': [
-                            p
-                            for n, p in self.student.named_parameters()
-                            if any(nd in n for nd in no_decay)
-                            and any(nd in n for nd in [g])
-                        ],
-                        'weight_decay': 0.0,
-                        'lr': l,
-                    }
-                )
-
-            group_all_parameters = [
-                {
-                    'params': [
-                        p
-                        for n, p in self.student.named_parameters()
-                        if not any(nd in n for nd in no_decay)
-                        and not any(nd in n for nd in group_all)
-                    ],
-                    'weight_decay': self.params.weight_decay,
-                },
-                {
-                    'params': [
-                        p
-                        for n, p in self.student.named_parameters()
-                        if any(nd in n for nd in no_decay)
-                        and not any(nd in n for nd in group_all)
-                    ],
-                    'weight_decay': 0.0,
-                },
-            ]
-            optimizer_grouped_parameters = (
-                no_decay_optimizer_parameters
-                + decay_optimizer_parameters
-                + group_all_parameters
-            )
-        else:
-            optimizer_grouped_parameters = [
-                {
-                    'params': [
-                        p
-                        for n, p in self.student.named_parameters()
-                        if not any(nd in n for nd in no_decay)
-                    ],
-                    'weight_decay': self.params.weight_decay,
-                },
-                {
-                    'params': [
-                        p
-                        for n, p in self.student.named_parameters()
-                        if any(nd in n for nd in no_decay)
-                    ],
-                    'weight_decay': 0.0,
-                },
-            ]
-        return optimizer_grouped_parameters
+    def _get_student_parameters(self) -> List[OptimizerParameterGroup]:
+        return build_optimizer_parameter_groups(
+            self.student.named_parameters(),
+            discriminative_learning=self.params.discriminative_learning,
+            learning_rates=self.params.learning_rates,
+            layer_lr_decay=self.params.get("layer_lr_decay", 1.0),
+            weight_decay=self.params.weight_decay,
+        )
 
     def configure_optimizers(self) -> Tuple[List, List]:
         """
@@ -188,13 +99,21 @@ def configure_optimizers(self) -> Tuple[List, List]:
             scheduler = ReduceLROnPlateau(optimizer)
             lr_scheduler = {
                 'scheduler': scheduler,
-                'name': 'NeptuneLogger',
-                'monitor': 'loss',
+                'name': 'learning_rate',
+                'monitor': self.params.get("lr_scheduler_monitor", "train/epoch_loss"),
             }
             return [optimizer], [lr_scheduler]
 
         return [optimizer], []
 
+    def _log_training_loss(self, loss: torch.Tensor) -> None:
+        self.log(
+            "train/epoch_loss",
+            loss,
+            on_step=False,
+            on_epoch=True,
+        )
+
     def training_step(self, batch, _) -> torch.Tensor:
         raise NotImplementedError()
 

bert_squeeze/distillation/seq2seq_distiller.py

@@ -1,9 +1,7 @@
-from typing import Any, Dict, TypeVar, Union
+from typing import Dict, Optional, Union
 
 import lightning.pytorch as pl
-import numpy as np
 import torch
-import torch.nn.functional as F
 from omegaconf import DictConfig
 from overrides import overrides
 from torch.nn import CrossEntropyLoss
@@ -35,7 +33,7 @@ def __init__(
         teacher: Union["pl.LightningModule", "torch.nn.Module"],
         student: Union["pl.LightningModule", "torch.nn.Module"],
         training_config: DictConfig,
-        teacher_checkpoint: str = None,
+        teacher_checkpoint: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(teacher, student, training_config, teacher_checkpoint, **kwargs)
@@ -127,6 +125,7 @@ def training_step(self, batch, _) -> torch.Tensor:
                 for key, val in self.s_scorer.losses.items()
             }
             self.log_dict(logging_loss)
+        self._log_training_loss(loss.full_loss)
         return loss.full_loss
 
     @overrides

bert_squeeze/distillation/sequence_classification_distiller.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import lightning.pytorch as pl
 import matplotlib.pyplot as plt
@@ -43,7 +43,7 @@ def __init__(
         student: Union["pl.LightningModule", "torch.nn.Module"],
         training_config: DictConfig,
         labels: Union[List[str], List[int]],
-        teacher_checkpoint: str = None,
+        teacher_checkpoint: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(teacher, student, training_config, teacher_checkpoint, **kwargs)
@@ -158,7 +158,7 @@ def __init__(
         student: Union["pl.LightningModule", "torch.nn.Module"],
         training_config: DictConfig,
         labels: Union[List[str], List[int]],
-        teacher_checkpoint: str = None,
+        teacher_checkpoint: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(
@@ -220,7 +220,7 @@ def loss(
         self,
         teacher_logits: torch.Tensor,
         student_logits: torch.Tensor,
-        labels: torch.Tensor = None,
+        labels: Optional[torch.Tensor] = None,
         ignore_index: int = -100,
         *args,
         **kwargs,
@@ -268,6 +268,7 @@ def training_step(self, batch, _) -> torch.Tensor:
             self.log_dict(logging_loss)
 
             self.log("train/acc", self.scorer.acc)
+        self._log_training_loss(loss.full_loss)
         return loss.full_loss
 
     @overrides
@@ -350,7 +351,7 @@ def __init__(
         student: Union["pl.LightningModule", "torch.nn.Module"],
         training_config: DictConfig,
         labels: Union[List[str], List[int]],
-        teacher_checkpoint: str = None,
+        teacher_checkpoint: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(
@@ -464,6 +465,7 @@ def training_step(self, batch, _) -> torch.Tensor:
         s_logits_original, s_logits_translated = self.get_student_logits(batch)
 
         loss = self.loss(t_logits, s_logits_original, s_logits_translated)
+        self._log_training_loss(loss.full_loss)
         return loss.full_loss
 
     @overrides