Add loss generating token counts

llmfoundry/data/utils.py

@@ -83,7 +83,7 @@ def get_data_spec(
 
 def get_tokens_per_batch_func(
     decoder_only: bool = True,
-) -> Callable[[Batch], int]:
+) -> Callable[[Batch], Union[int, dict[str, int]]]:
     """Returns a callable that counts the number of tokens in a batch.
 
     Args:
@@ -114,13 +114,28 @@ def get_num_tokens_in_batch(batch: Batch) -> int:
         else:
             input_ids_tokens = batch['input_ids'].numel()
 
+        loss_generating_tokens = 0
+        if 'labels' in batch:
+            loss_generating_tokens = int(
+                torch.sum(batch['labels'] != -100).item(),
+            )
+
+            # Subtract one for each example in the batch because the labels
+            # will be shifted by one
+            loss_generating_tokens -= batch['labels'].shape[0]
+
         # For encoder decoder models only
         decoder_input_ids_tokens = 0
         if not decoder_only:
             decoder_input_ids_tokens = int(
                 torch.sum(batch['decoder_attention_mask']).item(),
             )
 
+        if loss_generating_tokens != 0:
+            return {
+                'total': input_ids_tokens + decoder_input_ids_tokens,
+                'loss_generating': loss_generating_tokens,
+            }
         return input_ids_tokens + decoder_input_ids_tokens
 
     return get_num_tokens_in_batch

tests/data/test_dataloader.py

@@ -1186,6 +1186,7 @@ def test_token_counting_func_dataloader_setting(
 
     batch_strings = []
     expected_token_count = 0
+    expected_loss_generating_token_count = 0
     for _ in range(batch_size):
         # Get randomly different lengths if we are going to add padding
         sample_length = random.randint(1, model_max_length // 4) if (
@@ -1208,8 +1209,14 @@ def test_token_counting_func_dataloader_setting(
         for b in batch_tokenized:
             b['labels'] = b['input_ids'].copy()  # type: ignore
         batch_tokenized = [{'turns': [b]} for b in batch_tokenized]
+        expected_loss_generating_token_count = expected_token_count
         expected_token_count *= 2
         expected_token_count += 1 * batch_size  # for the eos token
+        expected_loss_generating_token_count += 1 * batch_size  # for the eos token
+    else:
+        expected_loss_generating_token_count = expected_token_count
+
+    expected_loss_generating_token_count -= 1 * batch_size  # because the labels will be shifted
 
     common_args = {
         'drop_last': False,
@@ -1311,9 +1318,11 @@ def build_from_hf(
         raise NotImplementedError()
 
     batch_collated = dl.dataloader.collate_fn(batch_tokenized)  # type: ignore
-    actual_token_count = dl.get_num_tokens_in_batch(batch_collated)
+    actual_total_token_count = dl.get_num_tokens_in_batch(batch_collated, token_type='total')
+    actual_loss_generating_token_count = dl.get_num_tokens_in_batch(batch_collated, token_type='loss_generating')
 
-    assert actual_token_count == expected_token_count
+    assert actual_total_token_count == expected_token_count
+    assert actual_loss_generating_token_count == expected_loss_generating_token_count
 
 
 def test_build_unknown_dataloader():