fix: add a mechanism for allowing more max tokens (#662)

<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION
& AFFILIATES. All rights reserved. -->
<!-- SPDX-License-Identifier: Apache-2.0 -->

<!-- Thank you for contributing to Safe Synthesizer! -->

# Summary
<!-- Brief description of changes -->

## Pre-Review Checklist

<!-- These checks should be completed before a PR is reviewed, -->
<!-- but you can submit a draft early to indicate that the issue is
being worked on. -->

Ensure that the following pass:

- [ ] `mise run format && mise run check` or via prek validation.
- [ ] `mise run test` passes locally
- [ ] `mise run test:e2e` passes locally
- [ ] `mise run test:ci-container` passes locally (recommended)
- [ ] GPU CI status check passes -- comment `/sync` on this PR to
trigger a run (auto-triggers on ready-for-review)

## Pre-Merge Checklist

<!-- These checks need to be completed before a PR is merged, -->
<!-- but as PRs often change significantly during review, -->
<!-- it's OK for them to be incomplete when review is first requested.
-->

- [ ] New or updated tests for any fix or new behavior
- [ ] Updated documentation for new features and behaviors, including
docstrings for API docs.

## Other Notes

<!-- Please add the issue number that should be closed when this PR is
merged. -->
- Closes #<issue>


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added `max_tokens_multiplier` to control the per-request generation
token budget, defaulting to **1.2**.
* Enforced validation to allow only finite values strictly **greater
than 0**.
* **Bug Fixes**
* Improved consistency of max-token budgeting across generation backends
by applying the configured multiplier while keeping existing
context-window clamping behavior.
* **Tests**
* Added unit tests for the new multiplier field (default, valid/invalid
values including non-finite cases) and for multiplier behavior
(including clamping).
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Matt Kornfield <mkornfield@nvidia.com>

Author: mckornfield

src/nemo_safe_synthesizer/config/generate.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import math
 import warnings
 from collections.abc import Mapping
 from typing import Annotated, Any, ClassVar, Literal, Self
@@ -264,6 +265,18 @@ class GenerateParameters(Parameters, BaseModel):
         ),
     ] = 0.8
 
+    max_tokens_multiplier: Annotated[
+        float,
+        ValueValidator(value_func=lambda v: math.isfinite(v) and v > 0),
+        Field(
+            title="max_tokens_multiplier",
+            description=(
+                "Multiplier on the longest training example when sizing per-sample "
+                "max_tokens. Must be a finite value > 0. Default 1.2."
+            ),
+        ),
+    ] = 1.2  # mirrors llm.metadata.GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER (kept a literal to avoid a config->llm import cycle)
+
     structured_generation: StructuredGenerationParameters = Field(
         description="Structured generation parameters controlling schema-constrained output.",
         default_factory=StructuredGenerationParameters,

src/nemo_safe_synthesizer/generation/timeseries_backend.py

@@ -921,7 +921,10 @@ def generate(
             top_p=self.config.generation.top_p,
             top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"],
             min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"],
-            max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()),
+            max_tokens=self.model_metadata.generation_max_tokens_for(
+                self._get_prompt_token_count(),
+                multiplier=self.config.generation.max_tokens_multiplier,
+            ),
             skip_special_tokens=True,
             include_stop_str_in_output=False,
             ignore_eos=False,

src/nemo_safe_synthesizer/generation/vllm_backend.py

@@ -746,7 +746,10 @@ def _run_generation(self, data_actions_fn: utils.DataActionsFn | None) -> None:
             top_p=self.config.generation.top_p,
             top_k=FIXED_RUNTIME_GENERATE_ARGS["top_k"],
             min_p=FIXED_RUNTIME_GENERATE_ARGS["min_p"],
-            max_tokens=self.model_metadata.generation_max_tokens_for(self._get_prompt_token_count()),
+            max_tokens=self.model_metadata.generation_max_tokens_for(
+                self._get_prompt_token_count(),
+                multiplier=self.config.generation.max_tokens_multiplier,
+            ),
             skip_special_tokens=not need_special_token_outputs,
             include_stop_str_in_output=need_special_token_outputs,
             ignore_eos=False,

src/nemo_safe_synthesizer/llm/metadata.py

@@ -470,13 +470,13 @@ def max_seq_length(self) -> int:
             rsf = self.rope_scaling.factor
         return int((self.base_max_seq_length or DEFAULT_MAX_SEQ_LENGTH) * rsf)
 
-    def generation_max_tokens_for(self, prompt_len: int) -> int:
+    def generation_max_tokens_for(self, prompt_len: int, multiplier: float | None = None) -> int:
         """Per-sample ``max_tokens`` ceiling, prompt-aware.
 
         Returns the smaller of:
 
-        1. ``int(max_tokens_per_example * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)``
-           when the assembler stat is populated, else ``max_seq_length``.
+        1. ``int(max_tokens_per_example * multiplier)`` when the assembler stat
+           is populated, else ``max_seq_length``.
         2. ``max_seq_length - prompt_len`` -- vLLM raises when
            ``len(prompt) + max_tokens > max_model_len``
            (`vllm#33418 <https://github.com/vllm-project/vllm/issues/33418>`_).
@@ -488,15 +488,29 @@ def generation_max_tokens_for(self, prompt_len: int) -> int:
         clamp is a defensive belt for legacy adapters where the assembler
         stat is missing and for prompts longer than those seen in training.
 
+        The default ``multiplier`` (``GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER``)
+        adds only a small jitter margin, which is enough for most tables but
+        too tight for long, unbounded free-text columns: a model that
+        over-generates slightly past the longest training example truncates
+        mid-JSON and yields no parseable record. Callers wire the user-facing
+        ``generation.max_tokens_multiplier`` knob through here to widen the
+        budget (bounded by the context window) for such datasets.
+
         Args:
             prompt_len: Tokenized length of the prompt this sample will
                 run against. Pass ``0`` to disable the prompt clamp.
+            multiplier: Margin applied to ``max_tokens_per_example``. Defaults
+                to ``GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER`` when ``None`` so
+                non-generation callers (e.g. the training eval callback) keep
+                the legacy sizing.
 
         Returns:
             Non-negative ``max_tokens`` value safe to feed to ``SamplingParams``.
         """
+        if multiplier is None:
+            multiplier = GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
         if self.max_tokens_per_example and self.max_tokens_per_example > 0:
-            sized = int(self.max_tokens_per_example * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+            sized = int(self.max_tokens_per_example * multiplier)
         else:
             sized = self.max_seq_length
         return max(0, min(sized, self.max_seq_length - prompt_len))

tests/config/test_generate.py

@@ -217,3 +217,22 @@ def test_legacy_keys_with_no_nested_dict_are_migrated(self) -> None:
         assert params.structured_generation.backend == "xgrammar"
         assert params.structured_generation.schema_method == "structural_tag"
         assert params.structured_generation.use_single_sequence is True
+
+
+@pytest.mark.unit
+class TestMaxTokensMultiplier:
+    def test_default_matches_metadata_constant(self) -> None:
+        """The config default mirrors the metadata safety-margin constant."""
+        from nemo_safe_synthesizer.llm.metadata import GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
+
+        assert GenerateParameters().max_tokens_multiplier == GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER
+
+    def test_accepts_widened_value(self) -> None:
+        """Users can widen the budget for long free-text datasets."""
+        assert GenerateParameters(max_tokens_multiplier=1.8).max_tokens_multiplier == 1.8
+
+    @pytest.mark.parametrize("value", [0, -0.5, float("inf"), float("-inf"), float("nan")])
+    def test_rejects_non_positive(self, value: float) -> None:
+        """Non-positive and non-finite multipliers are rejected by the validator."""
+        with pytest.raises(ValidationError):
+            GenerateParameters(max_tokens_multiplier=value)

tests/generation/test_vllm_backend.py

@@ -1048,8 +1048,10 @@ def capture_and_stop(**kwargs):
         assert captured["max_tokens"] == 4_200
         # Engine is not initialized in this plumbing test, so the cached
         # prompt-token count falls back to 0; the helper is still called
-        # exactly once with that value.
-        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(0)
+        # exactly once with that value plus the configured budget multiplier.
+        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(
+            0, multiplier=base_params.generation.max_tokens_multiplier
+        )
 
     def test_passes_cached_prompt_token_count_when_engine_initialized(
         self, base_params, mock_model_metadata, mock_schema, mock_workdir
@@ -1079,7 +1081,9 @@ def capture_and_stop(**kwargs):
             backend.generate()
 
         assert captured["max_tokens"] == 4_096
-        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(37)
+        mock_model_metadata.generation_max_tokens_for.assert_called_once_with(
+            37, multiplier=base_params.generation.max_tokens_multiplier
+        )
         # Cached: a second access does not retokenize.
         assert backend._get_prompt_token_count() == 37
         fake_tokenizer.encode.assert_called_once_with(backend.prompt)

tests/llm/test_metadata.py

@@ -689,6 +689,32 @@ def test_metadata_generation_max_tokens_for_never_returns_negative(self, sample_
         # Beyond-context prompts still produce a safe value vLLM can accept.
         assert sample_model_metadata.generation_max_tokens_for(sample_model_metadata.max_seq_length + 64) == 0
 
+    def test_metadata_generation_max_tokens_for_none_multiplier_uses_default(self, sample_model_metadata):
+        """``multiplier=None`` reproduces the legacy safety-margin sizing."""
+        sample_model_metadata.max_tokens_per_example = 1000
+        expected = int(1000 * GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+        assert sample_model_metadata.generation_max_tokens_for(10, multiplier=None) == expected
+        # Explicitly passing the default constant matches the None path.
+        assert (
+            sample_model_metadata.generation_max_tokens_for(10, multiplier=GENERATION_MAX_TOKENS_SAFETY_MULTIPLIER)
+            == expected
+        )
+
+    def test_metadata_generation_max_tokens_for_custom_multiplier_widens_budget(self, sample_model_metadata):
+        """A larger multiplier widens the stat-derived budget (the long-text fix)."""
+        sample_model_metadata.max_tokens_per_example = 1000
+        # 1000 * 1.8 = 1800, still within the 2048 window given a small prompt.
+        assert sample_model_metadata.generation_max_tokens_for(10, multiplier=1.8) == 1800
+
+    def test_metadata_generation_max_tokens_for_custom_multiplier_still_clamped_to_window(self, sample_model_metadata):
+        """The prompt/window clamp still binds even with an aggressive multiplier."""
+        assert sample_model_metadata.max_seq_length == 2048
+        sample_model_metadata.max_tokens_per_example = 1500  # * 4.0 = 6000, far past the window
+        prompt_len = 48
+        assert sample_model_metadata.generation_max_tokens_for(prompt_len, multiplier=4.0) == (
+            sample_model_metadata.max_seq_length - prompt_len
+        )
+
     def test_metadata_max_tokens_per_example_round_trips_through_metadata_json(self, sample_model_metadata):
         """``max_tokens_per_example`` persists through save → load."""
         sample_model_metadata.max_tokens_per_example = 1500