Class FALanguageCrossEntropy

src/flex_bert.py:54–79

Torchmetric that computes cross entropy on language modeling outputs using flash_attn's Cross Entropy. Adds metric state variables: sum_loss (float): The sum of the per-example loss in the batch. total_items (float): The number of batches to average across. Args: di

@rename_class("LanguageCrossEntropy")
class FALanguageCrossEntropy(LanguageCrossEntropy):
    """Torchmetric that computes cross entropy on language modeling outputs using flash_attn's Cross Entropy.

    Adds metric state variables:
        sum_loss (float): The sum of the per-example loss in the batch.
        total_items (float): The number of batches to average across.

    Args:
        dist_sync_on_step (bool, optional): Synchronize metric state across processes at
            each forward() before returning the value at the step. Default: ``False``.
        ignore_index (int, optional): The class index to ignore. Default: ``-100``.
    """

    # Make torchmetrics call update only once
    full_state_update = False

    def __init__(self, dist_sync_on_step: bool = False, ignore_index: int = -100):
        super().__init__(dist_sync_on_step=dist_sync_on_step)

        if CrossEntropyLoss is None:
            raise ImportError("flash_attn is not installed. Please install flash_attn to use FALanguageCrossEntropy.")

        self.ignore_index = ignore_index
        self.loss_fn = CrossEntropyLoss(ignore_index=ignore_index, reduction="sum")
        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")


@rename_class("LanguageCrossEntropy")