Log documents_seen / num_documents and offer it as a W&B x-axis

fast_llm/data/document/config.py

@@ -30,11 +30,6 @@ class LengthPreprocessingConfig(BatchPreprocessingConfig):
     return_position_index: bool = Field(default=False)
 
 
-@config_class()
-class TokenPreprocessingConfig(LengthPreprocessingConfig):
-    return_document_count: bool = Field(default=False)
-
-
 @config_class()
 class ImageNormalizationConfig(Config):
     scale: float = Field(default=255.0)
@@ -68,7 +63,7 @@ def get_batch_meta(self, size: int = 1) -> "PatchBatch":
 
 
 @config_class()
-class LanguageModelBatchPreprocessingConfig(TokenPreprocessingConfig):
+class LanguageModelBatchPreprocessingConfig(LengthPreprocessingConfig):
     _abstract = False
     phase: PhaseType = Field(default=PhaseType.training)
     micro_batch_splits: int = Field(default=1)

fast_llm/data/document/token.py

@@ -6,7 +6,7 @@
 from fast_llm.core.distributed import allreduce_scalar
 from fast_llm.data.document.abstract import Batch, Document
 from fast_llm.data.document.block import BlockModelInput, LengthModelInputPreprocessor
-from fast_llm.data.document.config import TokenPreprocessingConfig
+from fast_llm.data.document.config import LengthPreprocessingConfig
 from fast_llm.engine.distributed.distributed import Distributed
 from fast_llm.layers.language_model.config import LanguageModelKwargs
 from fast_llm.tensor import TensorMeta
@@ -36,7 +36,7 @@ class TokenModelInput(BlockModelInput, TokenDocument):
 
     @classmethod
     def share_batch_data(cls, model_inputs: "list[TokenModelInput]", distributed: "Distributed"):
-        if model_inputs[0].num_documents is not None and model_inputs[0].num_documents_in_batch is None:
+        if model_inputs[0].num_documents_in_batch is None:
             # We sum over sequences but not within a sequence.
             num_documents_in_batch = allreduce_scalar(
                 sum(model_input.num_documents for model_input in model_inputs),
@@ -98,18 +98,17 @@ def _get_cropped_lengths(self, begin: int, end: int) -> tuple[list[int], int, in
 
         return lengths, first_document_begin, document_end
 
-    def _get_model_input(self, begin: int, end: int, config: TokenPreprocessingConfig, *, is_first_for_rank: bool):
+    def _get_model_input(self, begin: int, end: int, config: LengthPreprocessingConfig, *, is_first_for_rank: bool):
         model_input = self._model_input_class(tokens=self.tokens[begin:end])
         lengths, first_document_begin, last_document_end = self._get_cropped_lengths(begin, end)
 
-        if config.return_document_count:
-            # Set the global whole-batch count on every rank's first microbatch; `share_batch_data`
-            # will sum across DP via `batch_data_group`. (`begin == 0` would only set it on the
-            # SDP rank-0 first microbatch, leaving other SDP ranks with 0 after the DP-only sum.)
-            # Exclude the padding "length" from the document count.
-            model_input.num_documents = (
-                len(self.lengths) - (1 if self.unpadded_length < len(self.tokens) else 0) if is_first_for_rank else 0
-            )
+        # Number of real documents on this rank (excluding the padding "length").
+        num_documents = len(self.lengths) - (1 if self.unpadded_length < len(self.tokens) else 0)
+
+        # Set this rank's document count on its first microbatch; `share_batch_data` DP-sums these
+        # contributions into the whole-batch count. (`begin == 0` would only set it on the SDP
+        # rank-0 first microbatch, leaving other SDP ranks with 0 after the DP-only sum.)
+        model_input.num_documents = num_documents if is_first_for_rank else 0
 
         if config.return_document_index:
             # Globally-consistent 1-based document index per token, computed from the unsliced
@@ -124,10 +123,7 @@ def _get_model_input(self, begin: int, end: int, config: TokenPreprocessingConfi
                 side="right",
                 out_int32=True,
             )
-            # Exclude the padding "length" from the count.
-            model_input.num_documents_in_sequence = len(self.lengths) - (
-                1 if self.unpadded_length < len(self.tokens) else 0
-            )
+            model_input.num_documents_in_sequence = num_documents
 
         LengthModelInputPreprocessor(
             lengths=lengths,

fast_llm/engine/evaluation/evaluator.py

@@ -121,7 +121,7 @@ def run(
         begin_time = time.perf_counter()
         total_losses = {loss_def.name: 0.0 for loss_def in self._loss_definitions}
         for iter_ in range(self._config.iterations):
-            iter_losses, _, _ = self._runner.run_step(
+            iter_losses, _, _, _ = self._runner.run_step(
                 self._data_iterator, self._schedule, iteration=completed_evaluation_steps + iter_
             )
             for name, value in iter_losses.items():

fast_llm/engine/inference/runner.py

@@ -61,7 +61,7 @@ def forward(
         self, model_input: ModelInput, *, iteration: int = 1, return_metrics: bool = False
     ) -> tuple[dict[str, float | int], dict[str, typing.Any] | None]:
         # TODO: Return an actual model output.
-        reduced_losses, update_successful, metrics = self._runner.run_step(
+        reduced_losses, update_successful, metrics, _ = self._runner.run_step(
             iter(((model_input,),)),
             self._schedule,
             iteration=iteration,

fast_llm/engine/schedule/runner.py

@@ -66,6 +66,8 @@ def __repr__(self):
 
 class ScheduleRunner[ConfigType: ScheduleConfig](Configurable[ConfigType]):
     _is_setup: bool = False
+    # Whole-step document count (DP-summed) from the last `run_step`.
+    _num_documents_in_batch: int
     _compute_stream: torch.cuda.Stream | MockStream
     _data_stream: torch.cuda.Stream | MockStream
     _pipeline_stream: torch.cuda.Stream | MockStream
@@ -148,7 +150,7 @@ def run_step(
         *,
         iteration: int = 1,
         return_metrics: bool = False,
-    ) -> tuple[dict[str, float | int], bool, dict[str, typing.Any] | None]:
+    ) -> tuple[dict[str, float | int], bool, dict[str, typing.Any] | None, int]:
         assert self._is_setup
         assert schedule._config is self._config  # Noqa
         if schedule.phase.is_training:
@@ -222,7 +224,7 @@ def run_step(
         self._record_event(context, EventType.compute_wait_data, None)
 
         if not context.is_training or self._config.skip_step:
-            return self._reduce_losses(context), True, metrics
+            return self._reduce_losses(context), True, metrics, self._num_documents_in_batch
 
         for name, tied_parameter in self._tied_parameters.items():
             if tied_parameter.group is not None:
@@ -281,7 +283,7 @@ def run_step(
                 lambda: log_memory_usage(f"End of {context.phase} iteration {iteration}", str)
             )
 
-        return self._reduce_losses(context), update_successful, metrics
+        return self._reduce_losses(context), update_successful, metrics, self._num_documents_in_batch
 
     def _reduce_losses(self, context: BatchContext) -> dict[str, float | int]:
         reduced_losses = {
@@ -322,6 +324,7 @@ def _preprocess_data(
         model_inputs[0][0].share_batch_data(
             [model_input for model_inputs_ in model_inputs for model_input in model_inputs_], self._distributed
         )
+        self._num_documents_in_batch = model_inputs[0][0].num_documents_in_batch
 
         for micro_batch, model_inputs_ in enumerate(model_inputs):
             Assert.eq(len(model_inputs_), self._config.micro_batch_splits)

fast_llm/engine/training/trainer.py

@@ -39,6 +39,7 @@ class Trainer[ConfigType: TrainerConfig](Configurable[ConfigType], abc.ABC):
     _wandb: Wandb
     _optimizer: Optimizer | None
     _completed_steps: int
+    _documents_seen: int
     _schedule: Schedule
 
     def __init__(self, config: TrainerConfig):
@@ -220,13 +221,15 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
 
                 # TODO: Data loader hates getting all micro-batches at once.
                 #   (Also preprocessing adds overhead)
-                reduced_losses, update_successful, train_metrics = self._runner.run_step(
+                reduced_losses, update_successful, train_metrics, step_num_documents = self._runner.run_step(
                     train_iterator,
                     self._schedule,
                     iteration=self._completed_steps,
                     return_metrics=is_logging,
                 )
 
+                self._documents_seen += step_num_documents
+
                 # Advanced, skipped, and Nan iterations.
                 if update_successful:
                     advanced_iters += 1
@@ -262,6 +265,8 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                         metrics_key = PhaseType.training
                         metrics[metrics_key] = {
                             "batch_size": self._batch_size,
+                            "num_documents": step_num_documents,
+                            "documents_seen": self._documents_seen,
                             **{
                                 name: (value / advanced_iters if advanced_iters > 0 else float("nan"))
                                 for name, value in total_losses.items()
@@ -356,6 +361,7 @@ def _prepare_training_state(self) -> None:
             if self._do_train:
                 self._optimizer.reset_state()
             self._completed_steps = 0
+            self._documents_seen = 0
         else:
             log_main_rank(lambda: f"Loading checkpoint from iteration {last_iteration}...")
             self._load_checkpoint(self._config.training.checkpoint, last_iteration)
@@ -393,6 +399,7 @@ def _save_checkpoint(
         metadata = {
             "optimizer": self._optimizer.save(),
             "completed_steps": self._completed_steps,
+            "documents_seen": self._documents_seen,
         }
         if metrics is not None:
             metadata["metrics"] = metrics
@@ -436,6 +443,7 @@ def _load_checkpoint(self, config: TrainingCheckpointConfig, iteration: int) ->
             self._completed_steps = metadata["schedules"][PhaseType.training]["completed_steps"]
         else:
             self._completed_steps = metadata["completed_steps"]
+        self._documents_seen = metadata.get("documents_seen", 0)
         # TODO: Move barrier, ok file to FastLLMModel
         safe_barrier(
             self._distributed.world_group,

fast_llm/layers/language_model/loss/policy_gradient.py

@@ -136,7 +136,7 @@ def get_loss_definitions(self) -> list[LossDef]:
         return defs
 
     def get_preprocessing_config(self) -> dict[str, typing.Any]:
-        return {"use_grpo_data": True, "return_label_counts": True, "return_document_count": True}
+        return {"use_grpo_data": True, "return_label_counts": True}
 
     @functools.cached_property
     def _logprob_metric_name(self) -> str:

tests/data/test_preprocessing.py

@@ -63,7 +63,6 @@ class PreprocessingTestConfig:
     return_prediction_mask: bool = False
     return_label_counts: bool = False
     return_position_index: bool = False
-    return_document_count: bool = False
     return_cumulative_sequence_lengths: bool = False
 
     @functools.cached_property
@@ -76,7 +75,6 @@ def config_kwargs(self) -> dict:
             "return_prediction_mask": self.return_prediction_mask,
             "return_label_counts": self.return_label_counts,
             "return_position_index": self.return_position_index,
-            "return_document_count": self.return_document_count,
             "return_cumulative_sequence_lengths": self.return_cumulative_sequence_lengths,
         }
 
@@ -237,11 +235,8 @@ def expected_cumulative_lengths(self) -> list[tuple[torch.Tensor | None, torch.T
         return result
 
     @functools.cached_property
-    def expected_num_documents(self) -> list[int | None]:
-        if self.return_document_count:
-            return [len(self.tokens) if split_index == 0 else 0 for split_index in range(self.micro_batch_splits)]
-        else:
-            return [None] * self.micro_batch_splits
+    def expected_num_documents(self) -> list[int]:
+        return [len(self.tokens) if split_index == 0 else 0 for split_index in range(self.micro_batch_splits)]
 
 
 _BASE_TEST_CASES = [
@@ -313,13 +308,11 @@ def expected_num_documents(self) -> list[int | None]:
     "prediction_mask": {"return_prediction_mask": True},
     "label_counts": {"return_label_counts": True},
     "position_index": {"return_position_index": True},
-    "document_count": {"return_document_count": True},
     "cumulative_sequence_lengths": {"return_cumulative_sequence_lengths": True},
     "all": {
         "return_prediction_mask": True,
         "return_label_counts": True,
         "return_position_index": True,
-        "return_document_count": True,
         "return_cumulative_sequence_lengths": True,
     },
 }