From ab176fbce0e1daf8cb877f0998e701af2c38b739 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Tue, 7 Jul 2026 14:14:38 -0400
Subject: [PATCH 01/10] Log documents_seen / num_documents and offer it as a
 W&B x-axis

Reintroduce the per-step document count and cumulative documents-seen
metrics on the static-schedule path, independent of `docs_per_step`.

`ScheduleRunner` exposes the whole-step DP-summed document count
(`num_documents_in_batch`, the loss-normalization divisor) as
`_num_documents_in_batch`, set once per step in `_preprocess_data`; it is
`None` when the data carries no document counts (non-RL runs). The trainer
accumulates it into `_documents_seen` every step (persisted in checkpoint
metadata) and logs `num_documents` + `documents_seen` when available.

`Wandb` lazily declares `training/documents_seen` as an alternative
`step_metric` for `training/*` on the first log that carries it, so it can
be used as a cross-run x-axis without affecting non-RL runs. The global
`wandb.log(step=...)` step is unchanged, so both axes remain available.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 fast_llm/engine/schedule/runner.py  |  5 +++++
 fast_llm/engine/training/trainer.py | 15 +++++++++++++++
 fast_llm/engine/training/wandb.py   | 15 +++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py
index 196ff1577..d9d17bc53 100644
--- a/fast_llm/engine/schedule/runner.py
+++ b/fast_llm/engine/schedule/runner.py
@@ -66,6 +66,9 @@ def __repr__(self):
 
 class ScheduleRunner[ConfigType: ScheduleConfig](Configurable[ConfigType]):
     _is_setup: bool = False
+    # Whole-step document count (DP-summed) from the last `run_step`, or None when the data does not
+    # provide document counts (i.e. no loss requested `return_document_count`).
+    _num_documents_in_batch: int | None = None
     _compute_stream: torch.cuda.Stream | MockStream
     _data_stream: torch.cuda.Stream | MockStream
     _pipeline_stream: torch.cuda.Stream | MockStream
@@ -322,6 +325,8 @@ 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
         )
+        # Whole-step DP-summed document count (the loss-normalization divisor), for `documents_seen`.
+        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)
diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
index 1ed18c449..e61296f3a 100644
--- a/fast_llm/engine/training/trainer.py
+++ b/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 = 0
     _schedule: Schedule
 
     def __init__(self, config: TrainerConfig):
@@ -227,6 +228,12 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                     return_metrics=is_logging,
                 )
 
+                # Cumulative document count (the RL x-axis / model-version clock); `None` when the
+                # data provides no document counts (non-RL runs).
+                step_num_documents = self._runner._num_documents_in_batch
+                if step_num_documents is not None:
+                    self._documents_seen += step_num_documents
+
                 # Advanced, skipped, and Nan iterations.
                 if update_successful:
                     advanced_iters += 1
@@ -257,6 +264,11 @@ 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}
+                                if step_num_documents is not None
+                                else {}
+                            ),
                             **{
                                 name: (value / advanced_iters if advanced_iters > 0 else float("nan"))
                                 for name, value in total_losses.items()
@@ -350,6 +362,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)
@@ -387,6 +400,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
@@ -430,6 +444,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,
diff --git a/fast_llm/engine/training/wandb.py b/fast_llm/engine/training/wandb.py
index 9257fe9c0..c60f0c992 100644
--- a/fast_llm/engine/training/wandb.py
+++ b/fast_llm/engine/training/wandb.py
@@ -5,6 +5,7 @@
 
 from fast_llm.config import Config
 from fast_llm.engine.config_utils.run import Run
+from fast_llm.engine.distributed.config import PhaseType
 from fast_llm.engine.training.config import WandbConfig
 
 
@@ -13,6 +14,9 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
         self._config = config
         self._is_setup = True
         self._run = run
+        # Whether the `documents_seen` custom x-axis has been declared (RL runs only, done lazily on
+        # the first log that carries it so non-RL runs are unaffected).
+        self._defined_documents_axis = False
         if self._config.entity_name is not None and self._run.is_main_rank:
             import wandb
             import wandb.sdk.lib.runid
@@ -51,6 +55,17 @@ def log_metrics(self, completed_steps: int, metrics: dict[str, dict[str, float |
         if self._wandb is not None:
             import wandb
 
+            if not self._defined_documents_axis and any(
+                isinstance(values, dict) and "documents_seen" in values for values in metrics.values()
+            ):
+                # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
+                # `wandb.log(step=...)` still records the global step, so both axes remain available.
+                self._wandb.define_metric(f"{PhaseType.training}/documents_seen")
+                self._wandb.define_metric(
+                    f"{PhaseType.training}/*", step_metric=f"{PhaseType.training}/documents_seen"
+                )
+                self._defined_documents_axis = True
+
             wandb.log(metrics, step=completed_steps, commit=commit)  # noqa
 
     def alert(self, title, text, level="INFO", wait=0.001) -> None:

From 68ad8a3e086214b84863abffdd69ccf312360662 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Wed, 8 Jul 2026 17:35:38 -0400
Subject: [PATCH 02/10] Review fixes for #559: return num_documents_in_batch
 from run_step, drop dead check, reword comments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Return `num_documents_in_batch` as a 4th element of `ScheduleRunner.run_step`'s tuple instead of
  having the trainer reach into the runner's private `_num_documents_in_batch` attribute — the
  runner already returns everything else it computes (`reduced_losses`, `update_successful`,
  `metrics`) this way, and the private attribute is also mutated by evaluation's `run_step` calls
  on the same shared runner later in the same iteration.
- Drop the `isinstance(values, dict)` guard in `Wandb.log_metrics`: `metrics` is always
  `dict[str, dict[str, float | int]]` per its own type signature.
- Reword comments that named "RL" as the consumer to describe the value instead.
- Dedupe the repeated `f"{PhaseType.training}/documents_seen"` literal into a local variable.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/engine/evaluation/evaluator.py |  2 +-
 fast_llm/engine/inference/runner.py     |  2 +-
 fast_llm/engine/schedule/runner.py      |  8 ++++----
 fast_llm/engine/training/trainer.py     |  7 +++----
 fast_llm/engine/training/wandb.py       | 15 ++++++---------
 5 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/fast_llm/engine/evaluation/evaluator.py b/fast_llm/engine/evaluation/evaluator.py
index 0f1fcda03..1ccecd5c5 100644
--- a/fast_llm/engine/evaluation/evaluator.py
+++ b/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():
diff --git a/fast_llm/engine/inference/runner.py b/fast_llm/engine/inference/runner.py
index d9ed695ec..23608dbdf 100644
--- a/fast_llm/engine/inference/runner.py
+++ b/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,
diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py
index d9d17bc53..e58d2252a 100644
--- a/fast_llm/engine/schedule/runner.py
+++ b/fast_llm/engine/schedule/runner.py
@@ -151,7 +151,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 | None]:
         assert self._is_setup
         assert schedule._config is self._config  # Noqa
         if schedule.phase.is_training:
@@ -225,7 +225,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:
@@ -284,7 +284,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 = {
@@ -325,7 +325,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
         )
-        # Whole-step DP-summed document count (the loss-normalization divisor), for `documents_seen`.
+        # Whole-step DP-summed document count (the loss-normalization divisor).
         self._num_documents_in_batch = model_inputs[0][0].num_documents_in_batch
 
         for micro_batch, model_inputs_ in enumerate(model_inputs):
diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
index e61296f3a..27bb898c3 100644
--- a/fast_llm/engine/training/trainer.py
+++ b/fast_llm/engine/training/trainer.py
@@ -221,16 +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,
                 )
 
-                # Cumulative document count (the RL x-axis / model-version clock); `None` when the
-                # data provides no document counts (non-RL runs).
-                step_num_documents = self._runner._num_documents_in_batch
+                # Cumulative document count across training steps; `None` when the data provides no
+                # document counts.
                 if step_num_documents is not None:
                     self._documents_seen += step_num_documents
 
diff --git a/fast_llm/engine/training/wandb.py b/fast_llm/engine/training/wandb.py
index c60f0c992..3a4ae81b1 100644
--- a/fast_llm/engine/training/wandb.py
+++ b/fast_llm/engine/training/wandb.py
@@ -14,8 +14,8 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
         self._config = config
         self._is_setup = True
         self._run = run
-        # Whether the `documents_seen` custom x-axis has been declared (RL runs only, done lazily on
-        # the first log that carries it so non-RL runs are unaffected).
+        # Whether the `documents_seen` custom x-axis has been declared. Done lazily on the first log
+        # that carries a `documents_seen` value, since not every run produces one.
         self._defined_documents_axis = False
         if self._config.entity_name is not None and self._run.is_main_rank:
             import wandb
@@ -55,15 +55,12 @@ def log_metrics(self, completed_steps: int, metrics: dict[str, dict[str, float |
         if self._wandb is not None:
             import wandb
 
-            if not self._defined_documents_axis and any(
-                isinstance(values, dict) and "documents_seen" in values for values in metrics.values()
-            ):
+            if not self._defined_documents_axis and any("documents_seen" in values for values in metrics.values()):
                 # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
                 # `wandb.log(step=...)` still records the global step, so both axes remain available.
-                self._wandb.define_metric(f"{PhaseType.training}/documents_seen")
-                self._wandb.define_metric(
-                    f"{PhaseType.training}/*", step_metric=f"{PhaseType.training}/documents_seen"
-                )
+                documents_seen_metric = f"{PhaseType.training}/documents_seen"
+                self._wandb.define_metric(documents_seen_metric)
+                self._wandb.define_metric(f"{PhaseType.training}/*", step_metric=documents_seen_metric)
                 self._defined_documents_axis = True
 
             wandb.log(metrics, step=completed_steps, commit=commit)  # noqa

From b74479739f25553c87bb7dee582508a78290c200 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Wed, 8 Jul 2026 17:57:06 -0400
Subject: [PATCH 03/10] Make document counting the default instead of an opt-in
 loss preprocessing flag

Document counting was only enabled indirectly, by whichever loss's preprocessing
config happened to request it (currently only GRPO/GSPO). Flip
`TokenPreprocessingConfig.return_document_count` to default `True` so
`documents_seen`/`num_documents` populate for any run, and drop the now-redundant
explicit request from the policy-gradient loss config.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/data/document/config.py                       | 2 +-
 fast_llm/layers/language_model/loss/policy_gradient.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/fast_llm/data/document/config.py b/fast_llm/data/document/config.py
index a90bcdebc..69463c718 100644
--- a/fast_llm/data/document/config.py
+++ b/fast_llm/data/document/config.py
@@ -32,7 +32,7 @@ class LengthPreprocessingConfig(BatchPreprocessingConfig):
 
 @config_class()
 class TokenPreprocessingConfig(LengthPreprocessingConfig):
-    return_document_count: bool = Field(default=False)
+    return_document_count: bool = Field(default=True)
 
 
 @config_class()
diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py
index c07c81e05..ce532f00d 100644
--- a/fast_llm/layers/language_model/loss/policy_gradient.py
+++ b/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:

From 64675ac7996318c85210545e7eb0b1ba29d39df7 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Wed, 8 Jul 2026 18:39:20 -0400
Subject: [PATCH 04/10] Remove now-vestigial return_document_count flag

Now that document counting is unconditional (previous commit), the flag itself
has no consumer that can turn it off. Drop the field and make the count
computation in TokenBatch._get_model_input unconditional; update the
preprocessing test's cross-product accordingly.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/data/document/config.py   |  2 +-
 fast_llm/data/document/token.py    | 15 +++++++--------
 fast_llm/engine/schedule/runner.py |  2 +-
 tests/data/test_preprocessing.py   | 11 ++---------
 4 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/fast_llm/data/document/config.py b/fast_llm/data/document/config.py
index 69463c718..4bc6bcdea 100644
--- a/fast_llm/data/document/config.py
+++ b/fast_llm/data/document/config.py
@@ -32,7 +32,7 @@ class LengthPreprocessingConfig(BatchPreprocessingConfig):
 
 @config_class()
 class TokenPreprocessingConfig(LengthPreprocessingConfig):
-    return_document_count: bool = Field(default=True)
+    pass
 
 
 @config_class()
diff --git a/fast_llm/data/document/token.py b/fast_llm/data/document/token.py
index 4d1af453d..2127a7724 100644
--- a/fast_llm/data/document/token.py
+++ b/fast_llm/data/document/token.py
@@ -102,14 +102,13 @@ def _get_model_input(self, begin: int, end: int, config: TokenPreprocessingConfi
         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
-            )
+        # 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
+        )
 
         if config.return_document_index:
             # Globally-consistent 1-based document index per token, computed from the unsliced
diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py
index e58d2252a..7ad56671b 100644
--- a/fast_llm/engine/schedule/runner.py
+++ b/fast_llm/engine/schedule/runner.py
@@ -67,7 +67,7 @@ def __repr__(self):
 class ScheduleRunner[ConfigType: ScheduleConfig](Configurable[ConfigType]):
     _is_setup: bool = False
     # Whole-step document count (DP-summed) from the last `run_step`, or None when the data does not
-    # provide document counts (i.e. no loss requested `return_document_count`).
+    # provide document counts (i.e. non-token data).
     _num_documents_in_batch: int | None = None
     _compute_stream: torch.cuda.Stream | MockStream
     _data_stream: torch.cuda.Stream | MockStream
diff --git a/tests/data/test_preprocessing.py b/tests/data/test_preprocessing.py
index ae58121ae..1c0d6acab 100644
--- a/tests/data/test_preprocessing.py
+++ b/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,
     },
 }

From b5e1dbd63e9b4a70517098fe627cf4045bf57752 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 13:44:51 -0400
Subject: [PATCH 05/10] Drop dead num_documents-is-not-None check in
 TokenModelInput.share_batch_data

num_documents is now always set by _get_model_input (no opt-in flag left to
gate it), and no other production or test path constructs a TokenModelInput
bypassing _get_model_input, so the check was always true.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/data/document/token.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fast_llm/data/document/token.py b/fast_llm/data/document/token.py
index 2127a7724..de2327125 100644
--- a/fast_llm/data/document/token.py
+++ b/fast_llm/data/document/token.py
@@ -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),

From a1b63a6943eaab6630b847294f871e26b352a373 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 13:58:04 -0400
Subject: [PATCH 06/10] Collapse now-empty TokenPreprocessingConfig into
 LengthPreprocessingConfig

TokenPreprocessingConfig added no fields of its own after the previous commit
removed return_document_count. Use LengthPreprocessingConfig (its only base)
directly, matching the sibling PatchPreprocessingConfig.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/data/document/config.py | 7 +------
 fast_llm/data/document/token.py  | 4 ++--
 2 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/fast_llm/data/document/config.py b/fast_llm/data/document/config.py
index 4bc6bcdea..ab184d4d8 100644
--- a/fast_llm/data/document/config.py
+++ b/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):
-    pass
-
-
 @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)
diff --git a/fast_llm/data/document/token.py b/fast_llm/data/document/token.py
index de2327125..c6da7101b 100644
--- a/fast_llm/data/document/token.py
+++ b/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
@@ -98,7 +98,7 @@ 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)
 

From 083dfabf3a6b83fa5ff11e31077e90a77f306b0b Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 14:01:48 -0400
Subject: [PATCH 07/10] Define the documents_seen W&B x-axis unconditionally in
 __init__

documents_seen is now always present in the training metrics (document
counting is unconditional), so the lazy define-on-first-log dance and its
_defined_documents_axis flag were dead weight. Declare the custom x-axis once,
at the same point as the rest of the one-time wandb setup.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/engine/training/wandb.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/fast_llm/engine/training/wandb.py b/fast_llm/engine/training/wandb.py
index 3a4ae81b1..819577490 100644
--- a/fast_llm/engine/training/wandb.py
+++ b/fast_llm/engine/training/wandb.py
@@ -14,9 +14,6 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
         self._config = config
         self._is_setup = True
         self._run = run
-        # Whether the `documents_seen` custom x-axis has been declared. Done lazily on the first log
-        # that carries a `documents_seen` value, since not every run produces one.
-        self._defined_documents_axis = False
         if self._config.entity_name is not None and self._run.is_main_rank:
             import wandb
             import wandb.sdk.lib.runid
@@ -47,6 +44,11 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
                     wandb_path.write_text(yaml.safe_dump(wandb_config))
             # TODO: Does wandb work with nested configs?
             self._wandb = wandb.init(config=experiment_config.to_dict(), **wandb_config)
+            # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
+            # `wandb.log(step=...)` still records the global step, so both axes remain available.
+            documents_seen_metric = f"{PhaseType.training}/documents_seen"
+            self._wandb.define_metric(documents_seen_metric)
+            self._wandb.define_metric(f"{PhaseType.training}/*", step_metric=documents_seen_metric)
         else:
             self._wandb = None
 
@@ -55,14 +57,6 @@ def log_metrics(self, completed_steps: int, metrics: dict[str, dict[str, float |
         if self._wandb is not None:
             import wandb
 
-            if not self._defined_documents_axis and any("documents_seen" in values for values in metrics.values()):
-                # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
-                # `wandb.log(step=...)` still records the global step, so both axes remain available.
-                documents_seen_metric = f"{PhaseType.training}/documents_seen"
-                self._wandb.define_metric(documents_seen_metric)
-                self._wandb.define_metric(f"{PhaseType.training}/*", step_metric=documents_seen_metric)
-                self._defined_documents_axis = True
-
             wandb.log(metrics, step=completed_steps, commit=commit)  # noqa
 
     def alert(self, title, text, level="INFO", wait=0.001) -> None:

From 006eb9c1b1b0e5dc2decc08e95c39361d0c84ad2 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 14:05:43 -0400
Subject: [PATCH 08/10] Always include num_documents/documents_seen in logged
 training metrics

Document counting is unconditional now, so the presence check was dead;
num_documents/documents_seen show up in both the console log line and W&B
on every logging step instead of only when the check happened to pass.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 fast_llm/engine/training/trainer.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
index 27bb898c3..094abb0f2 100644
--- a/fast_llm/engine/training/trainer.py
+++ b/fast_llm/engine/training/trainer.py
@@ -263,11 +263,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}
-                                if step_num_documents is not None
-                                else {}
-                            ),
+                            "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()

From baffd3aa08271502f494628361c187d5990fe549 Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 15:20:23 -0400
Subject: [PATCH 09/10] Review fixes: trust int document-count boundary, keep
 step as default W&B x-axis

- run_step's document count is always an int (all top-level inputs are
  TokenModelInput); drop the illusory int|None plumbing and the trainer's
  None guard.
- Drop the consumer-naming duplicate comment on _num_documents_in_batch.
- Drop the redundant class-level _documents_seen = 0 default.
- Drop the define_metric block so global step stays the default x-axis;
  documents_seen remains a logged, selectable metric.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 fast_llm/engine/schedule/runner.py  | 8 +++-----
 fast_llm/engine/training/trainer.py | 8 +++-----
 fast_llm/engine/training/wandb.py   | 6 ------
 3 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py
index 7ad56671b..fb3ef538d 100644
--- a/fast_llm/engine/schedule/runner.py
+++ b/fast_llm/engine/schedule/runner.py
@@ -66,9 +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`, or None when the data does not
-    # provide document counts (i.e. non-token data).
-    _num_documents_in_batch: int | None = None
+    # 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
@@ -151,7 +150,7 @@ def run_step(
         *,
         iteration: int = 1,
         return_metrics: bool = False,
-    ) -> tuple[dict[str, float | int], bool, dict[str, typing.Any] | None, int | 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:
@@ -325,7 +324,6 @@ 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
         )
-        # Whole-step DP-summed document count (the loss-normalization divisor).
         self._num_documents_in_batch = model_inputs[0][0].num_documents_in_batch
 
         for micro_batch, model_inputs_ in enumerate(model_inputs):
diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
index 094abb0f2..4b206afae 100644
--- a/fast_llm/engine/training/trainer.py
+++ b/fast_llm/engine/training/trainer.py
@@ -39,7 +39,7 @@ class Trainer[ConfigType: TrainerConfig](Configurable[ConfigType], abc.ABC):
     _wandb: Wandb
     _optimizer: Optimizer | None
     _completed_steps: int
-    _documents_seen: int = 0
+    _documents_seen: int
     _schedule: Schedule
 
     def __init__(self, config: TrainerConfig):
@@ -228,10 +228,8 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                     return_metrics=is_logging,
                 )
 
-                # Cumulative document count across training steps; `None` when the data provides no
-                # document counts.
-                if step_num_documents is not None:
-                    self._documents_seen += step_num_documents
+                # Cumulative document count across training steps.
+                self._documents_seen += step_num_documents
 
                 # Advanced, skipped, and Nan iterations.
                 if update_successful:
diff --git a/fast_llm/engine/training/wandb.py b/fast_llm/engine/training/wandb.py
index 819577490..9257fe9c0 100644
--- a/fast_llm/engine/training/wandb.py
+++ b/fast_llm/engine/training/wandb.py
@@ -5,7 +5,6 @@
 
 from fast_llm.config import Config
 from fast_llm.engine.config_utils.run import Run
-from fast_llm.engine.distributed.config import PhaseType
 from fast_llm.engine.training.config import WandbConfig
 
 
@@ -44,11 +43,6 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
                     wandb_path.write_text(yaml.safe_dump(wandb_config))
             # TODO: Does wandb work with nested configs?
             self._wandb = wandb.init(config=experiment_config.to_dict(), **wandb_config)
-            # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
-            # `wandb.log(step=...)` still records the global step, so both axes remain available.
-            documents_seen_metric = f"{PhaseType.training}/documents_seen"
-            self._wandb.define_metric(documents_seen_metric)
-            self._wandb.define_metric(f"{PhaseType.training}/*", step_metric=documents_seen_metric)
         else:
             self._wandb = None
 

From 6176cc95e7ddd488ea3f72e43211764b694040dd Mon Sep 17 00:00:00 2001
From: Joel Lamy-Poirier <joel.lamy-poirier@servicenow.com>
Date: Thu, 9 Jul 2026 15:57:20 -0400
Subject: [PATCH 10/10] Review fixes: dedup document-count expression, drop
 restating comments

- Extract the real-document-count local in _get_model_input and reuse it
  for both num_documents and num_documents_in_sequence.
- Reword the num_documents comment to describe this rank's contribution.
- Drop the comment restating the documents_seen accumulation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 fast_llm/data/document/token.py     | 19 ++++++++-----------
 fast_llm/engine/training/trainer.py |  1 -
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/fast_llm/data/document/token.py b/fast_llm/data/document/token.py
index c6da7101b..380dfd0ff 100644
--- a/fast_llm/data/document/token.py
+++ b/fast_llm/data/document/token.py
@@ -102,13 +102,13 @@ def _get_model_input(self, begin: int, end: int, config: LengthPreprocessingConf
         model_input = self._model_input_class(tokens=self.tokens[begin:end])
         lengths, first_document_begin, last_document_end = self._get_cropped_lengths(begin, end)
 
-        # 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
@@ -123,10 +123,7 @@ def _get_model_input(self, begin: int, end: int, config: LengthPreprocessingConf
                 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,
diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py
index 4b206afae..9ddafe5b3 100644
--- a/fast_llm/engine/training/trainer.py
+++ b/fast_llm/engine/training/trainer.py
@@ -228,7 +228,6 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                     return_metrics=is_logging,
                 )
 
-                # Cumulative document count across training steps.
                 self._documents_seen += step_num_documents
 
                 # Advanced, skipped, and Nan iterations.