Class EfficientZLoss

src/flex_bert.py:114–141

Torchmetric that grabs the precomputed z_loss value from the model outputs

@rename_class("ZLoss")
class EfficientZLoss(Metric):
    """Torchmetric that grabs the precomputed z_loss value from the model outputs"""

    # Make torchmetrics call update only once
    full_state_update = False

    def __init__(self, dist_sync_on_step: bool = False):
        super().__init__(dist_sync_on_step=dist_sync_on_step)
        self.add_state("sum_loss", default=torch.tensor(0.0), dist_reduce_fx="sum")
        self.add_state("total_items", default=torch.tensor(0), dist_reduce_fx="sum")

    def update(self, loss: Tensor) -> None:
        """Updates the internal state with results from a new batch.

        Args:
            loss (~torch.Tensor): A Tensor of loss values to compare against.
        """
        self.sum_loss += loss
        self.total_items += 1

    def compute(self) -> Tensor:
        """Aggregate the state over all processes to compute the metric.

        Returns:
            loss: The loss averaged across all batches as a :class:`~torch.Tensor`.
        """
        # Return average loss over entire dataset
        return self.sum_loss / self.total_items  # type: ignore (third-party)


class EfficientHuggingFaceModel(HuggingFaceModel):