Use FA's CrossEntropyLoss for metrics calculation

composer/metrics/nlp.py

@@ -102,7 +102,15 @@ def __init__(self, dist_sync_on_step: bool = False, ignore_index: int = -100):
         super().__init__(dist_sync_on_step=dist_sync_on_step)
 
         self.ignore_index = ignore_index
-        self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduction='sum')
+        try:
+            from flash_attn.losses.cross_entropy import CrossEntropyLoss as FusedCrossEntropyLoss
+            self.loss_fn = FusedCrossEntropyLoss(ignore_index=ignore_index, reduction='sum')
+        except:
+            log.debug(
+                'Package `flash_attn` not installed. Using torch.nn.CrossEntropyLoss ' +
+                'to compute LanguageCrossEntropy metric, which will be slower.',
+            )
+            self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduction='sum')
         self.add_state('sum_loss', default=torch.tensor(0.), dist_reduce_fx='sum')
         self.add_state('total_items', default=torch.tensor(0), dist_reduce_fx='sum')
 

pyproject.toml

@@ -1,6 +1,6 @@
 # build requirements
 [build-system]
-requires = ["setuptools < 68.0.0"]
+requires = ["setuptools < 68.0.0", "packaging >= 21.3.0, < 24.1"]
 build-backend = "setuptools.build_meta"
 
 # iSort

setup.py

@@ -225,6 +225,10 @@ def package_files(prefix: str, directory: str, extension: str):
     'mlflow>=2.11.1,<3.0',
 ]
 
+extra_deps['flash-attn'] = [
+    'flash-attn==2.5.8',
+]
+
 extra_deps['pandas'] = ['pandas>=2.0.0,<3.0']
 
 extra_deps['databricks'] = ['databricks-sdk==0.25.1']