| 68 | full_state_update = False |
| 69 | |
| 70 | def __init__(self, dist_sync_on_step: bool = False, ignore_index: int = -100): |
| 71 | super().__init__(dist_sync_on_step=dist_sync_on_step) |
| 72 | |
| 73 | if CrossEntropyLoss is None: |
| 74 | raise ImportError("flash_attn is not installed. Please install flash_attn to use FALanguageCrossEntropy.") |
| 75 | |
| 76 | self.ignore_index = ignore_index |
| 77 | self.loss_fn = CrossEntropyLoss(ignore_index=ignore_index, reduction="sum") |
| 78 | self.add_state("sum_loss", default=torch.tensor(0.0), dist_reduce_fx="sum") |
| 79 | self.add_state("total_items", default=torch.tensor(0), dist_reduce_fx="sum") |
| 80 | |
| 81 | |
| 82 | @rename_class("LanguageCrossEntropy") |