Loss masking

.gitignore

@@ -37,3 +37,6 @@ devenv.*
 
 # direnv
 .direnv
+
+# wandb
+wandb/

fast_llm/data/sample/language_model.py

@@ -100,21 +100,41 @@ def __init__(
         chosen_spans: RangeBatch | None = None,
         rejected_spans: RangeBatch | None = None,
         image_patches: PatchBatch | None = None,
+        valid_tokens: int | None = None,
     ):
         self.tokens = tokens
         self.loss_masking_spans = loss_masking_spans
         self.chosen_spans = chosen_spans
         self.rejected_spans = rejected_spans
         self.image_patches = image_patches
+        self.valid_tokens = valid_tokens
 
     @classmethod
     def from_samples(cls, samples: typing.Iterable[LanguageModelSample]) -> typing.Self:
+        samples = list(samples)
+        token_batch = TokenBatch.from_samples([sample.tokens for sample in samples])
+        loss_masking_spans = _merge_optional(
+            RangeBatch.from_samples, [sample.loss_masking_spans for sample in samples]
+        )
+
+        # Calculate valid tokens for this batch (used for gradient accumulation weighting)
+        valid_tokens = None
+        if loss_masking_spans is not None:
+            batch_size, sequence_length = token_batch.tokens.shape
+            # Start with all tokens
+            valid_tokens = batch_size * sequence_length
+            # Subtract masked tokens
+            for sample_ranges in loss_masking_spans.ranges:
+                for begin, end in sample_ranges:
+                    valid_tokens -= end - begin
+
         return cls(
-            TokenBatch.from_samples([sample.tokens for sample in samples]),
-            _merge_optional(RangeBatch.from_samples, [sample.loss_masking_spans for sample in samples]),
+            token_batch,
+            loss_masking_spans,
             _merge_optional(RangeBatch.from_samples, [sample.chosen_spans for sample in samples]),
             _merge_optional(RangeBatch.from_samples, [sample.rejected_spans for sample in samples]),
             _merge_optional(PatchBatch.from_samples, [sample.image_patches for sample in samples]),
+            valid_tokens,
         )
 
     def crop(self, begin: int, end: int) -> typing.Self:
@@ -124,6 +144,7 @@ def crop(self, begin: int, end: int) -> typing.Self:
             _crop_optional(self.chosen_spans, begin, end),
             _crop_optional(self.rejected_spans, begin, end),
             _crop_optional(self.image_patches, begin, end),
+            valid_tokens=None,  # Cropped batches don't have valid token counts
         )
 
     def to_device_(self, device: "torch.device | str"):

fast_llm/data/sample/token.py

@@ -142,7 +142,7 @@ def get_document(self, index: int, begin: int, end: int) -> Sample:
         begin_ = self._size_cumsums[index].item()
         # Torch doesn't support type promotion between signed and unsigned types, so we convert here to avoid issues.
         # Convert begin and end to int to avoid numpy dtype overflow when adding to begin_
-        return TokenSample(self._tokens[begin_ + begin : begin_ + end].to(torch.int64), [end - begin])
+        return TokenSample(self._tokens[begin_ + int(begin) : begin_ + int(end)].to(torch.int64), [end - begin])
 
     def get_document_sizes(self) -> torch.Tensor:
         return self._size_cumsums[1:] - self._size_cumsums[:-1]

fast_llm/engine/evaluation/config.py

@@ -8,6 +8,7 @@
 
 if typing.TYPE_CHECKING:
     from fast_llm.engine.evaluation.evaluator import Evaluator, EvaluatorLmEval, LossEvaluator
+    from fast_llm.engine.evaluation.forward_kl.evaluator import ForwardKLEvaluator
 
 
 @config_class()
@@ -119,3 +120,58 @@ def get_evaluator(
         from fast_llm.engine.evaluation.lm_eval.evaluator import LmEvalEvaluator
 
         return LmEvalEvaluator(name, self, batch_config, data_load_num_proc, train_iters)
+
+
+@config_class(dynamic_type={EvaluatorConfig: "forward_kl"})
+class ForwardKLEvaluatorConfig(EvaluatorConfig):
+    _abstract: typing.ClassVar[bool] = False
+
+    dataset_path: str | None = Field(
+        default=None,
+        desc="HuggingFace dataset path containing teacher traces.",
+        hint=FieldHint.core,
+    )
+    split: str = Field(
+        default="validation",
+        desc="Dataset split to evaluate on. Use 'train+validation' syntax to combine multiple splits.",
+        hint=FieldHint.optional,
+    )
+    seed: int = Field(
+        default=42,
+        desc="Random seed for shuffling traces. Ensures reproducible evaluation across runs.",
+        hint=FieldHint.optional,
+    )
+    num_samples: int | None = Field(
+        default=None,
+        desc="Maximum number of traces to evaluate (after shuffling). None for all.",
+        hint=FieldHint.optional,
+        valid=skip_valid_if_none(check_field(Assert.gt, 0)),
+    )
+    batch_size: int = Field(
+        default=8,
+        desc="Batch size for forward passes.",
+        hint=FieldHint.performance,
+        valid=check_field(Assert.gt, 0),
+    )
+    trust_remote_code: bool = Field(
+        default=False,
+        desc="Trust remote code when loading dataset.",
+        hint=FieldHint.optional,
+    )
+    inference_mixer: str | None = Field(
+        default=None,
+        desc="Name of the mixer to use during evaluation (for StochasticMixer models). "
+        "If None, uses the model's default main_mixer_name.",
+        hint=FieldHint.optional,
+    )
+
+    def get_evaluator(
+        self,
+        name: str,
+        batch_config: BatchConfig,
+        data_load_num_proc: int,
+        train_iters: int | None = None,
+    ) -> "ForwardKLEvaluator":
+        from fast_llm.engine.evaluation.forward_kl.evaluator import ForwardKLEvaluator
+
+        return ForwardKLEvaluator(name, self, batch_config, data_load_num_proc, train_iters)