Skip to content

Fix GenerationConfig continuous batching serialization #47038

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
VectorPeak wants to merge 4 commits into
base: main
Choose a base branch
from
+40 −2
Open

Fix GenerationConfig continuous batching serialization #47038

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e0172b5
Fix continuous batching config serialization
VectorPeak Jul 3, 2026
3cafc9c
Merge branch 'main' into fix-continuous-batching-config-serialization
VectorPeak Jul 3, 2026
dac3f93
Use generic dataclass serialization fallback
VectorPeak Jul 3, 2026
599f658
Merge branch 'main' into fix-continuous-batching-config-serialization
VectorPeak Jul 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions src/transformers/generation/configuration_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,8 @@ def __init__(self, **kwargs):
self.disable_compile = kwargs.pop("disable_compile", None)

self.continuous_batching_config = kwargs.pop("continuous_batching_config", None)
if isinstance(self.continuous_batching_config, dict):
self.continuous_batching_config = ContinuousBatchingConfig(**self.continuous_batching_config)

# Deprecated (moved to the Hub). TODO remove for v5
self.low_memory = kwargs.pop("low_memory", None)
Expand Down Expand Up @@ -1798,6 +1800,11 @@ class ContinuousBatchingConfig:
max_cached_graphs: int | None = None

def __post_init__(self):
if isinstance(self.varlen_compile_config, dict):
self.varlen_compile_config = CompileConfig(**self.varlen_compile_config)
if isinstance(self.decode_compile_config, dict):
self.decode_compile_config = CompileConfig(**self.decode_compile_config)

# Only turn off graph mixing support if TP is on
graph_mixing_supported = os.environ.get("NCCL_GRAPH_MIXING_SUPPORT", "1") == "1"
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
Expand Down Expand Up @@ -1836,3 +1843,12 @@ def cuda_graph_booleans(self) -> tuple[bool, bool]:
def fallback_max_blocks_per_request(self) -> int:
"""Fallback if no user-hint is given and decode path is available."""
return 32

def to_dict(self) -> dict[str, Any]:
"""Serializes this instance to a Python dictionary."""
output = copy.deepcopy(self.__dict__)
if self.varlen_compile_config is not None:
output["varlen_compile_config"] = self.varlen_compile_config.to_dict()
if self.decode_compile_config is not None:
output["decode_compile_config"] = self.decode_compile_config.to_dict()
return output

@zucchini-nlp zucchini-nlp Jul 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rather than manually doing it per key, lets force all dataclasses to resolve as dicts. I copied this from a few lines above

def convert_dataclass_to_dict(obj):
            if isinstance(obj, dict):
                return {key: convert_dataclass_to_dict(value) for key, value in obj.items()}
            elif is_dataclass(obj):
                # Some of our dataclasses have a custom `to_dict()` method, and we prefer it
                if hasattr(obj, "to_dict"):
                    return obj.to_dict()
            else:
                return obj

@VectorPeak VectorPeak Jul 3, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the suggestion! I updated the patch to use a generic dataclass fallback in convert_dataclass_to_dict() and removed the per-key ContinuousBatchingConfig.to_dict() handling.

My initial intent with the manual keys was to avoid widening the behavioral surface, but this shared fallback is cleaner and matches the direction you suggested while still preferring custom to_dict() implementations when present.

Validation rerun locally:

python -m pytest tests/generation/test_configuration_utils.py::GenerationConfigSerializationTest::test_serialize_generation_continuous_batching_config -q
python -m pytest tests/generation/test_configuration_utils.py::GenerationConfigSerializationTest::test_serialize_generation_watermarking_config -q
python -m ruff format --check src/transformers/generation/configuration_utils.py tests/generation/test_configuration_utils.py
python -m ruff check src/transformers/generation/configuration_utils.py tests/generation/test_configuration_utils.py
git diff --check

@zucchini-nlp zucchini-nlp Jul 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @remi-or , to make sure if saving continuous-batch config is intended

32 changes: 31 additions & 1 deletion tests/generation/test_configuration_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,14 @@
from huggingface_hub import create_pull_request
from parameterized import parameterized

from transformers import AutoConfig, GenerationConfig, WatermarkingConfig, is_torch_available
from transformers import (
AutoConfig,
CompileConfig,
ContinuousBatchingConfig,
GenerationConfig,
WatermarkingConfig,
is_torch_available,
)
from transformers import logging as transformers_logging


Expand Down Expand Up @@ -755,6 +762,29 @@ def test_serialize_generation_watermarking_config(self):
self.assertEqual(watermark.seeding_scheme, seeding_scheme)
self.assertEqual(watermark.context_width, context_width)

def test_serialize_generation_continuous_batching_config(self):
"""Tests that GenerationConfig serializes ContinuousBatchingConfig parameters."""
continuous_batching_config = ContinuousBatchingConfig(
block_size=128,
default_compile_level=2,
varlen_compile_config=CompileConfig(dynamic=True),
decode_compile_config=CompileConfig(mode="default"),
)
generation_config = GenerationConfig(continuous_batching_config=continuous_batching_config)

with tempfile.TemporaryDirectory("test-generation-config") as tmp_dir:
generation_config.save_pretrained(tmp_dir)
new_config = GenerationConfig.from_pretrained(tmp_dir)

self.assertIsInstance(new_config.continuous_batching_config, ContinuousBatchingConfig)
self.assertEqual(new_config.continuous_batching_config.block_size, 128)
self.assertEqual(new_config.continuous_batching_config.default_compile_level, 2)
self.assertIsInstance(new_config.continuous_batching_config.varlen_compile_config, CompileConfig)
self.assertTrue(new_config.continuous_batching_config.varlen_compile_config.dynamic)
self.assertNotIn("_compile_all_devices", new_config.continuous_batching_config.varlen_compile_config.to_dict())
self.assertIsInstance(new_config.continuous_batching_config.decode_compile_config, CompileConfig)
self.assertEqual(new_config.continuous_batching_config.decode_compile_config.mode, "default")


@is_staging_test
class ConfigPushToHubTester(unittest.TestCase):
Expand Down
Loading