Add sampling observer auxiliary output hooks (#35747)

Co-authored-by: Alec Solder <alecs@fb.com>
2026-08-21 19:28:18 -07:00
parent 5662c03363
commit fbafd1b123
17 changed files with 1525 additions and 18 deletions
@@ -198,6 +198,77 @@ class TestApplyLogitsBias(CustomTestCase):
        info.apply_logits_bias(logits)
        self.assertTrue(torch.equal(logits, original))

+    def test_apply_logits_bias_without_penalizer_orchestrator(self):
+        info = _make_info(batch_size=1, penalizer_orchestrator=None)
+        logits = torch.zeros(1, VOCAB_SIZE)
+
+        info.apply_logits_bias(logits)
+
+        self.assertTrue(torch.equal(logits, torch.zeros_like(logits)))
+
+    def test_observer_sees_production_constraint_boundary(self):
+        events = []
+
+        class Observer:
+            def before_grammar(self, logits, sampling_info):
+                events.append(("before", logits.clone()))
+                return object()
+
+        grammar = MagicMock()
+        grammar.apply_vocab_mask.side_effect = lambda logits, vocab_mask: logits.fill_(
+            -4.0
+        )
+        info = _make_info(batch_size=1)
+        info.acc_additive_penalties = torch.ones(1, VOCAB_SIZE)
+        info.grammar_mask = GrammarMask(grammar, torch.ones(1, VOCAB_SIZE))
+        info.logit_bias = torch.full((1, VOCAB_SIZE), 2.0)
+        logits = torch.zeros(1, VOCAB_SIZE)
+
+        state = info.apply_logits_bias_with_observer(logits, observer=Observer())
+
+        self.assertIsNotNone(state)
+        self.assertTrue(torch.equal(events[0][1], torch.ones_like(logits)))
+        self.assertEqual(len(events), 1)
+        self.assertTrue(torch.equal(logits, torch.full_like(logits, -2.0)))
+        grammar.apply_vocab_mask.assert_called_once()
+
+    def test_observer_path_preserves_production_logit_transforms(self):
+        class Observer:
+            def before_grammar(self, logits, sampling_info):
+                return object()
+
+        def make_info():
+            grammar = MagicMock()
+            grammar.apply_vocab_mask.side_effect = (
+                lambda logits, vocab_mask: logits.add_(vocab_mask)
+            )
+            info = _make_info(batch_size=1)
+            info.acc_additive_penalties = torch.linspace(
+                -0.5, 0.5, VOCAB_SIZE
+            ).unsqueeze(0)
+            info.acc_scaling_penalties = torch.linspace(1.0, 1.5, VOCAB_SIZE).unsqueeze(
+                0
+            )
+            info.grammar_mask = GrammarMask(
+                grammar,
+                torch.linspace(-2.0, 0.0, VOCAB_SIZE).unsqueeze(0),
+            )
+            info.logit_bias = torch.linspace(0.0, 1.0, VOCAB_SIZE).unsqueeze(0)
+            return info
+
+        ordinary = make_info()
+        observed = make_info()
+        ordinary_logits = torch.linspace(-3.0, 3.0, VOCAB_SIZE).unsqueeze(0)
+        observed_logits = ordinary_logits.clone()
+
+        ordinary.apply_logits_bias(ordinary_logits)
+        observed.apply_logits_bias_with_observer(
+            observed_logits,
+            observer=Observer(),
+        )
+
+        self.assertTrue(torch.equal(observed_logits, ordinary_logits))
+

 # update_penalties
 class TestUpdatePenalties(CustomTestCase):